app/services/nucleic_resolver.py

from __future__ import annotations

import asyncio
import hashlib
import logging
import re
import secrets
import time
from dataclasses import dataclass, field
from typing import Any, Literal
from urllib.parse import urlparse

import httpx

from app.schemas.visual_lesson import (
    EvidenceClaim,
    EvidenceSource,
    MolecularResolutionCandidate,
    NucleicFeature,
    NucleicFormParameters,
    NucleobaseAtom,
    NucleobaseBond,
    NucleobaseMolecule,
    NucleicStructure,
    ProteinChain,
    VisualizationRequest,
)
from app.services.visual_lesson_store import VisualLessonStore


PDB_RE = re.compile(r"(?i)(?:\bPDB\s*[:#]?\s*|\b)(?=[0-9A-Za-z]{4}\b)(?=[0-9A-Za-z]*[A-Za-z])([0-9][A-Za-z0-9]{3})\b")
SEQUENCE_RE = re.compile(r"(?i)(?<![A-Z])([ACGTU]{6,60})(?![A-Z])")
EXPLICIT_SEQUENCE_RE = re.compile(r"(?i)\b(?:sequence|seq)\s*[:=]?\s*([A-Z]{4,})\b")
DNA_ALPHABET = set("ACGTRYSWKMBDHVN")
RNA_ALPHABET = set("ACGURYSWKMBDHVN")
DNA_COMPLEMENT = str.maketrans("ACGTRYSWKMBDHVN", "TGCAYRSWMKVHDBN")
RNA_COMPLEMENT = str.maketrans("ACGURYSWKMBDHVN", "UGCAYRSWMKVHDBN")
ILLUSTRATIVE_DNA = "ATGCCGTAACGT"
ILLUSTRATIVE_RNA = "AUGCCGUAACGU"
ALLOWED_STRUCTURE_HOSTS = {"models.rcsb.org", "files.rcsb.org"}
MAX_COORDINATE_BYTES = 50 * 1024 * 1024
NUCLEOBASES = {
    "adenine": ("A", "C5H5N5"),
    "cytosine": ("C", "C4H5N3O"),
    "guanine": ("G", "C5H5N5O"),
    "thymine": ("T", "C5H6N2O2"),
    "uracil": ("U", "C4H4N2O2"),
}
NUCLEOBASE_ROLES = {
    "A": "Adenine pairs with thymine through two hydrogen bonds in canonical DNA base pairing.",
    "C": "Cytosine pairs with guanine through three hydrogen bonds in canonical DNA base pairing.",
    "G": "Guanine pairs with cytosine through three hydrogen bonds in canonical DNA base pairing.",
    "T": "Thymine pairs with adenine through two hydrogen bonds in canonical DNA base pairing.",
    "U": "Uracil pairs with adenine through two hydrogen bonds in canonical RNA base pairing.",
}
ELEMENT_SYMBOLS = {1: "H", 6: "C", 7: "N", 8: "O", 15: "P", 16: "S"}

logger = logging.getLogger(__name__)

FORM_PARAMETERS = {
    "a_dna": NucleicFormParameters(form="a_dna", handedness="right", bases_per_turn=11.0, rise_angstrom=2.6, radius_angstrom=11.5, strand_offset_degrees=132.0),
    "b_dna": NucleicFormParameters(form="b_dna", handedness="right", bases_per_turn=10.5, rise_angstrom=3.4, radius_angstrom=10.0, strand_offset_degrees=144.0),
    "z_dna": NucleicFormParameters(form="z_dna", handedness="left", bases_per_turn=12.0, rise_angstrom=3.7, radius_angstrom=9.0, strand_offset_degrees=150.0),
    "rna": NucleicFormParameters(form="rna", handedness="right", bases_per_turn=11.0, rise_angstrom=2.8, radius_angstrom=11.0, strand_offset_degrees=130.0),
}


class NucleicResolutionError(RuntimeError):
    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code


@dataclass
class NucleicIntent:
    mode: Literal["concept", "sequence", "comparison", "structure"]
    molecule: Literal["dna", "rna", "dna_rna", "dna_forms"]
    sequence: str
    complement: str
    sequence_is_illustrative: bool
    forms: list[NucleicFormParameters]
    focus_base: str | None = None
    assumptions: list[str] = field(default_factory=list)


