"""annotate_genes — what HGNC and UniProt say each gene and its protein are.

HGNC is the authority for the symbol (approved name, locus type, Ensembl id) and for the
gene → UniProt accession link. UniProt is then asked once per batch of up to 100 accessions,
reviewed (Swiss-Prot) human entries only, for the recommended name, function and subcellular
location. A gene with no reviewed entry — a pseudogene such as SORD2P — keeps its HGNC row with
``protein: null`` and is listed in ``unresolved``; it is never filled in from anything else.
"""

from __future__ import annotations

from typing import Any

from src import payloads
from src.endpoints._common import (
    Budget,
    Collector,
    according_to,
    get_client,
    humanise_list,
    resolve_symbols,
    stamped,
)
from src.identifiers import parse_gene_list

ENDPOINT = "annotate_genes"
UNIPROT_SEARCH = "UniProt_search"
UNIPROT_BATCH = 100
UNIPROT_FIELDS = [
    "accession",
    "gene_primary",
    "protein_name",
    "cc_function",
    "cc_subcellular_location",
]
MAX_FUNCTION_CHARS = 600


def annotate_genes(genes: str = "") -> dict[str, Any]:
    parsed = parse_gene_list(genes, allow_ensembl=False)
    if not parsed.ok:
        return payloads.refused(ENDPOINT, parsed.reason, parsed.message)
    client, budget, out = get_client(), Budget(), Collector()
    symbols = [i.value for i in parsed.identifiers]
    resolved = resolve_symbols(symbols, client, budget, out)

    accessions = sorted({acc for g in resolved for acc in g.uniprot_ids})
    entries = fetch_reviewed_entries(accessions, client, budget, out)

    with_protein = 0
    for gene in resolved:
        protein = None
        for acc in gene.uniprot_ids:
            if acc in entries:
                entry, result = entries[acc]
                protein = stamped(entry_fields(entry), result, ref=acc, url=uniprot_url(acc))
                break
        if protein is not None:
            with_protein += 1
            status = "reviewed UniProt entry"
        elif not gene.uniprot_ids:
            status = f"HGNC links no UniProt entry (locus type: {gene.locus_type or 'not stated'})"
            out.unresolved_input(gene.input, f"no reviewed UniProt entry: {status}")
        else:
            status = "no reviewed (Swiss-Prot) human entry for " + ", ".join(gene.uniprot_ids)
            out.unresolved_input(gene.input, status)
        out.rows.append(
            {
                "input": gene.input,
                "symbol": gene.symbol,
                "hgnc_id": gene.hgnc_id,
                "name": gene.name,
                "locus_type": gene.locus_type,
                "ensembl_gene_id": gene.ensembl_gene_id,
                "protein": protein,
                "protein_status": status,
                "provenance": gene.provenance_row["provenance"],
            }
        )

    if not out.rows and out.source_errors:
        return payloads.error(
            ENDPOINT,
            "all_sources_failed",
            "HGNC did not answer; nothing was looked up.",
            source_errors=out.source_errors,
        )
    lead = according_to(out.sources_used or {"hgnc"}, out.retrieved)
    approved = "an approved HGNC gene" if len(symbols) == 1 else "approved HGNC genes"
    sentences = [
        (
            f"{lead}, {len(resolved)} of {len(symbols)} input{'s' if len(symbols) != 1 else ''} "
            f"{'is' if len(symbols) == 1 else 'are'} {approved}, and {with_protein} "
            f"{'has' if with_protein == 1 else 'have'} a reviewed human UniProt entry."
        )
    ]
    if out.unresolved:
        sentences.append(
            "Not resolved: "
            + humanise_list([f"{u['input']} ({u['reason']})" for u in out.unresolved], limit=8)
            + "."
        )
    if out.source_errors:
        sentences.append(
            "Source problems: "
            + "; ".join(f"{e['source']} {e['reason']}" for e in out.source_errors)
            + "."
        )
    return payloads.ok(
        ENDPOINT,
        query={"genes": symbols},
        rows=out.rows,
        sources_used=out.sources_used or {"hgnc"},
        narrative=" ".join(sentences),
        unresolved=out.unresolved,
        source_errors=out.source_errors,
    )


def uniprot_url(accession: str) -> str:
    return f"https://www.uniprot.org/uniprotkb/{accession}/entry"


def fetch_reviewed_entries(
    accessions: list[str], client, budget: Budget, out: Collector
) -> dict[str, tuple[dict[str, Any], Any]]:
    """Reviewed human UniProt entries by accession, in batches of ``UNIPROT_BATCH``."""
    found: dict[str, tuple[dict[str, Any], Any]] = {}
    for start in range(0, len(accessions), UNIPROT_BATCH):
        batch = accessions[start : start + UNIPROT_BATCH]
        query = (
            "("
            + " OR ".join(f"accession:{a}" for a in batch)
            + ") AND reviewed:true AND organism_id:9606"
        )
        result = out.call(
            client,
            budget,
            UNIPROT_SEARCH,
            {"query": query, "limit": len(batch), "fields": UNIPROT_FIELDS},
        )
        if result is None:
            continue
        for entry in result.data.get("results") or []:
            acc = entry.get("primaryAccession")
            if acc in batch:
                found[acc] = (entry, result)
    return found


def entry_fields(entry: dict[str, Any]) -> dict[str, Any]:
    comments = entry.get("comments") or []
    function = " ".join(
        t.get("value", "")
        for c in comments
        if c.get("commentType") == "FUNCTION"
        for t in c.get("texts") or []
    ).strip()
    locations = []
    for c in comments:
        if c.get("commentType") != "SUBCELLULAR LOCATION":
            continue
        for loc in c.get("subcellularLocations") or []:
            value = (loc.get("location") or {}).get("value")
            if value and value not in locations:
                locations.append(value)
    recommended = ((entry.get("proteinDescription") or {}).get("recommendedName") or {}).get(
        "fullName"
    ) or {}
    genes = entry.get("genes") or []
    return {
        "accession": entry.get("primaryAccession"),
        "gene": ((genes[0].get("geneName") or {}).get("value") if genes else None),
        "protein_name": recommended.get("value"),
        "function": function[:MAX_FUNCTION_CHARS]
        + ("…" if len(function) > MAX_FUNCTION_CHARS else ""),
        "function_truncated": len(function) > MAX_FUNCTION_CHARS,
        "subcellular_location": locations,
    }