from __future__ import annotations
import anthropic
from ..models import Fact, ClaudeFact, ClaudeCitation
from .citation_mapper import map_citations

SYSTEM_PROMPT = (
    "You are a helpful assistant that answers questions about documents. "
    "Be conversational and concise. Always cite the exact text from the document "
    "that supports your answer. If the answer is not in the document, say so clearly "
    "and explain what the document does cover."
)

MAX_TEXT_LENGTH = 50_000
INPUT_COST_PER_TOKEN = 3.0 / 1_000_000
INPUT_CACHE_WRITE_COST = 3.75 / 1_000_000
INPUT_CACHE_READ_COST = 0.30 / 1_000_000
OUTPUT_COST_PER_TOKEN = 15.0 / 1_000_000

RESPOND_TOOL = {
    "name": "respond_with_citations",
    "description": "Respond to the user and cite exact quotes from the document.",
    "input_schema": {
        "type": "object",
        "properties": {
            "message": {"type": "string"},
            "citations": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {"exact_quote": {"type": "string"}},
                    "required": ["exact_quote"],
                },
            },
            "not_in_document": {"type": "boolean"},
            "not_in_document_explanation": {"type": "string"},
        },
        "required": ["message", "citations"],
    },
}


class ChatResult:
    def __init__(
        self,
        message: str,
        facts: list[Fact],
        not_in_document: bool | None,
        not_in_document_explanation: str | None,
        usage: dict,
    ):
        self.message = message
        self.facts = facts
        self.notInDocument = not_in_document
        self.notInDocumentExplanation = not_in_document_explanation
        self.usage = usage


async def chat_with_document(
    client: anthropic.AsyncAnthropic,
    doc_text: str,
    history: list[dict],
    message: str,
) -> ChatResult:
    truncated = doc_text[:MAX_TEXT_LENGTH]

    response = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=1024,
        system=[
            {"type": "text", "text": SYSTEM_PROMPT},
            {
                "type": "text",
                "text": f"<document>\n{truncated}\n</document>",
                "cache_control": {"type": "ephemeral"},
            },
        ],
        messages=[
            *[{"role": m["role"], "content": m["content"]} for m in history],
            {"role": "user", "content": message},
        ],
        tools=[RESPOND_TOOL],
        tool_choice={"type": "tool", "name": "respond_with_citations"},
    )

    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")

    parsed = tool_block.input
    raw_citations = parsed.get("citations", [])

    facts: list[Fact] = []
    if raw_citations:
        claude_facts = [
            ClaudeFact(
                fact=f"Citation {i + 1}",
                citations=[ClaudeCitation(exact_quote=c["exact_quote"])],
            )
            for i, c in enumerate(raw_citations)
        ]
        facts = map_citations(doc_text, claude_facts)

    u = response.usage
    cache_write = getattr(u, "cache_creation_input_tokens", 0) or 0
    cache_read = getattr(u, "cache_read_input_tokens", 0) or 0
    cost_usd = (
        u.input_tokens * INPUT_COST_PER_TOKEN
        + cache_write * INPUT_CACHE_WRITE_COST
        + cache_read * INPUT_CACHE_READ_COST
        + u.output_tokens * OUTPUT_COST_PER_TOKEN
    )

    return ChatResult(
        message=parsed["message"],
        facts=facts,
        not_in_document=parsed.get("not_in_document"),
        not_in_document_explanation=parsed.get("not_in_document_explanation"),
        usage={
            "inputTokens": u.input_tokens,
            "outputTokens": u.output_tokens,
            "cacheWriteTokens": cache_write,
            "cacheReadTokens": cache_read,
            "costUsd": cost_usd,
        },
    )
