from __future__ import annotations

import asyncio
import logging
import os

from app.schemas.visual_lesson import EvidenceClaim, EvidenceSource, VisualizationRequest
from app.services.visual_evidence import VisualEvidenceResolver, _clean


CANONICAL_NUCLEIC_SOURCES = [
    EvidenceSource(
        source_id="builtin:nucleic-chemistry",
        origin="builtin",
        title="NCBI Bookshelf: Molecular Biology of the Cell — DNA structure",
        url="https://www.ncbi.nlm.nih.gov/books/NBK26821/",
        excerpt="Canonical reference metadata for nucleotides, antiparallel strands, complementary base pairing, and the sugar-phosphate backbone.",
        authority="canonical",
    ),
    EvidenceSource(
        source_id="builtin:nucleic-forms",
        origin="builtin",
        title="RCSB PDB-101: Nucleic acid structure",
        url="https://pdb101.rcsb.org/learn/guide-to-understanding-pdb-data/introduction",
        excerpt="Canonical reference metadata for archived nucleic-acid structures and common DNA/RNA structural representations.",
        authority="canonical",
    ),
]

logger = logging.getLogger(__name__)


class NucleicEvidenceResolver:
    """Resolve only context that can affect the requested nucleic visualization."""

    def __init__(self, base: VisualEvidenceResolver | None = None) -> None:
        self.base = base or VisualEvidenceResolver()

    @staticmethod
    def _needs_context_research(request: VisualizationRequest) -> bool:
        prompt_text = request.prompt.lower()
        explicit_context_request = any(
            term in prompt_text
            for term in ("paper", "selected", "according to", "reported", "study", "figure", "this molecule", "this structure")
        )
        # Surrounding PDF text is transport context, not an instruction to use
        # the active paper. Named canonical molecules must not fan out into an
        # unrelated project merely because a document is open.
        named_canonical = any(
            term in prompt_text
            for term in ("adenine", "cytosine", "guanine", "thymine", "uracil")
        )
        if named_canonical and not explicit_context_request:
            return False
        return bool(
            explicit_context_request
            or request.selection_text
            or request.selection_snippets
            or request.selection_image_base64
        )

    @staticmethod
    def _dedupe(sources: list[EvidenceSource]) -> list[EvidenceSource]:
        output: list[EvidenceSource] = []
        seen: set[tuple[str, str]] = set()
        limits = {"selection": 4, "project": 4, "web": 3, "builtin": 2, "database": 3}
        counts: dict[str, int] = {}
        for source in sources:
            key = (source.source_id, source.excerpt[:240])
            if key in seen or counts.get(source.origin, 0) >= limits[source.origin]:
                continue
            seen.add(key)
            counts[source.origin] = counts.get(source.origin, 0) + 1
            output.append(source)
        return output

    async def resolve_context(self, project_id: str, request: VisualizationRequest) -> tuple[list[EvidenceSource], list[EvidenceClaim], list[str]]:
        warnings: list[str] = []
        sources = self.base._selection_sources(request)
        if self._needs_context_research(request):
            queries = list(dict.fromkeys(filter(None, [
                _clean(request.prompt),
                _clean(request.selection_text),
                "DNA RNA nucleotide nucleobase molecular structure",
            ])))
            logger.info("Nucleic evidence contextual fan-out request=%s queries=%d", request.request_id, len(queries))
            project_task = (
                asyncio.create_task(self.base._project_search(project_id, queries, request.active_document_ids))
                if request.project_context_enabled else None
            )
            image_task = asyncio.create_task(self.base._image_source(request))
            web_rows: list[dict] = []
            if os.getenv("TAVILY_API_KEY"):
                try:
                    search = self.base._web_search or self.base._default_web_search
                    batches = await asyncio.gather(*(search(query) for query in queries[:2]), return_exceptions=True)
                    for batch in batches:
                        if isinstance(batch, list):
                            web_rows.extend(batch)
                except Exception:
                    warnings.append("Live contextual research failed; canonical and official database evidence were still used.")
            if project_task is not None:
                try:
                    project_rows = await project_task
                    sources.extend(self.base._candidate_sources(project_rows, "project"))
                except Exception:
                    warnings.append("Project context retrieval failed; canonical and official database evidence were still used.")
            image_source = await image_task
            if image_source:
                sources.append(image_source)
            sources.extend(self.base._candidate_sources(web_rows, "web"))
        else:
            logger.info(
                "Nucleic evidence canonical-only request=%s active_documents=%d",
                request.request_id,
                len(request.active_document_ids),
            )
        sources.extend(CANONICAL_NUCLEIC_SOURCES)
        sources = self._dedupe(sources)
        claims = [
            EvidenceClaim(
                claim_id="nucleic-backbone",
                text="Nucleic-acid strands have a repeating sugar-phosphate backbone and directional 5-prime and 3-prime ends.",
                claim_type="standard_definition",
                support_level="canonical",
                source_ids=["builtin:nucleic-chemistry"],
            ),
            EvidenceClaim(
                claim_id="nucleic-pairing",
                text="The visualization uses canonical A–T and G–C DNA pairing, or A–U and G–C RNA pairing.",
                claim_type="standard_definition",
                support_level="canonical",
                source_ids=["builtin:nucleic-chemistry"],
            ),
            EvidenceClaim(
                claim_id="nucleic-forms",
                text="A-, B-, and Z-DNA are shown with distinct canonical helix parameters; Z-DNA is left-handed.",
                claim_type="standard_definition",
                support_level="canonical",
                source_ids=["builtin:nucleic-forms"],
            ),
        ]
        # Retrieved passages remain planning context only. A literal selection
        # anchor proves provenance, not that the passage supports a displayed
        # molecular fact, so it is never promoted to a claim automatically.
        logger.info(
            "Nucleic evidence resolved request=%s sources=%d claims=%d warnings=%d",
            request.request_id,
            len(sources),
            len(claims),
            len(warnings),
        )
        return sources, claims, warnings