@dataclass
class ResolvedNucleic:
    intent: NucleicIntent
    structure: NucleicStructure | None
    features: list[NucleicFeature]
    sources: list[EvidenceSource]
    claims: list[EvidenceClaim]
    nucleobase_molecule: NucleobaseMolecule | None = None
    warnings: list[str] = field(default_factory=list)


class NucleicResolver:
    def __init__(self, store: VisualLessonStore | None = None, client: httpx.AsyncClient | None = None) -> None:
        self.store = store or VisualLessonStore()
        self._client = client

    @staticmethod
    def parse_intent(request: VisualizationRequest) -> NucleicIntent:
        text = f"{request.prompt} {request.selection_text}".upper()
        explicit_sequence = EXPLICIT_SEQUENCE_RE.search(text)
        sequence_match = explicit_sequence or SEQUENCE_RE.search(text)
        sequence = sequence_match.group(1).upper() if sequence_match else ""
        if len(sequence) > 60:
            raise NucleicResolutionError("sequence_too_long", "Sequence visualization supports at most 60 nucleotides.")
        if "T" in sequence and "U" in sequence:
            raise NucleicResolutionError("mixed_nucleic_alphabet", "A sequence cannot mix DNA thymine (T) and RNA uracil (U).")
        is_rna = "RNA" in text or "U" in sequence
        if sequence and not set(sequence) <= (RNA_ALPHABET if is_rna else DNA_ALPHABET):
            raise NucleicResolutionError("invalid_nucleic_sequence", "The sequence contains unsupported IUPAC nucleotide symbols.")
        compare_rna = ("DNA" in text and "RNA" in text) or "DNA VS RNA" in text or "DNA VERSUS RNA" in text
        compare_forms = (
            "DNA" in text
            and any(term in text for term in ("COMPARE", "FORM"))
            and all(any(token in text for token in (f"{letter}-DNA", f"{letter}-", f"{letter} DNA")) for letter in ("A", "B", "Z"))
        )
        chemistry = any(term in text for term in ("CHEMISTRY", "CHEMICAL", "NUCLEOTIDE", "ATOM", "BOND"))
        named_base = next((name for name in NUCLEOBASES if name.upper() in text), "")
        if compare_forms:
            molecule, mode, forms = "dna_forms", "comparison", [FORM_PARAMETERS["a_dna"], FORM_PARAMETERS["b_dna"], FORM_PARAMETERS["z_dna"]]
        elif compare_rna:
            molecule, mode, forms = "dna_rna", "comparison", [FORM_PARAMETERS["b_dna"], FORM_PARAMETERS["rna"]]
        else:
            molecule = "rna" if is_rna else "dna"
            mode = "sequence" if sequence else "concept"
            forms = [FORM_PARAMETERS["rna" if is_rna else "b_dna"]]
        illustrative = not sequence
        if not sequence:
            sequence = ILLUSTRATIVE_DNA if compare_rna else (ILLUSTRATIVE_RNA if is_rna else ILLUSTRATIVE_DNA)
        complement = sequence.translate(RNA_COMPLEMENT if molecule == "rna" else DNA_COMPLEMENT)
        assumptions = (
            ["The PubChem conformer supplies molecular connectivity and coordinates; atom radii, colors, labels, and lighting are display conventions."]
            if named_base
            else ["The parametric helix is illustrative geometry, not an atomistic coordinate model."]
        )
        if molecule == "rna":
            complement = ""
            assumptions.append("No RNA fold is inferred from sequence; the strand shape is illustrative.")
        if chemistry and not named_base:
            assumptions.append("Chemical structures are schematic 2D formulas; official coordinates are required for an atomistic 3D view.")
        return NucleicIntent(
            mode=mode,
            molecule=molecule,
            sequence=sequence,
            complement=complement,
            sequence_is_illustrative=illustrative,
            forms=forms,
            focus_base=NUCLEOBASES[named_base][0] if named_base else None,
            assumptions=assumptions,
        )

    async def _json(self, url: str) -> Any:
        owns = self._client is None
        client = self._client or httpx.AsyncClient(timeout=15.0, follow_redirects=False)
        try:
            response = await client.get(url, headers={"Accept": "application/json"})
            if response.status_code == 404:
                raise NucleicResolutionError("identifier_not_found", "The requested PDB structure was not found.")
            response.raise_for_status()
            return response.json()
        except NucleicResolutionError:
            raise
        except Exception as exc:
            raise NucleicResolutionError("official_database_unavailable", f"Official molecular metadata could not be retrieved: {exc}") from exc
        finally:
            if owns:
                await client.aclose()

    async def _coordinates(self, url: str) -> bytes:
        parsed = urlparse(url)
        if parsed.scheme != "https" or parsed.hostname not in ALLOWED_STRUCTURE_HOSTS:
            raise NucleicResolutionError("unsafe_structure_url", "The coordinate host is not allowed.")
        owns = self._client is None
        client = self._client or httpx.AsyncClient(timeout=30.0, follow_redirects=False)
        try:
            async with client.stream("GET", url, headers={"Accept": "application/octet-stream"}) as response:
                response.raise_for_status()
                if int(response.headers.get("content-length") or 0) > MAX_COORDINATE_BYTES:
                    raise NucleicResolutionError("invalid_structure", "The coordinate file exceeds 50 MB.")
                data = bytearray()
                async for chunk in response.aiter_bytes():
                    data.extend(chunk)
                    if len(data) > MAX_COORDINATE_BYTES:
                        raise NucleicResolutionError("invalid_structure", "The coordinate file exceeds 50 MB.")
            if len(data) < 64 or b"<html" in bytes(data[:512]).lower():
                raise NucleicResolutionError("invalid_structure", "The official source returned invalid coordinates.")
            return bytes(data)
        except NucleicResolutionError:
            raise
        except Exception as exc:
            raise NucleicResolutionError("structure_download_failed", f"The coordinate file could not be downloaded: {exc}") from exc
        finally:
            if owns:
                await client.aclose()

    async def classify_pdb(self, pdb_id: str) -> tuple[bool, bool, dict[str, Any], list[tuple[str, dict[str, Any]]]]:
        pdb_id = pdb_id.upper()
        entry = await self._json(f"https://data.rcsb.org/rest/v1/core/entry/{pdb_id}")
        ids = (entry.get("rcsb_entry_container_identifiers") or {}).get("polymer_entity_ids") or []
        rows = await asyncio.gather(*(self._json(f"https://data.rcsb.org/rest/v1/core/polymer_entity/{pdb_id}/{entity_id}") for entity_id in ids))
        entities = [(str(entity_id), row) for entity_id, row in zip(ids, rows)]
        kinds = [f"{(row.get('entity_poly') or {}).get('rcsb_entity_polymer_type', '')} {(row.get('entity_poly') or {}).get('type', '')}".lower() for _, row in entities]
        has_protein = any("protein" in kind or "polypeptide" in kind for kind in kinds)
        has_nucleic = any(any(term in kind for term in ("dna", "rna", "ribonucleotide", "deoxyribonucleotide")) for kind in kinds)
        return has_protein, has_nucleic, entry, entities

    async def _pubchem_nucleobase(self, name: str) -> NucleobaseMolecule:
        logger.info("PubChem nucleobase resolution start name=%s", name)
        payload = await self._json(
            f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{name}/JSON?record_type=3d"
        )
        compound = (payload.get("PC_Compounds") or [None])[0]
        if not isinstance(compound, dict):
            raise NucleicResolutionError("small_molecule_unavailable", f"PubChem returned no 3D conformer for {name}.")
        atom_block = compound.get("atoms") or {}
        coordinate_block = (compound.get("coords") or [None])[0] or {}
        conformer = (coordinate_block.get("conformers") or [None])[0] or {}
        atom_ids = list(atom_block.get("aid") or [])
        atomic_numbers = list(atom_block.get("element") or [])
        coordinate_ids = list(coordinate_block.get("aid") or [])
        xs, ys, zs = list(conformer.get("x") or []), list(conformer.get("y") or []), list(conformer.get("z") or [])
        coordinate_index = {int(atom_id): index for index, atom_id in enumerate(coordinate_ids)}
        atoms: list[NucleobaseAtom] = []
        for atom_id, atomic_number in zip(atom_ids, atomic_numbers):
            index = coordinate_index.get(int(atom_id), -1)
            element = ELEMENT_SYMBOLS.get(int(atomic_number))
            if index < 0 or element is None or index >= len(xs) or index >= len(ys) or index >= len(zs):
                continue
            atoms.append(NucleobaseAtom(atom_id=int(atom_id), element=element, x=float(xs[index]), y=float(ys[index]), z=float(zs[index])))
        bonds_block = compound.get("bonds") or {}
        bonds = [
            NucleobaseBond(atom_a=int(atom_a), atom_b=int(atom_b), order=int(order))
            for atom_a, atom_b, order in zip(
                bonds_block.get("aid1") or [], bonds_block.get("aid2") or [], bonds_block.get("order") or []
            )
        ]
        if not atoms or not bonds or len(atoms) > 64:
            raise NucleicResolutionError("small_molecule_unavailable", f"PubChem returned incomplete 3D data for {name}.")
        symbol, formula = NUCLEOBASES[name]
        cid = int(((compound.get("id") or {}).get("id") or {}).get("cid") or 0)
        molecule = NucleobaseMolecule(
            name=name,
            symbol=symbol,
            molecular_formula=formula,
            pubchem_cid=cid,
            source_url=f"https://pubchem.ncbi.nlm.nih.gov/compound/{cid}",
            atoms=atoms,
            bonds=bonds,
        )
        logger.info(
            "PubChem nucleobase resolution done name=%s cid=%d atoms=%d bonds=%d",
            name,
            molecule.pubchem_cid,
            len(molecule.atoms),
            len(molecule.bonds),
        )
        return molecule

    async def resolve(self, project_id: str, request: VisualizationRequest, sources: list[EvidenceSource], claims: list[EvidenceClaim], scope: str = "nucleic_acid") -> ResolvedNucleic:
        pdb_match = PDB_RE.search(f"{request.prompt} {request.selection_text}")
        if not pdb_match:
            intent = self.parse_intent(request)
            molecule = None
            warnings: list[str] = []
            if intent.focus_base:
                name = next(name for name, (symbol, _) in NUCLEOBASES.items() if symbol == intent.focus_base)
                claims.append(EvidenceClaim(
                    claim_id="nucleobase-role",
                    text=NUCLEOBASE_ROLES[intent.focus_base],
                    claim_type="standard_definition",
                    support_level="canonical",
                    source_ids=["builtin:nucleic-chemistry"],
                ))
                try:
                    molecule = await self._pubchem_nucleobase(name)
                    source_id = f"pubchem:{molecule.pubchem_cid}"
                    sources.append(EvidenceSource(
                        source_id=source_id,
                        origin="database",
                        title=f"PubChem Compound {molecule.pubchem_cid}: {name}",
                        url=molecule.source_url,
                        excerpt=f"PubChem 3D conformer and molecular graph for {name} ({molecule.molecular_formula}).",
                        authority="official",
                    ))
                    claims.append(EvidenceClaim(
                        claim_id="nucleobase-structure",
                        text=(
                            f"The displayed {name} molecular graph and computed 3D conformer come from "
                            f"PubChem CID {molecule.pubchem_cid}."
                        ),
                        claim_type="source_fact",
                        support_level="direct",
                        source_ids=[source_id],
                    ))
                    claims.append(EvidenceClaim(
                        claim_id="nucleobase-representation",
                        text=(
                            "Atom sizes, element colors, labels, lighting, and bond-cylinder spacing are "
                            "illustrative display choices; they are not measured atomic radii or electron density."
                        ),
                        claim_type="illustrative_choice",
                        support_level="illustrative",
                        source_ids=[],
                    ))
                except NucleicResolutionError as exc:
                    warnings.append(f"Official 3D nucleobase data was unavailable ({exc}); the chemistry schematic remains available.")
            return ResolvedNucleic(
                intent=intent,
                structure=None,
                features=self._generic_features(intent),
                sources=sources,
                claims=claims,
                nucleobase_molecule=molecule,
                warnings=warnings,
            )
        pdb_id = pdb_match.group(1).upper()
        has_protein, has_nucleic, entry, entities = await self.classify_pdb(pdb_id)
        if not has_nucleic and scope != "whole_assembly":
            raise NucleicResolutionError("unsupported_biomolecule", f"PDB {pdb_id} does not contain a DNA or RNA polymer.")
        nucleic_entities = []
        for entity_id, row in entities:
            kind = f"{(row.get('entity_poly') or {}).get('rcsb_entity_polymer_type', '')} {(row.get('entity_poly') or {}).get('type', '')}".lower()
            if scope == "whole_assembly" or any(term in kind for term in ("dna", "rna", "ribonucleotide", "deoxyribonucleotide")):
                nucleic_entities.append((entity_id, row, kind))
        coordinate_url = f"https://models.rcsb.org/{pdb_id.lower()}.bcif"
        coordinate_bytes = await self._coordinates(coordinate_url)
        asset_id = self.store.save_asset(project_id, coordinate_bytes, "bcif")
        chains: list[ProteinChain] = []
        polymer_types: list[str] = []
        features: list[NucleicFeature] = []
        for entity_id, row, kind in nucleic_entities:
            container = row.get("rcsb_polymer_entity_container_identifiers") or {}
            description = str((row.get("rcsb_polymer_entity") or {}).get("pdbx_description") or f"Polymer entity {entity_id}")
            asym_ids = list(container.get("asym_ids") or [])
            auth_ids = list(container.get("auth_asym_ids") or [])
            length = int((row.get("entity_poly") or {}).get("rcsb_sample_sequence_length") or 0)
            sequence = re.sub(r"[^ACGTU]", "", str((row.get("entity_poly") or {}).get("pdbx_seq_one_letter_code_can") or "").upper())
            polymer_types.append(kind.strip())
            for index, chain_id in enumerate(asym_ids):
                auth_id = str(auth_ids[index] if index < len(auth_ids) else chain_id)
                chains.append(ProteinChain(chain_id=str(chain_id), auth_chain_id=auth_id, entity_id=entity_id, description=description, sequence_length=length))
                features.append(NucleicFeature(
                    feature_id=f"chain:{chain_id}", kind="chain", label=f"Chain {auth_id}", color="#315E8A",
                    claim_ids=["nucleic-structure-chains"], chain_id=str(chain_id),
                ))
                features.extend(
                    NucleicFeature(
                        feature_id=f"base:{chain_id}:{position}", kind="base", label=f"{base}{position}",
                        color={"A": "#59B88A", "T": "#E56F9A", "U": "#E56F9A", "G": "#C566E5", "C": "#7482EA"}.get(base, "#8D9BAA"),
                        claim_ids=["nucleic-structure-bases"], chain_id=str(chain_id), start=position, end=position, comp_id=base,
                    )
                    for position, base in enumerate(sequence[:60], start=1)
                )
        identifiers = entry.get("rcsb_entry_container_identifiers") or {}
        resolution_values = (entry.get("rcsb_entry_info") or {}).get("resolution_combined") or []
        method = str(((entry.get("exptl") or [{}])[0]).get("method") or "")
        title = str((entry.get("struct") or {}).get("title") or f"PDB {pdb_id}")
        structure = NucleicStructure(
            pdb_id=pdb_id,
            assembly_id=str((identifiers.get("assembly_ids") or [""])[0]),
            title=title,
            coordinate_asset_id=asset_id,
            coordinate_sha256=hashlib.sha256(coordinate_bytes).hexdigest(),
            coordinate_format="bcif",
            source_url=coordinate_url,
            experimental_method=method,
            resolution_angstrom=float(resolution_values[0]) if resolution_values else 0.0,
            polymer_types=polymer_types,
            chains=chains,
        )
        source_id = f"rcsb:{pdb_id}"
        sources.append(EvidenceSource(
            source_id=source_id, origin="database", title=f"RCSB PDB {pdb_id}",
            url=f"https://www.rcsb.org/structure/{pdb_id}", excerpt=f"{title}; {method or 'method not reported'}.", authority="official",
        ))
        claims.extend([
            EvidenceClaim(claim_id="nucleic-structure-provenance", text=f"{pdb_id} is an archived experimental molecular structure.", claim_type="source_fact", support_level="direct", source_ids=[source_id]),
            EvidenceClaim(claim_id="nucleic-structure-chains", text=f"The selected structure view contains {len(chains)} resolved polymer chain instance(s).", claim_type="source_fact", support_level="direct", source_ids=[source_id]),
            EvidenceClaim(claim_id="nucleic-structure-bases", text="Displayed nucleotide selectors follow the archived polymer sequence and label numbering.", claim_type="source_fact", support_level="direct", source_ids=[source_id]),
        ])
        intent = NucleicIntent(mode="structure", molecule="dna" if not any("rna" in kind and "dna" not in kind for kind in polymer_types) else "rna", sequence="", complement="", sequence_is_illustrative=False, forms=[])
        warnings = ["This view includes the whole molecular assembly." ] if scope == "whole_assembly" and has_protein else []
        return ResolvedNucleic(
            intent=intent,
            structure=structure,
            features=features,
            sources=sources,
            claims=claims,
            warnings=warnings,
        )

    @staticmethod
    def _generic_features(intent: NucleicIntent) -> list[NucleicFeature]:
        if intent.focus_base:
            return []
        claim = ["nucleic-backbone"]
        features = [
            NucleicFeature(feature_id="backbone", kind="backbone", label="Sugar–phosphate backbone", color="#D7A642", claim_ids=claim),
            NucleicFeature(feature_id="direction-5", kind="direction", label="5′ end", color="#1A3557", claim_ids=claim),
            NucleicFeature(feature_id="direction-3", kind="direction", label="3′ end", color="#1A3557", claim_ids=claim),
            NucleicFeature(feature_id="sugar", kind="sugar", label="Ribose" if intent.molecule == "rna" else "Deoxyribose", color="#D7A642", claim_ids=claim),
            NucleicFeature(feature_id="phosphate", kind="phosphate", label="Phosphate", color="#315E8A", claim_ids=claim),
        ]
        if intent.complement:
            features.extend([
                NucleicFeature(feature_id="base-pair", kind="base_pair", label="Complementary base pair", color="#59B88A", claim_ids=["nucleic-pairing"]),
                NucleicFeature(feature_id="hydrogen-bond", kind="hydrogen_bond", label="Hydrogen bonds", color="#8D9BAA", claim_ids=["nucleic-pairing"]),
                NucleicFeature(feature_id="major-groove", kind="major_groove", label="Major groove", color="#7A5AA6", claim_ids=["nucleic-forms"]),
                NucleicFeature(feature_id="minor-groove", kind="minor_groove", label="Minor groove", color="#4A7FB5", claim_ids=["nucleic-forms"]),
            ])
        return features


