"""Module A: Knowledge Novelty — measures how much of the document Claude already knows."""

from __future__ import annotations
import json
import anthropic
from kvl.ingestor import Document

_CLAIM_PROMPT = """You are analyzing a document to evaluate how novel its knowledge is to AI language models.

Extract exactly {n} specific, verifiable factual claims from the document below. Focus on:
- Precise facts, findings, statistics, or assertions
- Claims that are specific (not general background knowledge)
- Claims that could be tested with a yes/no or short answer question

Return ONLY a JSON array of objects with keys "claim" and "question", where "question" is the specific question whose answer is the claim.

Example format:
[
  {{"claim": "The study found a 34% reduction in soil erosion.", "question": "What reduction in soil erosion did the study find?"}},
  ...
]

Document:
{document}"""

_JUDGE_PROMPT = """You are evaluating whether an AI model's answer demonstrates prior knowledge of a specific claim.

Claim from document: {claim}
AI model's closed-book answer: {answer}

Rate how well the model's answer already captures the claim WITHOUT having seen the document.
Score 0-1 where:
- 1.0 = model's answer is accurate and complete — this is not novel knowledge
- 0.5 = model has partial or approximate knowledge
- 0.0 = model doesn't know this — this IS novel knowledge

Return ONLY a JSON object: {{"score": <float 0-1>, "reason": "<one sentence>"}}"""


def _call_claude(client: anthropic.Anthropic, system: str, user: str, model: str = "claude-sonnet-4-6") -> str:
    msg = client.messages.create(
        model=model,
        max_tokens=2048,
        messages=[{"role": "user", "content": user}],
        system=system,
    )
    return msg.content[0].text.strip()


def _extract_claims(client: anthropic.Anthropic, doc: Document, n: int = 12) -> list[dict]:
    # Use the full document text but cap at ~6000 words to stay within limits
    text = " ".join(doc.raw.split()[:6000])
    prompt = _CLAIM_PROMPT.format(n=n, document=text)
    raw = _call_claude(client, "You extract factual claims from documents.", prompt)
    # Strip markdown code fences if present
    raw = raw.strip()
    if raw.startswith("```"):
        raw = "\n".join(raw.split("\n")[1:])
        raw = raw.rsplit("```", 1)[0]
    try:
        claims = json.loads(raw)
        return claims[:n]
    except json.JSONDecodeError:
        return []


def _closed_book_answer(client: anthropic.Anthropic, question: str) -> str:
    return _call_claude(
        client,
        "Answer the question using only your pre-trained knowledge. Do not make up information. If unsure, say so.",
        question,
        model="claude-haiku-4-5-20251001",  # cheaper for bulk closed-book queries
    )


def _judge_novelty(client: anthropic.Anthropic, claim: str, answer: str) -> dict:
    raw = _call_claude(
        client,
        "You are an expert evaluator assessing AI knowledge coverage.",
        _JUDGE_PROMPT.format(claim=claim, answer=answer),
    )
    raw = raw.strip()
    if raw.startswith("```"):
        raw = "\n".join(raw.split("\n")[1:])
        raw = raw.rsplit("```", 1)[0]
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        return {"score": 0.5, "reason": "Could not parse judge response."}


def _evaluate_claim(args):
    client, item = args
    answer = _closed_book_answer(client, item["question"])
    judgment = _judge_novelty(client, item["claim"], answer)
    return {
        "claim": item["claim"],
        "question": item["question"],
        "model_answer": answer,
        "known_score": judgment["score"],
        "reason": judgment.get("reason", ""),
    }


def evaluate(client: anthropic.Anthropic, doc: Document, progress_cb=None, max_workers: int = 6) -> dict:
    """Return novelty score (0-100) and detailed results."""
    from concurrent.futures import ThreadPoolExecutor

    if progress_cb:
        progress_cb("Extracting factual claims from document...")

    claims = _extract_claims(client, doc)
    if not claims:
        return {"score": 50, "details": [], "summary": "Could not extract claims from document."}

    if progress_cb:
        progress_cb(f"Testing {len(claims)} claims in parallel...")

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        results = list(pool.map(_evaluate_claim, [(client, item) for item in claims]))

    avg_known = sum(r["known_score"] for r in results) / len(results)
    novelty_score = round((1 - avg_known) * 100)

    return {
        "score": novelty_score,
        "details": results,
        "summary": f"Tested {len(results)} claims. Model already knows ~{round(avg_known*100)}% of this knowledge.",
    }