from __future__ import annotations

import functools
import json
import logging
from collections.abc import Callable
from pathlib import Path

from mcp.server.fastmcp import FastMCP

from quarry.backends import get_embedding_backend
from quarry.collections import derive_collection
from quarry.config import Settings, configure_logging, get_settings
from quarry.database import (
    count_chunks,
    delete_collection as db_delete_collection,
    delete_document as db_delete_document,
    get_db,
    get_page_text,
    list_collections as db_list_collections,
    list_documents,
    search,
)
from quarry.pipeline import ingest_document, ingest_text as pipeline_ingest_text
from quarry.registry import (
    deregister_directory as registry_deregister,
    list_registrations as registry_list,
    open_registry,
    register_directory as registry_register,
)
from quarry.sync import sync_all as engine_sync_all
from quarry.types import LanceDB

configure_logging(get_settings())
logger = logging.getLogger(__name__)

mcp = FastMCP("quarry-mcp")


def _handle_errors(fn: Callable[..., str]) -> Callable[..., str]:
    """Catch exceptions at the MCP boundary, log, and return error string."""

    @functools.wraps(fn)
    def wrapper(*args: object, **kwargs: object) -> str:
        try:
            return fn(*args, **kwargs)
        except Exception as exc:
            logger.exception("Error in %s", fn.__name__)
            return f"Error: {type(exc).__name__}: {exc}"

    return wrapper


def _settings() -> Settings:
    return get_settings()


def _db() -> LanceDB:
    return get_db(_settings().lancedb_path)


@mcp.tool()
@_handle_errors
def search_documents(
    query: str,
    limit: int = 10,
    document_filter: str = "",
    collection: str = "",
) -> str:
    """Search indexed documents using semantic similarity.

    Args:
        query: Natural language search query.
        limit: Maximum number of results (default 10, max 50).
        document_filter: Optional exact document name to filter by.
        collection: Optional collection name to search within.
    """
    limit = min(limit, 50)
    settings = _settings()
    db = _db()

    query_vector = get_embedding_backend(settings).embed_query(query)

    results = search(
        db,
        query_vector,
        limit=limit,
        document_filter=document_filter or None,
        collection_filter=collection or None,
    )

    formatted = [
        {
            "document_name": r["document_name"],
            "collection": r["collection"],
            "page_number": r["page_number"],
            "chunk_index": r["chunk_index"],
            "text": r["text"],
            "similarity": round(1 - float(str(r.get("_distance", 0))), 4),
        }
        for r in results
    ]

    return json.dumps(
        {
            "query": query,
            "total_results": len(formatted),
            "results": formatted,
        },
        indent=2,
    )


@mcp.tool()
@_handle_errors
def ingest(
    file_path: str,
    overwrite: bool = False,
    collection: str = "",
) -> str:
    """Ingest a document: OCR, chunk, embed, and index for search.

    Supported formats: PDF, TXT, MD, TEX, DOCX.

    Args:
        file_path: Absolute path to the document file.
        overwrite: If true, replace existing data for this document.
        collection: Collection name. Auto-derived from parent directory if empty.
    """
    path = Path(file_path)
    settings = _settings()
    db = _db()
    col = derive_collection(path, explicit=collection or None)

    progress_lines: list[str] = []

    result = ingest_document(
        path,
        db,
        settings,
        overwrite=overwrite,
        collection=col,
        progress_callback=progress_lines.append,
    )

    progress_lines.append("")
    progress_lines.append(f"Result: {json.dumps(result, indent=2)}")
    return "\n".join(progress_lines)


@mcp.tool()
@_handle_errors
def ingest_text(
    content: str,
    document_name: str,
    overwrite: bool = False,
    collection: str = "default",
    format_hint: str = "auto",
) -> str:
    """Ingest raw text content: chunk, embed, and index for search.

    Args:
        content: The text content to ingest.
        document_name: Name for the document (e.g., 'notes.md').
        overwrite: If true, replace existing data for this document.
        collection: Collection name (default: 'default').
        format_hint: Format hint: 'auto', 'plain', 'markdown', 'latex'.
    """
    settings = _settings()
    db = _db()

    progress_lines: list[str] = []

    result = pipeline_ingest_text(
        content,
        document_name,
        db,
        settings,
        overwrite=overwrite,
        collection=collection,
        format_hint=format_hint,
        progress_callback=progress_lines.append,
    )

    progress_lines.append("")
    progress_lines.append(f"Result: {json.dumps(result, indent=2)}")
    return "\n".join(progress_lines)


@mcp.tool()
@_handle_errors
def get_documents(collection: str = "") -> str:
    """List all indexed documents with metadata.

    Args:
        collection: Optional collection name to filter by.
    """
    db = _db()
    docs = list_documents(db, collection_filter=collection or None)
    return json.dumps({"total_documents": len(docs), "documents": docs}, indent=2)


