import os
import json
import logging
import uuid
from typing import Any, Optional

import httpx
from fastmcp import FastMCP
from starlette.responses import JSONResponse

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    stream=__import__("sys").stderr,
)
logger = logging.getLogger(__name__)

mcp = FastMCP("kal")

BACKEND_API_BASE = os.getenv("BACKEND_API_URL", "http://localhost:5000/api")
MCP_API_KEY = os.getenv("MCP_API_KEY", "")
MCP_USER_ID = os.getenv("MCP_USER_ID", "")
USER_AGENT = "kal-mcp-server/2.0"

BACKEND_TIMEOUT = float(os.getenv("BACKEND_TIMEOUT", "180"))


# ---------------------------------------------------------------------------
# Health check
# ---------------------------------------------------------------------------

@mcp.custom_route("/health", methods=["GET"])
async def health_check(request):
    """Health check endpoint for Azure Container Apps liveness/readiness probes."""
    return JSONResponse({"status": "healthy", "service": "kal-mcp-server"})


# ---------------------------------------------------------------------------
# Backend HTTP helpers
# ---------------------------------------------------------------------------

async def make_backend_request(
    endpoint: str,
    method: str = "GET",
    params: Optional[dict] = None,
    json_body: Optional[dict] = None,
) -> dict[str, Any] | None:
    """Makes a request to the Kal backend API with auth headers."""
    headers: dict[str, str] = {
        "User-Agent": USER_AGENT,
        "Accept": "application/json",
    }
    if MCP_API_KEY:
        headers["Authorization"] = f"Bearer {MCP_API_KEY}"
    if MCP_USER_ID:
        headers["X-MCP-User-Id"] = MCP_USER_ID

    url = f"{BACKEND_API_BASE}{endpoint}"
    async with httpx.AsyncClient() as client:
        try:
            logger.info(f"Backend {method} {url}")
            response = await client.request(
                method,
                url,
                headers=headers,
                params=params,
                json=json_body,
                timeout=BACKEND_TIMEOUT,
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as e:
            logger.error(f"Backend HTTP {e.response.status_code}: {e.response.text[:200]}")
            return {"error": f"Backend returned {e.response.status_code}", "details": e.response.text[:500]}
        except httpx.TimeoutException:
            logger.error(f"Backend request timed out after {BACKEND_TIMEOUT}s")
            return {"error": "Backend request timed out. The query may be too complex -- try a simpler question."}
        except Exception as e:
            logger.error(f"Unexpected error calling backend: {e}")
            return {"error": f"Unexpected error: {e}"}


def _format_source_line(src: dict) -> str:
    """Format a structured source dict into a single readable line."""
    parts = []
    if src.get("source"):
        parts.append(f"Source: {src['source']}")
    if src.get("document_title"):
        parts.append(f"Title: {src['document_title']}")
    pages = src.get("pages") or src.get("page")
    if pages:
        parts.append(f"Pages: {pages}")
    if src.get("publishing_body"):
        parts.append(f"Publisher: {src['publishing_body']}")
    return " | ".join(parts) if parts else json.dumps(src)


def _check_error(data: dict | None, context: str) -> Optional[str]:
    """Return an error message string if the response indicates failure, else None."""
    if not data:
        return f"{context}: No response from backend."
    if data.get("error"):
        detail = data.get("details", data["error"])
        return f"{context}: {detail}"
    return None


# ---------------------------------------------------------------------------
# MCP Tools
# ---------------------------------------------------------------------------

@mcp.tool()
async def search_documents(
    query: str,
    target_documents: Optional[str] = None,
    target_tag: Optional[str] = None,
) -> str:
    """
    Search the Kal knowledge base for documents matching a query.
    Returns raw, graded retrieval results -- use this when you want to do
    your own analysis of the source material rather than getting a
    pre-synthesised answer.

    Results are filtered by the current user's document access permissions.

    Args:
        query: The search query (e.g. "earned value management techniques").
        target_documents: Optional comma-separated filenames to restrict
                          the search (e.g. "97R-18.pdf, BS6079-1.pdf").
        target_tag: Optional tag name to restrict search to a category
                    (e.g. "AACE", "Decision Intelligence").

    Returns:
        A text summary of search results with source citations, plus a
        confidence assessment and the number of results found.
    """
    payload: dict[str, Any] = {"query": query}
    if target_documents:
        docs = [d.strip() for d in target_documents.split(",") if d.strip()]
        payload["target_documents"] = docs
    if target_tag:
        payload["target_tag"] = target_tag

    data = await make_backend_request("/mcp/search", method="POST", json_body=payload)

    err = _check_error(data, "Search failed")
    if err:
        return err

    results = data.get("results", [])
    confidence = data.get("confidence", "unknown")
    total = data.get("total_results", 0)

    if not results:
        return f"No relevant results found for: {query}\nConfidence: {confidence}"

    lines = [f"Found {total} result(s) (confidence: {confidence}):\n"]
    for i, r in enumerate(results, 1):
        source_line = _format_source_line(r)
        snippet = (r.get("content") or "")[:500]
        lines.append(f"--- Result {i} ---")
        lines.append(source_line)
        lines.append(snippet)
        lines.append("")

    return "\n".join(lines)


@mcp.tool()
async def list_documents() -> str:
    """
    List all documents available in the Kal knowledge base for the
    current user. Returns document titles, references, publishers, and
    descriptions.

    Use this to discover what sources are available before searching.

    Returns:
        A formatted list of all accessible documents with metadata.
    """
    data = await make_backend_request("/mcp/documents", method="GET")

    err = _check_error(data, "Failed to list documents")
    if err:
        return err

    documents = data.get("documents", [])
    total = data.get("total", 0)

    if not documents:
        return "No documents are currently available."

    lines = [f"Available documents ({total} total):\n"]
    for doc in documents:
        title = doc.get("document_title") or doc.get("filename") or "Untitled"
        ref = doc.get("document_reference") or ""
        publisher = doc.get("publishing_body") or ""
        date = doc.get("publication_date") or ""
        desc = doc.get("description") or ""

        header = f"- **{title}**"
        if ref:
            header += f" [{ref}]"
        meta_parts = [p for p in [publisher, date] if p]
        if meta_parts:
            header += f" ({', '.join(meta_parts)})"
        lines.append(header)
        if desc:
            lines.append(f"  {desc[:200]}")

    return "\n".join(lines)


@mcp.tool()
async def answer_with_kal(query: str, session_id: str = "") -> str:
    """
    Ask the Kal agent for a fully researched, consolidated answer.
    Kal will plan a research strategy, search its document library,
    evaluate result quality (retrying if needed), and synthesise an
    answer with source citations.

    Use this when you want a complete answer. For raw search results,
    use search_documents instead.

    Args:
        query: The question to ask (e.g. "Compare AACE and BSI approaches
               to cost estimate classification").
        session_id: Optional session ID to maintain conversation context.
                    Leave empty for a standalone question.

    Returns:
        The synthesised answer with sources, confidence level, and any
        gaps or recommendations for follow-up.
    """
    if not session_id:
        session_id = f"mcp-{MCP_USER_ID or 'anon'}-{uuid.uuid4().hex[:12]}"

    payload: dict[str, Any] = {"query": query, "session_id": session_id}

    data = await make_backend_request("/mcp/chat", method="POST", json_body=payload)

    err = _check_error(data, "Chat failed")
    if err:
        return err

    answer = data.get("answer", "No answer received.")
    sources = data.get("sources", [])
    confidence = data.get("confidence", "unknown")
    gaps = data.get("gaps", [])
    recommendations = data.get("recommendations", [])
    plan = data.get("plan", {})

    parts = [answer]

    if sources:
        parts.append("\n**Sources:**")
        for s in sources:
            parts.append(f"- {_format_source_line(s)}")

    parts.append(f"\n**Confidence:** {confidence}")

    if plan.get("retry_count", 0) > 0:
        parts.append(f"**Retrieval passes:** {plan['retry_count'] + 1}")

    if gaps:
        parts.append("\n**Gaps:**")
        for g in gaps:
            parts.append(f"- {g}")

    if recommendations:
        parts.append("\n**Recommendations:**")
        for r in recommendations:
            parts.append(f"- {r}")

    return "\n".join(parts)


@mcp.tool()
async def get_document_content(source_filename: str) -> str:
    """
    Retrieve the full content of a specific document from the knowledge
    base. Returns the document's markdown text if available, or indicates
    that only a PDF is available.

    Args:
        source_filename: The filename of the document (e.g. "97R-18.pdf").

    Returns:
        The document's full text content, or an error if not found.
    """
    data = await make_backend_request(
        "/mcp/document/content",
        method="GET",
        params={"filename": source_filename},
    )

    err = _check_error(data, "Failed to retrieve document")
    if err:
        return err

    content_type = data.get("type", "")
    markdown = data.get("markdown_content")
    page_count = data.get("page_count")

    if markdown:
        header = f"Document: {source_filename}"
        if page_count:
            header += f" ({page_count} pages)"
        return f"{header}\n\n{markdown}"

    if content_type == "pdf":
        return (
            f"Document '{source_filename}' is available as PDF only "
            f"(no text extraction available). You can reference it but "
            f"cannot read its contents through this tool."
        )

    return f"No content found for '{source_filename}'."


# ---------------------------------------------------------------------------
# Server startup
# ---------------------------------------------------------------------------

app = mcp.http_app()


def main():
    import uvicorn

    port = int(os.getenv("PORT", "8080"))
    host = os.getenv("HOST", "0.0.0.0")
    logger.info(f"Starting Kal MCP server on {host}:{port}")
    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    main()
