dee/core/resolve.py

"""Paste-anything target resolution (Phase 3, M2).

One input box, three kinds of input — auto-detected and resolved to an
editable DNA sequence (plus a human-readable label and, where known, a
gene symbol that feeds base-edit AA consequences + the structure viewer):

  * raw DNA / FASTA  → used as-is (never leaves the Space).
  * gene symbol      → Ensembl canonical-transcript CDS (human / mouse),
                        via the existing exon.fetch_gene_structure().
  * accession        → Ensembl transcript/gene ID (ENST…/ENSG…) or RefSeq
                        (NM_/XM_/NR_/XR_) fetched from Ensembl / NCBI.

Privacy: for symbol/accession lookups, ONLY the (organism, identifier)
leaves the Space — never the user's pasted sequence. The raw-sequence
path makes no network calls at all.

Classification is conservative: a long, DNA-looking blob is always a
sequence; a short token that matches an ID pattern is an accession; an
alphanumeric token is a gene symbol (needs an organism). Ambiguous short
tokens fall through to a symbol lookup, which fails gracefully.
"""
from __future__ import annotations

import re
import urllib.parse
from typing import Dict, Tuple

from dee.core import exon as _exon

# NCBI taxonomy IDs for the UniProt structure lookup (M5).
_UNIPROT_TAXID = {"human": "9606", "mouse": "10090"}

# ─── Identifier patterns ─────────────────────────────────────────────
_ENSEMBL_TX   = re.compile(r"^ENS[A-Z]*T\d{6,}(?:\.\d+)?$", re.I)   # ENST…, ENSMUST…
_ENSEMBL_GENE = re.compile(r"^ENS[A-Z]*G\d{6,}(?:\.\d+)?$", re.I)   # ENSG…, ENSMUSG…
_REFSEQ       = re.compile(r"^[NX][MR]_\d+(?:\.\d+)?$", re.I)        # NM_/NR_/XM_/XR_
_SYMBOL       = re.compile(r"^[A-Za-z][A-Za-z0-9._-]{0,19}$")        # TP53, BRCA1, …

_NCBI_EFETCH = ("https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
                "?db=nuccore&rettype=fasta&retmode=text&id=")

MIN_TARGET_LEN = 23   # shortest usable target (Cas12a spacer+PAM)
MAX_TARGET_LEN = 1_000_000


def _clean_dna(text: str) -> str:
    """Strip FASTA headers + whitespace + non-ACGTN, uppercase."""
    lines = [ln for ln in text.splitlines() if not ln.strip().startswith(">")]
    return re.sub(r"[^ACGTNacgtn]", "", "".join(lines)).upper()


def classify(text: str) -> Tuple[str, str]:
    """Return (kind, value). kind ∈ {empty, sequence, symbol, ensembl_tx,
    ensembl_gene, refseq, unknown}. `value` is the cleaned sequence (for
    'sequence') or the identifier token otherwise."""
    t = (text or "").strip()
    if not t:
        return ("empty", "")
    if t.lstrip().startswith(">"):
        return ("sequence", _clean_dna(t))

    cleaned = _clean_dna(t)
    letters = re.sub(r"[^A-Za-z]", "", re.sub(r"\s", "", t))
    dna_ratio = len(cleaned) / max(1, len(letters))
    has_ws = bool(re.search(r"\s", t))
    # A long, DNA-dominant blob (or any multi-line / spaced DNA) is a sequence.
    if len(cleaned) >= MIN_TARGET_LEN and dna_ratio >= 0.9 and (has_ws or len(cleaned) > 20):
        return ("sequence", cleaned)

    token = t.split()[0] if t.split() else ""
    if _ENSEMBL_TX.match(token):
        return ("ensembl_tx", token)
    if _ENSEMBL_GENE.match(token):
        return ("ensembl_gene", token)
    if _REFSEQ.match(token):
        return ("refseq", token.upper())
    if _SYMBOL.match(token):
        return ("symbol", token.upper())
    if len(cleaned) >= MIN_TARGET_LEN and dna_ratio >= 0.9:
        return ("sequence", cleaned)
    return ("unknown", token)


def _err(msg: str) -> Dict:
    return {"ok": False, "error": msg, "kind": "", "sequence": "",
            "gene_symbol": "", "label": "", "source": ""}


def _ok(kind: str, sequence: str, label: str, source: str,
        gene_symbol: str = "") -> Dict:
    return {"ok": True, "kind": kind, "sequence": sequence,
            "gene_symbol": gene_symbol, "label": label, "source": source}


