import argparse
import requests
import re
import sys

# Парсинг аргументов командной строки
parser = argparse.ArgumentParser(description="Extract domain sequences from FASTA using UniProt API.")
parser.add_argument("-inp", required=True, help="Input FASTA file")
parser.add_argument("-domain", required=True, help="Main domain description to extract (e.g., 'CheW-like')")
parser.add_argument("-include_domains", nargs="*", default=[], help="List of additional domain descriptions that must be present (optional)")
parser.add_argument("-out", required=True, help="Output FASTA file with extracted domains")
parser.add_argument("-unextracted", help="Output FASTA file for full sequences that passed filters (without domain trimming)")
args = parser.parse_args()

# Функция для извлечения Accession из заголовка FASTA
def get_accession(header):
    # Предполагаем, что первое слово до '|' или пробела является Accession
    match = re.match(r'^>(\S+)', header)
    if match:
        acc = match.group(1)
        if '|' in acc:
            acc = acc.split('|')[0]
        return acc
    else:
        return None

# Чтение входного FASTA
try:
    with open(args.inp, 'r') as f:
        lines = f.readlines()
except Exception as e:
    print(f"Error reading input file: {e}")
    sys.exit(1)

# Разбиваем на записи
records = []
current_header = None
current_seq = []
for line in lines:
    line = line.strip()
    if line.startswith('>'):
        if current_header is not None:
            records.append((current_header, ''.join(current_seq)))
        current_header = line
        current_seq = []
    else:
        current_seq.append(line)
if current_header is not None:
    records.append((current_header, ''.join(current_seq)))

# Открываем выходные файлы
try:
    out_handle = open(args.out, 'w')
except Exception as e:
    print(f"Error opening output file {args.out}: {e}")
    sys.exit(1)

unextracted_handle = None
if args.unextracted:
    try:
        unextracted_handle = open(args.unextracted, 'w')
    except Exception as e:
        print(f"Error opening unextracted file {args.unextracted}: {e}")
        out_handle.close()
        sys.exit(1)

def write_fasta(handle, header, sequence):
    handle.write(header + "\n")
    for i in range(0, len(sequence), 60):
        handle.write(sequence[i:i+60] + "\n")

for header, seq in records:
    acc = get_accession(header)
    if not acc:
        print(f"Could not parse accession from header: {header}")
        continue

    # Запрос к UniProt API
    url = f"https://rest.uniprot.org/uniprotkb/{acc}.json"
    try:
        response = requests.get(url)
        if response.status_code != 200:
            print(f"API request failed for {acc}: HTTP {response.status_code}")
            continue
        data = response.json()
    except Exception as e:
        print(f"Error fetching data for {acc}: {e}")
        continue

    # Извлечение features типа Domain
    features = data.get('features', [])
    domain_features = [f for f in features if f.get('type') == 'Domain']

    # Проверка наличия дополнительных доменов (если указаны)
    if args.include_domains:
        required = set(args.include_domains)
        found_descriptions = {f.get('description', '') for f in domain_features}
        missing = required - found_descriptions
        if missing:
            print(f"{acc}: missing required domains: {', '.join(missing)}")
            continue

    # Находим основной домен
    main_feature = None
    for f in domain_features:
        if f.get('description') == args.domain:
            main_feature = f
            break
    if not main_feature:
        print(f"{acc}: main domain '{args.domain}' not found")
        continue

    # Извлечение координат
    loc = main_feature.get('location')
    if not loc:
        print(f"{acc}: location not found for domain")
        continue
    start = loc.get('start', {}).get('value')
    end = loc.get('end', {}).get('value')
    if start is None or end is None:
        print(f"{acc}: invalid coordinates for domain")
        continue

    # Извлечение полной последовательности белка
    full_seq = data.get('sequence', {}).get('value')
    if not full_seq:
        print(f"{acc}: sequence not found")
        continue

    # Вырезаем домен
    domain_seq = full_seq[start-1:end]
    if not domain_seq:
        print(f"{acc}: extracted domain sequence is empty")
        continue

    # Запись в основной выходной файл
    out_header = f">{acc}|{args.domain}|{start}-{end}"
    write_fasta(out_handle, out_header, domain_seq)

    # Если указан файл для необрезанных последовательностей, записываем полную
    if unextracted_handle is not None:
        # Используем оригинальный заголовок
        write_fasta(unextracted_handle, header, full_seq)

# Закрываем файлы
out_handle.close()
if unextracted_handle is not None:
    unextracted_handle.close()

print(f"Done. Output written to {args.out}")
if args.unextracted:
    print(f"Unextracted sequences written to {args.unextracted}")