from __future__ import annotations

import functools
import logging
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
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, load_settings, resolve_db_paths
from quarry.database import (
    count_chunks,
    delete_collection as db_delete_collection,
    delete_document as db_delete_document,
    discover_databases,
    get_db,
    get_page_text,
    list_collections as db_list_collections,
    list_documents,
    search,
)
from quarry.formatting import (
    format_collections,
    format_databases,
    format_document_detail,
    format_documents,
    format_registrations,
    format_search_results,
    format_status,
    format_switch_summary,
)
from quarry.pipeline import (
    ingest_auto as pipeline_ingest_auto,
    ingest_content as pipeline_ingest_content,
    ingest_document,
)
from quarry.sync import sync_all as engine_sync_all
from quarry.sync_registry import (
    deregister_directory as registry_deregister,
    list_registrations as registry_list,
    open_registry,
    register_directory as registry_register,
)
from quarry.types import LanceDB

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

mcp = FastMCP(
    "punt-quarry",
    instructions=(
        "Quarry is a local semantic search engine for your knowledge base. "
        "Use these tools to ingest, search, and manage documents.\n\n"
        "All quarry tool output is pre-formatted plain text using unicode "
        "characters for alignment. Always emit quarry output verbatim \u2014 "
        "never reformat, never convert to markdown tables, never wrap "
        "in code fences or boxes."
    ),
)

_db_name: str | None = None


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 resolve_db_paths(load_settings(), _db_name)


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


_executor = ThreadPoolExecutor(max_workers=4)


def _background(fn: Callable[..., object], *args: object, **kwargs: object) -> None:
    """Run *fn* in the bounded thread pool.  Exceptions are logged, not raised."""

    def _target() -> None:
        try:
            fn(*args, **kwargs)
        except Exception:
            logger.exception("Background task %s failed", fn.__name__)

    _executor.submit(_target)


