"""FASTA parsing and UniProt sequence fetching."""

from __future__ import annotations

from dataclasses import dataclass
from io import StringIO

import requests
from Bio import SeqIO

UNIPROT_FASTA_URL = "https://rest.uniprot.org/uniprotkb/{accession}.fasta"


@dataclass
class ProteinRecord:
    id: str
    description: str
    sequence: str


def parse_fasta(text: str) -> list[ProteinRecord]:
    records = SeqIO.parse(StringIO(text), "fasta")
    return [
        ProteinRecord(id=r.id, description=r.description, sequence=str(r.seq).upper())
        for r in records
    ]


def extract_accession(record_id: str) -> str:
    """Pull the accession out of a FASTA record id.

    UniProt headers follow the `db|accession|entry_name` convention (e.g.
    `sp|Q9UK17|KCNG1_HUMAN`), so the accession is the middle field. Anything
    else (a plain accession, a custom id, ...) is returned as-is.
    """
    parts = record_id.split("|")
    if len(parts) == 3:
        return parts[1]
    return record_id


def fetch_uniprot_sequence(accession: str, timeout: float = 10.0) -> ProteinRecord:
    """Fetch a single sequence from UniProt by accession (e.g. "Q9Y2W7")."""
    url = UNIPROT_FASTA_URL.format(accession=accession.strip())
    response = requests.get(url, timeout=timeout)
    response.raise_for_status()
    records = parse_fasta(response.text)
    if not records:
        raise ValueError(f"No sequence returned for accession {accession!r}")
    return records[0]