from __future__ import annotations
import asyncio
import anthropic
from ..models import ClaudeExtraction, ClaudeFact, TokenUsage

BASE_SYSTEM_PROMPT = (
    "You are a fact extraction assistant. Extract key factual claims from the document. "
    "For each fact, provide the exact quote(s) from the document that support it. "
    "Be precise with quotes - they must be verbatim from the text.\n\n"
    "When a specific focus topic is requested and that topic has NO presence in the document, "
    "set no_mention_found to true and use absence_evidence to describe what the document "
    "actually covers instead."
)

MAX_TEXT_LENGTH = 50_000
INPUT_COST_PER_TOKEN = 3.0 / 1_000_000
OUTPUT_COST_PER_TOKEN = 15.0 / 1_000_000

RECORD_FACTS_TOOL = {
    "name": "record_facts",
    "description": "Record the extracted facts and their supporting citations from the document.",
    "input_schema": {
        "type": "object",
        "properties": {
            "facts": {
                "type": "array",
                "description": "List of factual claims extracted from the document",
                "items": {
                    "type": "object",
                    "properties": {
                        "fact": {"type": "string"},
                        "citations": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {"exact_quote": {"type": "string"}},
                                "required": ["exact_quote"],
                            },
                        },
                    },
                    "required": ["fact", "citations"],
                },
            },
            "no_mention_found": {"type": "boolean"},
            "absence_evidence": {"type": "string"},
        },
        "required": ["facts"],
    },
}


async def extract_facts(
    client: anthropic.AsyncAnthropic,
    text: str,
    custom_prompt: str | None = None,
) -> tuple[ClaudeExtraction, TokenUsage]:
    """Extract facts from a document text (up to ``MAX_TEXT_LENGTH`` chars).

    For documents larger than ``MAX_TEXT_LENGTH``, use :func:`extract_facts_chunked`
    to process the full text across multiple Claude calls.
    """
    truncated = text[:MAX_TEXT_LENGTH]

    system_prompt = (
        f"{BASE_SYSTEM_PROMPT}\n\nFocus specifically on: {custom_prompt}"
        if custom_prompt else BASE_SYSTEM_PROMPT
    )
    user_message = (
        f"Extract facts from this document, focusing on: {custom_prompt}\n\n"
        "For each fact, include the exact verbatim quote(s). "
        "If this topic is not present, set no_mention_found to true.\n\n"
        f"<document>\n{truncated}\n</document>"
        if custom_prompt
        else f"Extract the key facts from this document. "
             f"For each fact, include the exact verbatim quote(s).\n\n"
             f"<document>\n{truncated}\n</document>"
    )

    response = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=4096,
        system=system_prompt,
        messages=[{"role": "user", "content": user_message}],
        tools=[RECORD_FACTS_TOOL],
        tool_choice={"type": "tool", "name": "record_facts"},
    )

    tool_block = next((b for b in response.content if b.type == "tool_use"), None)
    if tool_block is None:
        raise RuntimeError("Claude did not return a tool use response")

    extraction = ClaudeExtraction(**tool_block.input)
    cost_usd = (
        response.usage.input_tokens * INPUT_COST_PER_TOKEN
        + response.usage.output_tokens * OUTPUT_COST_PER_TOKEN
    )
    usage = TokenUsage(
        inputTokens=response.usage.input_tokens,
        outputTokens=response.usage.output_tokens,
        costUsd=cost_usd,
    )
    return extraction, usage


async def extract_facts_chunked(
    client: anthropic.AsyncAnthropic,
    text: str,
    custom_prompt: str | None = None,
    chunk_size: int = 45_000,
    overlap: int = 5_000,
) -> tuple[ClaudeExtraction, TokenUsage]:
    """Extract facts from a document of any size by splitting it into
    overlapping chunks and running extraction on each concurrently.

    Chunks are broken at paragraph boundaries where possible.  Facts whose
    citations share the same leading 60 characters as an already-seen citation
    are deduplicated (they likely come from the overlap region).

    Parameters
    ----------
    client:
        Anthropic async client.
    text:
        Full document text (no length limit).
    custom_prompt:
        Optional focus topic passed to each chunk extraction.
    chunk_size:
        Maximum characters per chunk (default 45,000).
    overlap:
        Characters of overlap between consecutive chunks so facts that span
        a chunk boundary aren't missed (default 5,000).
    """
    # Build overlapping chunks, preferring paragraph breaks
    chunks: list[str] = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        if end < len(text):
            # Try to break at the last newline in the second half of the chunk
            newline = text.rfind("\n", start + chunk_size // 2, end)
            if newline > start:
                end = newline + 1
        chunks.append(text[start:end])
        if end >= len(text):
            break
        start = end - overlap

    # Single chunk — just delegate directly
    if len(chunks) == 1:
        return await extract_facts(client, chunks[0], custom_prompt)

    # Run all chunks concurrently
    results: list[tuple[ClaudeExtraction, TokenUsage]] = await asyncio.gather(
        *[extract_facts(client, chunk, custom_prompt) for chunk in chunks]
    )

    # Merge, deduplicate by citation prefix, aggregate usage
    seen_quote_prefixes: set[str] = set()
    unique_facts: list[ClaudeFact] = []

    for extraction, _ in results:
        for fact in extraction.facts:
            key = "|".join(sorted(c.exact_quote[:60] for c in fact.citations))
            if key not in seen_quote_prefixes:
                seen_quote_prefixes.add(key)
                unique_facts.append(fact)

    total_input = sum(u.inputTokens for _, u in results)
    total_output = sum(u.outputTokens for _, u in results)
    total_cost = sum(u.costUsd for _, u in results)

    merged_extraction = ClaudeExtraction(facts=unique_facts)
    merged_usage = TokenUsage(
        inputTokens=total_input,
        outputTokens=total_output,
        costUsd=total_cost,
    )
    return merged_extraction, merged_usage