class PendingMolecularResolutions:
    def __init__(self) -> None:
        self._items: dict[str, tuple[float, VisualizationRequest, list[EvidenceSource], list[EvidenceClaim], list[MolecularResolutionCandidate]]] = {}

    def create(self, request: VisualizationRequest, sources: list[EvidenceSource], claims: list[EvidenceClaim], candidates: list[MolecularResolutionCandidate]) -> str:
        self._purge()
        token = f"molecular_{secrets.token_urlsafe(24)}"
        self._items[token] = (time.time() + 600, request, sources, claims, candidates)
        return token

    def consume(self, token: str, candidate_id: str) -> tuple[VisualizationRequest, list[EvidenceSource], list[EvidenceClaim], MolecularResolutionCandidate]:
        self._purge()
        item = self._items.pop(token, None)
        if not item:
            raise NucleicResolutionError("resolution_expired", "The molecular structure choice expired. Run the visualization again.")
        _, request, sources, claims, candidates = item
        candidate = next((value for value in candidates if value.candidate_id == candidate_id), None)
        if not candidate:
            raise NucleicResolutionError("invalid_resolution_choice", "The molecular structure choice is no longer available.")
        return request, sources, claims, candidate

    def _purge(self) -> None:
        now = time.time()
        self._items = {key: value for key, value in self._items.items() if value[0] > now}