"""SAIA Docling MCP server — document → Markdown/HTML/JSON conversion.

Wraps GWDG SAIA's Docling endpoint (POST /v1/documents/convert).

Async queue: large PDFs (OCR + image extraction) can take minutes and would
blow past LLM client tool-call timeouts. `convert_document` submits the job and
waits up to CONVERT_DEADLINE seconds inline; if it's still running it returns a
job_id to poll with `get_job_status`. The conversion keeps running in the
background regardless (asyncio.shield), so the result is never lost.

Defaults chosen for LLM/RAG ingestion:
  response_type            = markdown   (best for embedding / chat context)
  extract_tables_as_images = False      (tables stay structured, not pictures)
  image_resolution_scale   = 2          (images kept, reasonable size)
"""
import asyncio
import base64
import os
import time
import uuid
from urllib.parse import urlparse

import httpx
from mcp.server.fastmcp import FastMCP

API_KEY = os.environ.get("SAIA_API_KEY", "")
API_ENDPOINT = os.environ.get(
    "DOCLING_ENDPOINT", "https://chat-ai.academiccloud.de/v1/documents/convert"
)
CONVERT_DEADLINE = int(os.environ.get("CONVERT_DEADLINE", 60))
_VALID_TYPES = ("markdown", "html", "json", "tokens")

mcp = FastMCP("saia-docling", host="0.0.0.0", port=8000, stateless_http=True)

# job_id -> {status, result, error, submitted_at, filename}
_jobs: dict[str, dict] = {}
_tasks: set[asyncio.Task] = set()


async def _do_convert(filename: str, content: bytes, params: dict) -> dict:
    headers = {"accept": "application/json"}
    if API_KEY:
        headers["Authorization"] = f"Bearer {API_KEY}"
    async with httpx.AsyncClient(timeout=httpx.Timeout(600, connect=10)) as c:
        r = await c.post(
            API_ENDPOINT,
            params=params,
            headers=headers,
            files={"document": (filename, content)},
        )
    if r.status_code >= 400:
        raise RuntimeError(f"Docling API {r.status_code}: {r.text[:400]}")
    return r.json()


async def _run_job(job_id, filename, content, params, response_type):
    try:
        data = await _do_convert(filename, content, params)
        _jobs[job_id].update(
            status="done",
            result={
                "response_type": data.get("response_type", response_type.upper()),
                "filename": data.get("filename", filename),
                response_type: data.get(response_type, data.get("markdown", "")),
                "images": data.get("images", []),
            },
        )
    except Exception as e:  # noqa: BLE001
        _jobs[job_id].update(status="error", error=str(e)[:500])


def _result(job_id: str) -> dict:
    j = _jobs[job_id]
    if j["status"] == "done":
        return {"job_id": job_id, "status": "done", **j["result"]}
    if j["status"] == "error":
        return {"job_id": job_id, "status": "error", "error": j["error"]}
    return {
        "job_id": job_id,
        "status": "running",
        "elapsed_s": int(time.time() - j["submitted_at"]),
    }


@mcp.tool()
async def convert_document(
    source: str,
    filename: str = "document",
    response_type: str = "markdown",
    extract_tables_as_images: bool = False,
    image_resolution_scale: int = 2,
) -> dict:
    """Convert a document (PDF, DOCX, PPTX, images, …) to Markdown/HTML/JSON.

    Submits the job and waits inline up to ~60s. If conversion takes longer
    (big/scanned PDFs), returns a job_id — call `get_job_status(job_id)` to
    retrieve the result when ready (it keeps processing in the background).

    Args:
        source: An http(s) URL to the document, OR its base64-encoded bytes.
        filename: Name (with extension) for the upload; aids format detection.
        response_type: "markdown" (default), "html", "json", or "tokens".
        extract_tables_as_images: False (default) keeps tables as structured
            text; True returns them as images.
        image_resolution_scale: Image scaling factor 1–4 (default 2).

    Returns:
        dict with status "done" (incl. the converted text under its format key
        plus "images"), or status "running" with a job_id to poll.
    """
    if response_type not in _VALID_TYPES:
        raise ValueError(f"response_type must be one of {_VALID_TYPES}")
    if image_resolution_scale not in (1, 2, 3, 4):
        raise ValueError("image_resolution_scale must be 1, 2, 3, or 4")

    if source.startswith(("http://", "https://")):
        async with httpx.AsyncClient(timeout=60, follow_redirects=True) as c:
            resp = await c.get(source)
            resp.raise_for_status()
            content = resp.content
        if filename == "document":
            filename = urlparse(source).path.split("/")[-1] or "document"
    else:
        try:
            content = base64.b64decode(source, validate=True)
        except Exception as e:
            raise ValueError(
                "source is neither an http(s) URL nor valid base64"
            ) from e

    params = {
        "response_type": response_type,
        "extract_tables_as_images": str(extract_tables_as_images).lower(),
        "image_resolution_scale": str(image_resolution_scale),
    }

    job_id = f"doc-{uuid.uuid4().hex[:8]}"
    _jobs[job_id] = {
        "status": "running", "result": None, "error": None,
        "submitted_at": time.time(), "filename": filename,
    }
    if len(_jobs) > 200:
        for k in sorted(_jobs, key=lambda k: _jobs[k]["submitted_at"])[:50]:
            _jobs.pop(k, None)

    task = asyncio.create_task(
        _run_job(job_id, filename, content, params, response_type)
    )
    _tasks.add(task)
    task.add_done_callback(_tasks.discard)

    try:
        await asyncio.wait_for(asyncio.shield(task), timeout=CONVERT_DEADLINE)
    except asyncio.TimeoutError:
        return {
            "job_id": job_id,
            "status": "running",
            "note": (
                f"'{filename}' still converting after {CONVERT_DEADLINE}s. "
                f"Call get_job_status('{job_id}') to retrieve it when ready."
            ),
        }
    return _result(job_id)


@mcp.tool()
async def get_job_status(job_id: str) -> dict:
    """Retrieve a previously submitted conversion job by its job_id."""
    if job_id not in _jobs:
        return {"error": f"unknown job_id '{job_id}'"}
    return _result(job_id)


if __name__ == "__main__":
    mcp.run(transport="streamable-http")