"""
ai-citer MCP server (Python).

Run with:
    python -m mcp_server
or via the package script:
    python mcp_server.py
"""
from __future__ import annotations
import asyncio
import base64
import json
import sys
import uuid
from datetime import datetime, timezone

import anthropic
import httpx
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import TextContent, Tool

from ai_citer.config import settings
from ai_citer.db.client import create_pool
from ai_citer.db.documents import get_document_by_id, insert_document, list_documents
from ai_citer.db.extractions import get_latest_extraction_by_document_id, insert_extraction_result
from ai_citer.db.init import init_db
from ai_citer.ai.claude import extract_facts
from ai_citer.ai.citation_mapper import map_citations
from ai_citer.ai.post_process import assign_page_numbers
from ai_citer.models import Document
from ai_citer.parsers.pdf import parse_pdf
from ai_citer.parsers.web import parse_web

server = Server("ai-citer")

# State populated in main()
_pool = None
_client: anthropic.AsyncAnthropic | None = None


@server.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="upload_document_url",
            description="Fetch a URL (web page or PDF) and ingest it into ai-citer for later extraction.",
            inputSchema={
                "type": "object",
                "properties": {
                    "url": {"type": "string", "description": "The URL to fetch and ingest"},
                },
                "required": ["url"],
            },
        ),
        Tool(
            name="extract_facts",
            description="Run AI fact extraction on an ingested document.",
            inputSchema={
                "type": "object",
                "properties": {
                    "document_id": {"type": "string", "description": "ID of the document to extract from"},
                    "prompt": {
                        "type": "string",
                        "description": "Optional focus prompt to guide extraction (e.g. 'financial figures')",
                    },
                },
                "required": ["document_id"],
            },
        ),
        Tool(
            name="get_facts",
            description="Retrieve the latest extraction results for a document.",
            inputSchema={
                "type": "object",
                "properties": {
                    "document_id": {"type": "string", "description": "ID of the document"},
                },
                "required": ["document_id"],
            },
        ),
        Tool(
            name="list_documents",
            description="List all ingested documents.",
            inputSchema={
                "type": "object",
                "properties": {},
                "required": [],
            },
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    try:
        if name == "upload_document_url":
            url: str = arguments.get("url", "")
            if not url:
                raise ValueError("url is required")

            doc_id = str(uuid.uuid4())
            doc: Document

            is_pdf = url.lower().endswith(".pdf")
            if not is_pdf:
                try:
                    async with httpx.AsyncClient(follow_redirects=True, timeout=10.0) as hx:
                        head_res = await hx.head(url)
                    ct = head_res.headers.get("content-type", "")
                    is_pdf = "application/pdf" in ct
                except Exception:
                    pass

            if is_pdf:
                async with httpx.AsyncClient(follow_redirects=True, timeout=30.0) as hx:
                    pdf_res = await hx.get(url)
                if pdf_res.status_code >= 400:
                    raise RuntimeError(f"Failed to fetch PDF: {pdf_res.status_code}")
                buffer = pdf_res.content
                content = parse_pdf(buffer)
                pdf_data = base64.b64encode(buffer).decode()
                title = url.split("/")[-1] or url
                doc = Document(
                    id=doc_id,
                    title=title,
                    type="pdf",
                    content=content,
                    sourceUrl=url,
                    pdfData=pdf_data,
                    createdAt=datetime.now(timezone.utc).isoformat(),
                )
            else:
                title, content = await parse_web(url)
                doc = Document(
                    id=doc_id,
                    title=title,
                    type="web",
                    content=content,
                    sourceUrl=url,
                    createdAt=datetime.now(timezone.utc).isoformat(),
                )

            await insert_document(_pool, doc)
            return [TextContent(type="text", text=json.dumps({"id": doc.id, "title": doc.title, "type": doc.type}))]

        if name == "extract_facts":
            document_id: str = arguments.get("document_id", "")
            prompt: str | None = arguments.get("prompt") or None
            if not document_id:
                raise ValueError("document_id is required")

            doc = await get_document_by_id(_pool, document_id)
            if doc is None:
                raise RuntimeError(f"Document not found: {document_id}")

            extraction, _usage = await extract_facts(_client, doc.content.rawText, prompt)
            facts = map_citations(doc.content.rawText, extraction.facts)
            assign_page_numbers(doc, facts)

            result = await insert_extraction_result(
                _pool,
                doc.id,
                facts,
                prompt,
                extraction.no_mention_found,
                extraction.absence_evidence,
            )
            return [
                TextContent(
                    type="text",
                    text=json.dumps({
                        "id": result.id,
                        "documentId": result.documentId,
                        "factCount": len(result.facts),
                        "extractedAt": result.extractedAt,
                    }),
                )
            ]

        if name == "get_facts":
            document_id = arguments.get("document_id", "")
            if not document_id:
                raise ValueError("document_id is required")

            result = await get_latest_extraction_by_document_id(_pool, document_id)
            if result is None:
                raise RuntimeError(f"No extraction results found for document: {document_id}")

            return [TextContent(type="text", text=json.dumps(result.model_dump()))]

        if name == "list_documents":
            documents = await list_documents(_pool)
            return [TextContent(type="text", text=json.dumps([d.model_dump() for d in documents]))]

        raise ValueError(f"Unknown tool: {name}")

    except Exception as exc:
        return [TextContent(type="text", text=f"Error: {exc}")]


async def main() -> None:
    global _pool, _client

    _pool = await create_pool(settings.database_url)
    _client = anthropic.AsyncAnthropic(api_key=settings.anthropic_api_key)
    await init_db(_pool)

    sys.stderr.write("ai-citer MCP server started\n")
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


if __name__ == "__main__":
    asyncio.run(main())