@mcp.tool()
@_handle_errors
def get_page(
    document_name: str,
    page_number: int,
    collection: str = "",
) -> str:
    """Retrieve the full raw OCR text for a specific document page.

    Args:
        document_name: Document filename (e.g., 'report.pdf').
        page_number: Page number (1-indexed).
        collection: Optional collection scope.
    """
    db = _db()
    text = get_page_text(db, document_name, page_number, collection=collection or None)

    if text is None:
        return f"No data found for {document_name} page {page_number}"

    return f"Document: {document_name}\nPage: {page_number}\n---\n{text}"


@mcp.tool()
@_handle_errors
def delete_document(
    document_name: str,
    collection: str = "",
) -> str:
    """Delete all indexed data for a document.

    Args:
        document_name: Document filename (e.g., 'report.pdf').
        collection: Optional collection scope. If empty, deletes across all collections.
    """
    db = _db()
    deleted = db_delete_document(db, document_name, collection=collection or None)
    return json.dumps(
        {
            "document_name": document_name,
            "collection": collection or "(all)",
            "chunks_deleted": deleted,
        },
        indent=2,
    )


@mcp.tool()
@_handle_errors
def list_collections() -> str:
    """List all collections with document and chunk counts."""
    db = _db()
    cols = db_list_collections(db)
    return json.dumps(
        {"total_collections": len(cols), "collections": cols},
        indent=2,
    )


@mcp.tool()
@_handle_errors
def delete_collection(collection: str) -> str:
    """Delete all indexed data for a collection.

    Args:
        collection: Collection name to delete.
    """
    db = _db()
    deleted = db_delete_collection(db, collection)
    return json.dumps(
        {
            "collection": collection,
            "chunks_deleted": deleted,
        },
        indent=2,
    )


@mcp.tool()
@_handle_errors
def register_directory(directory: str, collection: str = "") -> str:
    """Register a directory for incremental sync.

    Args:
        directory: Absolute path to the directory.
        collection: Collection name. Uses directory name if empty.
    """
    settings = _settings()
    path = Path(directory).resolve()
    col = collection or path.name
    conn = open_registry(settings.registry_path)
    try:
        reg = registry_register(conn, path, col)
    finally:
        conn.close()
    return json.dumps(
        {
            "directory": reg.directory,
            "collection": reg.collection,
            "registered_at": reg.registered_at,
        },
        indent=2,
    )


@mcp.tool()
@_handle_errors
def deregister_directory(
    collection: str,
    keep_data: bool = False,
) -> str:
    """Remove a directory registration.

    Args:
        collection: Collection name to deregister.
        keep_data: If true, keep indexed data in LanceDB.
    """
    settings = _settings()
    conn = open_registry(settings.registry_path)
    try:
        doc_names = registry_deregister(conn, collection)
    finally:
        conn.close()

    if not keep_data and doc_names:
        db = _db()
        for name in doc_names:
            db_delete_document(db, name, collection=collection)

    return json.dumps(
        {
            "collection": collection,
            "documents_removed": len(doc_names),
            "data_deleted": not keep_data,
        },
        indent=2,
    )


@mcp.tool()
@_handle_errors
def sync_all_registrations() -> str:
    """Sync all registered directories: ingest new/changed, remove deleted."""
    settings = _settings()
    db = _db()

    progress_lines: list[str] = []
    results = engine_sync_all(
        db,
        settings,
        progress_callback=progress_lines.append,
    )

    summary = {
        col: {
            "ingested": res.ingested,
            "deleted": res.deleted,
            "skipped": res.skipped,
            "failed": res.failed,
            "errors": res.errors,
        }
        for col, res in results.items()
    }

    return json.dumps(
        {
            "collections_synced": len(results),
            "results": summary,
            "progress": progress_lines,
        },
        indent=2,
    )


@mcp.tool()
@_handle_errors
def list_registrations() -> str:
    """List all registered directories."""
    settings = _settings()
    conn = open_registry(settings.registry_path)
    try:
        regs = registry_list(conn)
    finally:
        conn.close()
    return json.dumps(
        {
            "total_registrations": len(regs),
            "registrations": [
                {
                    "directory": r.directory,
                    "collection": r.collection,
                    "registered_at": r.registered_at,
                }
                for r in regs
            ],
        },
        indent=2,
    )


@mcp.tool()
@_handle_errors
def status() -> str:
    """Get database status: document/chunk counts, storage size, and model info."""
    settings = _settings()
    db = _db()

    docs = list_documents(db)
    chunks = count_chunks(db)
    cols = db_list_collections(db)

    if settings.registry_path.exists():
        conn = open_registry(settings.registry_path)
        try:
            regs = registry_list(conn)
        finally:
            conn.close()
    else:
        regs = []

    db_size_bytes = (
        sum(f.stat().st_size for f in settings.lancedb_path.rglob("*") if f.is_file())
        if settings.lancedb_path.exists()
        else 0
    )

    return json.dumps(
        {
            "document_count": len(docs),
            "collection_count": len(cols),
            "chunk_count": chunks,
            "registered_directories": len(regs),
            "database_path": str(settings.lancedb_path),
            "database_size_bytes": db_size_bytes,
            "embedding_model": settings.embedding_model,
            "embedding_dimension": settings.embedding_dimension,
        },
        indent=2,
    )


def main() -> None:
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
