"""FastMCP server for ARS Title 25 structured RAG."""

import logging

from mcp.server.fastmcp import FastMCP

from ars25.core.responses import (
    success_response,
    error_response,
    validation_error,
    not_found_error,
    internal_error,
    ErrorType,
)
from ars25.core.db import table_exists, get_table
from ars25.core.config import get_settings

logger = logging.getLogger(__name__)

mcp = FastMCP(
    "ARS-25",
    instructions="Structured RAG server for Arizona Revised Statutes Title 25 (Family Law)",
)


# ─── Search & Retrieval ──────────────────────────────────────────


@mcp.tool()
def search_statutes(
    query: str,
    top_k: int = 10,
    chapter: int | None = None,
    article: int | None = None,
) -> dict:
    """Search Arizona family law statutes using hybrid vector + full-text search.

    Auto-detects statute references (e.g. "25-403") and fetches them directly.
    Returns ranked chunks with relevance scores.

    Args:
        query: Natural language search query or statute reference
        top_k: Number of results to return (default 10)
        chapter: Optional chapter number filter (1-8)
        article: Optional article number filter
    """
    try:
        if not table_exists():
            return error_response(
                "Database not initialized. Run ingest_statutes first.",
                ErrorType.NOT_FOUND,
            )

        from ars25.search.hybrid import search

        results = search(query, top_k=top_k, chapter=chapter, article=article)

        if not results:
            return success_response(
                results=[],
                count=0,
                query=query,
                message="No results found.",
            )

        return success_response(
            results=[r.to_dict() for r in results],
            count=len(results),
            query=query,
        )
    except Exception as e:
        logger.exception("search_statutes failed")
        return internal_error(str(e))


@mcp.tool()
def get_statute(section: str) -> dict:
    """Look up a specific Arizona statute by section number.

    Returns the full statute text with all chunks concatenated.

    Args:
        section: Statute section number (e.g. "25-403", "§ 25-403.01", "A.R.S. 25-411")
    """
    try:
        if not table_exists():
            return error_response(
                "Database not initialized. Run ingest_statutes first.",
                ErrorType.NOT_FOUND,
            )

        from ars25.search.lookup import get_statute as _get_statute, sanitize_section

        clean = sanitize_section(section)
        if not clean:
            return validation_error(
                f"Invalid section format: '{section}'. Expected format: 25-NNN or 25-NNN.NN",
            )

        result = _get_statute(section)
        if not result:
            return not_found_error(f"Statute section {clean} not found in database.")

        return success_response(**result)
    except Exception as e:
        logger.exception("get_statute failed")
        return internal_error(str(e))


@mcp.tool()
def list_statutes(
    chapter: int | None = None,
    article: int | None = None,
) -> dict:
    """List all Arizona family law statute sections and titles.

    Returns metadata only (no full text). Optionally filter by chapter or article.

    Args:
        chapter: Optional chapter number filter (1-8)
        article: Optional article number filter
    """
    try:
        if not table_exists():
            return error_response(
                "Database not initialized. Run ingest_statutes first.",
                ErrorType.NOT_FOUND,
            )

        from ars25.search.lookup import list_statutes as _list_statutes

        statutes = _list_statutes(chapter=chapter, article=article)

        return success_response(
            statutes=statutes,
            count=len(statutes),
        )
    except Exception as e:
        logger.exception("list_statutes failed")
        return internal_error(str(e))


@mcp.tool()
def get_chapter_tree() -> dict:
    """Get the hierarchical structure of ARS Title 25.

    Returns chapters → articles → statutes with summary counts.
    No parameters needed.
    """
    try:
        if not table_exists():
            return error_response(
                "Database not initialized. Run ingest_statutes first.",
                ErrorType.NOT_FOUND,
            )

        from ars25.search.lookup import get_chapter_tree as _get_chapter_tree

        tree = _get_chapter_tree()

        total_statutes = sum(
            s_count
            for ch in tree
            for art in ch["articles"]
            for s_count in [art["statute_count"]]
        )

        return success_response(
            chapters=tree,
            chapter_count=len(tree),
            total_statutes=total_statutes,
        )
    except Exception as e:
        logger.exception("get_chapter_tree failed")
        return internal_error(str(e))


# ─── Ingestion & Status ──────────────────────────────────────────


@mcp.tool()
async def ingest_statutes(force: bool = False) -> dict:
    """Run the full ingestion pipeline: parse CSV → scrape azleg.gov → chunk → embed → store.

    Takes ~3-5 minutes on first run. Skips if database already exists unless force=True.

    Args:
        force: If True, rebuild the database even if it already exists
    """
    try:
        from ars25.ingestion.pipeline import run_ingestion

        result = await run_ingestion(force=force)

        if result.get("skipped"):
            return success_response(
                skipped=True,
                message=result["reason"],
            )

        if result.get("error"):
            return error_response(result["error"], ErrorType.INGESTION)

        return success_response(
            message="Ingestion complete",
            statutes_scraped=result["statutes_scraped"],
            statutes_total=result["statutes_total"],
            chunks_created=result["chunks_created"],
            fts_index=result["fts_index"],
            db_path=result["db_path"],
        )
    except Exception as e:
        logger.exception("ingest_statutes failed")
        return internal_error(str(e))


@mcp.tool()
def get_ingestion_status() -> dict:
    """Check the current state of the statute database.

    Reports whether the DB exists, chunk count, unique statutes,
    schema info, and embedding model details.
    """
    try:
        settings = get_settings()

        if not table_exists():
            return success_response(
                initialized=False,
                message="Database not initialized. Run ingest_statutes to build it.",
                db_path=settings.db_path,
                embedding_model=settings.embedding_model,
            )

        table = get_table()
        row_count = table.count_rows()

        # Get unique statute count
        rows = table.search().select(["section"]).limit(10000).to_list()
        unique_sections = len(set(r["section"] for r in rows))

        # Get schema info
        schema_fields = [
            {"name": f.name, "type": str(f.type)}
            for f in table.schema
            if f.name != "vector"
        ]

        return success_response(
            initialized=True,
            chunk_count=row_count,
            unique_statutes=unique_sections,
            schema=schema_fields,
            db_path=settings.db_path,
            embedding_model=settings.embedding_model,
            embedding_dim=settings.embedding_dim,
        )
    except Exception as e:
        logger.exception("get_ingestion_status failed")
        return internal_error(str(e))


# ─── Entry point ─────────────────────────────────────────────────


def main():
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    )
    mcp.run()


if __name__ == "__main__":
    main()
