from __future__ import annotations
import anthropic
from ..models import ClaudeExtraction, 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]:
    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