@mcp.tool()
@_handle_errors
def find(
    query: str,
    limit: int = 10,
    document_filter: str = "",
    collection: str = "",
    page_type: str = "",
    source_format: 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.
        page_type: Optional content type filter (text, code, spreadsheet, etc.).
        source_format: Optional source format filter (.pdf, .py, .xlsx, etc.).
    """
    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,
        page_type_filter=page_type or None,
        source_format_filter=source_format or None,
    )

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

    return format_search_results(query, formatted)


def _do_ingest(
    source: str, overwrite: bool, collection: str, settings: Settings, db: LanceDB
) -> None:
    """Blocking ingest — runs in background thread."""
    if source.startswith(("http://", "https://")):
        pipeline_ingest_auto(
            source, db, settings, overwrite=overwrite, collection=collection or ""
        )
        return

    path = Path(source)
    col = derive_collection(path, explicit=collection or None)
    ingest_document(path, db, settings, overwrite=overwrite, collection=col)


@mcp.tool()
@_handle_errors
def ingest(
    source: str,
    overwrite: bool = False,
    collection: str = "",
) -> str:
    """Ingest a file or URL into the knowledge base.

    Auto-detects the source type: URLs (http/https) use smart sitemap
    discovery with single-page fallback; local paths are ingested as files.

    Supported formats: PDF, images (PNG, JPG, TIFF, BMP, WebP), presentations
    (PPTX), spreadsheets (XLSX, CSV), HTML, TXT, MD, TEX, DOCX, and source
    code files.

    Returns immediately — ingestion runs in the background.

    Args:
        source: File path or HTTP(S) URL to ingest.
        overwrite: If true, replace existing data.
        collection: Collection name. Auto-derived if empty.
    """
    settings = _settings()
    db = _db()
    _background(_do_ingest, source, overwrite, collection, settings, db)
    return f"\u25b6  Ingesting {source} (background)"


def _do_remember(
    content: str,
    document_name: str,
    overwrite: bool,
    collection: str,
    format_hint: str,
    settings: Settings,
    db: LanceDB,
) -> None:
    """Blocking remember — runs in background thread."""
    pipeline_ingest_content(
        content,
        document_name,
        db,
        settings,
        overwrite=overwrite,
        collection=collection,
        format_hint=format_hint,
    )


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

    Use this instead of ingest when you have the text content directly
    (e.g., clipboard, API response, or sandbox-uploaded files in Claude Desktop).

    Returns immediately — indexing runs in the background.

    Args:
        content: The text content to remember.
        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()
    _background(
        _do_remember,
        content,
        document_name,
        overwrite,
        collection,
        format_hint,
        settings,
        db,
    )
    return f"\u25b6  Remembering {document_name} (background)"


@mcp.tool(name="list")
@_handle_errors
def list_resources(
    kind: str,
    collection: str = "",
) -> str:
    """List documents, collections, databases, or registrations.

    Args:
        kind: What to list — "documents", "collections", "databases",
              or "registrations".
        collection: Optional collection filter (only for kind="documents").
    """
    if kind == "documents":
        db = _db()
        docs = list_documents(db, collection_filter=collection or None)
        return format_documents(docs)
    if kind == "collections":
        db = _db()
        cols = db_list_collections(db)
        return format_collections(cols)
    if kind == "databases":
        settings = _settings()
        databases = discover_databases(settings.quarry_root)
        return format_databases(databases, current=_db_name or "default")
    if kind == "registrations":
        settings = _settings()
        conn = open_registry(settings.registry_path)
        try:
            regs = registry_list(conn)
        finally:
            conn.close()
        return format_registrations(
            [
                {
                    "directory": r.directory,
                    "collection": r.collection,
                    "registered_at": r.registered_at,
                }
                for r in regs
            ]
        )
    return (
        f"Error: unknown kind {kind!r}. "
        "Use documents, collections, databases, or registrations."
    )


@mcp.tool()
@_handle_errors
def show(
    document_name: str,
    page_number: int = 0,
    collection: str = "",
) -> str:
    """Show document metadata or retrieve a specific page's text.

    Without page_number: shows document metadata (pages, chunks, collection).
    With page_number: shows the full text for that page.

    Args:
        document_name: Document filename (e.g., 'report.pdf').
        page_number: Page number (1-indexed). 0 means show metadata only.
        collection: Optional collection scope.
    """
    db = _db()

    if page_number > 0:
        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}"

    docs = list_documents(db, collection_filter=collection or None)
    match = [d for d in docs if d["document_name"] == document_name]
    if not match:
        return f"Document {document_name!r} not found"
    return format_document_detail(match[0])


def _do_delete(name: str, kind: str, collection: str, db: LanceDB) -> None:
    """Blocking delete — runs in background thread."""
    if kind == "collection":
        db_delete_collection(db, name)
    else:
        db_delete_document(db, name, collection=collection or None)


@mcp.tool()
@_handle_errors
def delete(
    name: str,
    kind: str = "document",
    collection: str = "",
) -> str:
    """Delete indexed data for a document or collection.

    Returns immediately — deletion runs in the background.

    Args:
        name: Document filename or collection name to delete.
        kind: What to delete — "document" or "collection".
        collection: Optional collection scope (only for kind="document").
    """
    if kind not in ("document", "collection"):
        return f"Invalid kind {kind!r}. Must be 'document' or 'collection'."
    db = _db()
    _background(_do_delete, name, kind, collection, db)
    return f"\u25b6  Deleting {kind} {name!r} (background)"


def _do_register(directory: str, collection: str, settings: Settings) -> None:
    """Blocking register — runs in background thread."""
    path = Path(directory).resolve()
    col = collection or path.name
    conn = open_registry(settings.registry_path)
    try:
        registry_register(conn, path, col)
    finally:
        conn.close()


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

    Returns immediately — registration runs in the background.

    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
    _background(_do_register, directory, collection, settings)
    return f"\u25b6  Registering {path} as {col!r} (background)"


def _do_deregister(
    collection: str, keep_data: bool, settings: Settings, db: LanceDB
) -> None:
    """Blocking deregister — runs in background thread."""
    conn = open_registry(settings.registry_path)
    try:
        doc_names = registry_deregister(conn, collection)
    finally:
        conn.close()

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


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

    Returns immediately — deregistration runs in the background.

    Args:
        collection: Collection name to deregister.
        keep_data: If true, keep indexed data in LanceDB.
    """
    settings = _settings()
    db = _db()
    _background(_do_deregister, collection, keep_data, settings, db)
    return f"\u25b6  Deregistering {collection!r} (background)"


def _do_sync(settings: Settings, db: LanceDB) -> None:
    """Blocking sync — runs in background thread."""
    engine_sync_all(db, settings)


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

    Returns immediately — sync runs in the background.
    """
    settings = _settings()
    db = _db()
    _background(_do_sync, settings, db)
    return "\u25b6  Syncing all registrations (background)"


@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 format_status(
        {
            "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,
        }
    )


@mcp.tool(name="use")
@_handle_errors
def use_database(name: str) -> str:
    """Switch to a different named database for subsequent operations.

    All tools (find, ingest, sync, etc.) will use the selected database
    until changed again. Use list(kind="databases") to see available databases.

    Args:
        name: Database name (e.g., 'coding', 'work'). Use 'default' for
              the default database.
    """
    global _db_name
    previous = _db_name or "default"
    new_name = name if name != "default" else None
    # Validate before mutating: resolve_db_paths raises ValueError for
    # names containing path separators or traversal segments.
    test_settings = resolve_db_paths(load_settings(), new_name)
    _db_name = new_name
    return format_switch_summary(previous, name, str(test_settings.lancedb_path))


def main(db_name: str | None = None) -> None:
    global _db_name
    _db_name = db_name
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