def _fetch_ensembl_cds(identifier: str, is_gene: bool) -> Tuple[str, str]:
    """Fetch a CDS sequence for an Ensembl transcript or gene ID.
    Returns (sequence, label) or ("", "")."""
    tx_id = identifier
    if is_gene:
        # Resolve gene → canonical transcript first.
        info = _exon._http_get_json(
            f"{_exon.ENSEMBL_BASE}/lookup/id/{identifier}?expand=1")
        if not info:
            return ("", "")
        transcripts = info.get("Transcript", []) or []
        if not transcripts:
            return ("", "")
        canonical = next((t for t in transcripts if t.get("is_canonical") == 1),
                         transcripts[0])
        tx_id = canonical.get("id") or ""
        if not tx_id:
            return ("", "")
    fasta = _exon._http_get_text(
        f"{_exon.ENSEMBL_BASE}/sequence/id/{tx_id}?type=cds")
    seq = _exon._fasta_to_seq(fasta) if fasta else ""
    if not seq:
        return ("", "")
    label = (f"{identifier} → {tx_id} · CDS {len(seq):,} nt"
             if is_gene else f"{tx_id} · CDS {len(seq):,} nt")
    return (seq, label)


def _fetch_refseq(accession: str) -> str:
    fasta = _exon._http_get_text(_NCBI_EFETCH + accession)
    return _exon._fasta_to_seq(fasta) if fasta else ""


def resolve_uniprot(organism: str, gene_symbol: str) -> Dict:
    """Resolve a gene symbol (+ organism) to a UniProt accession and the
    AlphaFold-DB structure URLs, for the structure viewer (Phase 3, M5).

    Returns {ok, uniprot, alphafold_url, alphafold_page, error?}. Only the
    (organism, gene_symbol) leaves the Space. Uses Ensembl xrefs, which
    map the gene symbol → SwissProt accession for human/mouse.
    """
    taxid = _UNIPROT_TAXID.get((organism or "").lower())
    if not taxid:
        return {"ok": False, "error": "Structure view needs Human or Mouse."}
    if not gene_symbol:
        return {"ok": False, "error": "No gene symbol to look up."}
    # UniProt REST: reviewed (SwissProt) entry for this gene + organism.
    # Gene-level Ensembl xrefs don't carry SwissProt (that's protein-level),
    # so we query UniProt directly. Only (gene, organism) is sent.
    q = urllib.parse.quote(f"gene:{gene_symbol} AND organism_id:{taxid} AND reviewed:true")
    url = (f"https://rest.uniprot.org/uniprotkb/search?query={q}"
           f"&fields=accession&format=json&size=1")
    data = _exon._http_get_json(url)
    acc = ""
    results = (data or {}).get("results") if isinstance(data, dict) else None
    if results:
        acc = (results[0] or {}).get("primaryAccession", "")
    if not acc:
        return {"ok": False,
                "error": f"No reviewed UniProt entry found for {gene_symbol} ({organism})."}
    return {
        "ok": True,
        "uniprot": acc,
        # Fallback URL only — the client re-resolves the current model version
        # via the AlphaFold prediction API (the DB bumps versions: v4→v6→…).
        "alphafold_url":  f"https://alphafold.ebi.ac.uk/files/AF-{acc}-F1-model_v6.pdb",
        "alphafold_page": f"https://alphafold.ebi.ac.uk/entry/{acc}",
    }


def resolve_target(text: str, organism: str = "") -> Dict:
    """Resolve pasted text to an editable sequence.

    Returns a dict: {ok, kind, sequence, gene_symbol, label, source, error?}.
    """
    organism = (organism or "").lower().strip()
    kind, val = classify(text)

    if kind == "empty":
        return _err("Paste a DNA sequence, a gene symbol, or an accession.")

    if kind == "sequence":
        if len(val) < MIN_TARGET_LEN:
            return _err(f"Sequence is only {len(val)} nt — need at least "
                        f"{MIN_TARGET_LEN} nt. Check you pasted DNA, not protein.")
        if len(val) > MAX_TARGET_LEN:
            return _err("Sequence too long. Cap is 1 Mbp — paste just the "
                        "gene / region you're editing.")
        return _ok("sequence", val, f"pasted sequence · {len(val):,} nt", "input")

    if kind == "symbol":
        if organism not in ("human", "mouse"):
            return _err(f"To look up “{val}” by gene symbol, pick Human or "
                        f"Mouse — or paste the sequence directly.")
        gene = _exon.fetch_gene_structure(organism, val)
        if gene is None:
            return _err(f"Couldn't find “{val}” in {organism}. Check the "
                        f"symbol, or paste the sequence directly.")
        return _ok("gene", gene.cds_sequence,
                   f"{val} · {gene.transcript_id} · CDS {len(gene.cds_sequence):,} nt",
                   "ensembl", gene_symbol=val)

    if kind in ("ensembl_tx", "ensembl_gene"):
        seq, label = _fetch_ensembl_cds(val, is_gene=(kind == "ensembl_gene"))
        if not seq:
            return _err(f"Couldn't fetch “{val}” from Ensembl. Check the ID, "
                        f"or paste the sequence directly.")
        return _ok("ensembl", seq, label, "ensembl")

    if kind == "refseq":
        seq = _fetch_refseq(val)
        if not seq:
            return _err(f"Couldn't fetch “{val}” from NCBI. Check the "
                        f"accession, or paste the sequence directly.")
        return _ok("refseq", seq, f"{val} · {len(seq):,} nt", "ncbi")

    return _err("Unrecognized input. Paste a DNA sequence, a gene symbol "
                "(with Human/Mouse selected), or an accession (ENST…, NM_…).")