import asyncio
from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, Field

from app.agents.cerebras_client import CerebrasClient
from app.rag.chromadb_client import ChromaDBClient
from app.rag.models import ConsumerType, EvidenceRequest
from app.rag.retrieval_service import PaperEvidenceService
from app.rag.vector_indexes import PaperEvidenceIndex


class _ClaimItem(BaseModel):
    claim: str = Field(description="A single factual claim asserted in the draft")
    citation: str = Field(
        default="",
        description="The citation marker attached to the claim (e.g. '[1]', 'Smith 2020'), or '' if none",
    )


class _ClaimExtraction(BaseModel):
    claims: List[_ClaimItem] = Field(
        default_factory=list,
        description="1-6 key factual claims from the draft that carry a citation",
    )


class _Verdict(BaseModel):
    verdict: Literal["supported", "partial", "unsupported", "contradicted", "not-found"]
    evidence: str = Field(
        description="One sentence grounded in the source passages that justifies the verdict "
                    "(or notes their absence). Do not invent facts beyond the passages."
    )


class DraftAgent:
    """Verifies the cited claims in a student's draft against the project's own
    uploaded corpus (ChromaDB), one claim at a time, grounded -> never guesses.
    """

    def __init__(self, client: Optional[CerebrasClient] = None, db: Optional[ChromaDBClient] = None) -> None:
        self._client = client or CerebrasClient()
        self._db = db

    def _get_db(self) -> ChromaDBClient:
        if self._db is None:
            self._db = ChromaDBClient()
        return self._db

    async def verify_draft(self, project_id: str, draft_text: str) -> List[Dict[str, Any]]:
        loop = asyncio.get_event_loop()

        # 1. Extract cited claims from the draft.
        try:
            extraction = await loop.run_in_executor(None, lambda: self._client.structured_complete(
                [
                    {
                        "role": "system",
                        "content": (
                            "Extract 1-6 key factual claims from the student's draft that are attributed "
                            "to a source (a citation marker like [1], (Smith, 2020), etc.). Return each "
                            "claim as a self-contained sentence plus the citation marker it carries. "
                            "Ignore rhetorical or unsupported sentences with no citation."
                        ),
                    },
                    {"role": "user", "content": draft_text[:8000]},
                ],
                _ClaimExtraction,
            ))
            claims = extraction.claims
        except Exception as e:
            print("Draft claim extraction error:", e)
            claims = []

        if not claims:
            return [{
                "claim": "No cited claims found in the draft.",
                "verdict": "not-found",
                "evidence": "Add citations (e.g. [1]) next to your factual statements so they can be verified.",
                "citation": "N/A",
            }]

        # 2. Verify each claim against the project corpus, concurrently.
        return await asyncio.gather(*(self._verify_claim(project_id, c) for c in claims))

    async def _verify_claim(self, project_id: str, claim: _ClaimItem) -> Dict[str, Any]:
        loop = asyncio.get_event_loop()
        citation = claim.citation or "N/A"

        chunks = await loop.run_in_executor(None, lambda: self._retrieve(project_id, claim.claim))
        if not chunks:
            return {
                "claim": claim.claim,
                "verdict": "not-found",
                "evidence": "No matching material was found in this project's uploaded corpus.",
                "citation": citation,
            }

        ctx = "\n\n".join(f"[{c.get('source', '?')}] {c['text']}" for c in chunks)
        try:
            verdict = await loop.run_in_executor(None, lambda: self._client.structured_complete(
                [
                    {
                        "role": "system",
                        "content": (
                            "You verify a single claim against retrieved source passages from the "
                            "student's own corpus. Classify strictly from the passages:\n"
                            "- supported: the passages directly affirm the claim.\n"
                            "- partial: the passages affirm part of it or affirm it with caveats.\n"
                            "- unsupported: the passages are on-topic but do not establish the claim.\n"
                            "- contradicted: the passages assert the opposite.\n"
                            "- not-found: the passages are unrelated to the claim.\n"
                            "Ground your one-sentence evidence in the passages -> never invent facts."
                        ),
                    },
                    {"role": "user", "content": f"CLAIM: {claim.claim}\n\nSOURCE PASSAGES:\n{ctx}"},
                ],
                _Verdict,
            ))
            return {
                "claim": claim.claim,
                "verdict": verdict.verdict,
                "evidence": verdict.evidence,
                "citation": citation,
            }
        except Exception as e:
            print("Draft verify error:", e)
            return {
                "claim": claim.claim,
                "verdict": "not-found",
                "evidence": "Verification could not be completed for this claim.",
                "citation": citation,
            }

    def _retrieve(self, project_id: str, query: str, n: int = 4) -> List[Dict[str, Any]]:
        service = PaperEvidenceService(
            index=None if self._db is None else PaperEvidenceIndex(self._get_db())
        )
        result = service.retrieve(EvidenceRequest(
            query=query,
            consumer=ConsumerType.DRAFT,
            project_id=project_id,
            token_budget=max(512, n * 500),
            project_memory_policy="none",
            student_memory_policy="none",
        ))
        citations = {citation.evidence_id: citation for citation in result.citations}
        return [
            {
                "text": item.evidence.index_text,
                "source": citations.get(item.evidence.evidence_id).filename
                if item.evidence.evidence_id in citations else item.evidence.document_id,
                "evidence_id": item.evidence.evidence_id,
                "page_start": item.evidence.page_start,
            }
            for item in result.evidence[:n]
        ]