"""
Cerebion Rivet MCP Server

Exposes Rivet's local scan history and quantum risk data as MCP tools so an
MCP-compatible AI client (Claude Desktop, Claude Code, etc.) running on the
same machine can query monitoring data conversationally, and can trigger new
code scans.

Runs entirely over stdio — no network listener, no new port, no change to
Rivet's air-gap-friendly architecture. Reads the same local SQLite scan
database the desktop app already owns, via the same standalone-desktop
local user Rivet's FastAPI routes already use (see app/auth.py).

Requires an Enterprise license, same tier boundary as JSON/CSV/SARIF export.

trigger_scan runs in the background and returns immediately with a scan_id —
it does not block the calling turn. Confirmed by testing: a first-time scan
of a modest multi-language directory (~90 files) took well over a minute,
which is exactly why this can't be a synchronous tool call. Rescans of
already-scanned code hit the incremental file cache in scan_cache.py and
complete much faster. If the MCP client process is killed (e.g. the user
closes the chat session) before a triggered scan finishes, that scan is
orphaned and its partial work is lost — a known v1 limitation.
"""
import asyncio
import json
import logging
from typing import Optional

from mcp.server.fastmcp import FastMCP
from sqlmodel import Session, select

from app.api.reports import _extract_findings_count, _extract_risk_score
from app.database import engine, get_session_context
from app.models import ScanRecord
from app.services.dashboard_metrics import get_metrics_provider

_LOCAL_USER_ID = "local"

logger = logging.getLogger(__name__)

mcp = FastMCP("cerebion-rivet")

# asyncio.create_task() only holds a WEAK reference to the task — without
# keeping a strong reference somewhere, the task can be garbage-collected
# mid-execution. Keep in-flight background scans alive here until they finish.
_background_tasks: set = set()


@mcp.tool()
async def get_risk_summary(time_range: Optional[str] = None) -> str:
    """Get the aggregated quantum risk summary: overall metrics, findings by
    severity, and per-analyzer breakdown. Same data as the Dashboard page.

    Args:
        time_range: Optional filter, e.g. "7d", "30d", "90d". Omit for all-time.
    """
    with Session(engine) as session:
        provider = get_metrics_provider(session)
        metrics = await provider.get_metrics(_LOCAL_USER_ID, time_range)
    return json.dumps(metrics, indent=2, default=str)


# The tool's public vocabulary ("code") is friendlier than the raw ScanType
# enum value actually stored in the database ("code_analysis") -- without
# this alias, filtering by "code" (exactly as this tool's own docstring
# instructs) silently returned zero rows even when code scans existed.
_SCAN_TYPE_ALIASES = {"code": "code_analysis"}


@mcp.tool()
def list_recent_scans(scan_type: Optional[str] = None, limit: int = 50) -> str:
    """List recent scans with target, type, timestamp, finding count, and
    risk score. Use this to find which scan_id to inspect with get_scan_detail.

    Args:
        scan_type: Optional filter: "code", "quantum", "port", or "binary".
        limit: Max results to return (default 50, max 500).
    """
    limit = max(1, min(limit, 500))
    with Session(engine) as session:
        stmt = (
            select(ScanRecord)
            .where(ScanRecord.user_id == _LOCAL_USER_ID)
            .order_by(ScanRecord.timestamp.desc())
            .limit(limit)
        )
        if scan_type:
            stmt = stmt.where(ScanRecord.scan_type == _SCAN_TYPE_ALIASES.get(scan_type, scan_type))
        scans = session.exec(stmt).all()

        rows = []
        for s in scans:
            sd = s.scan_data or {}
            rows.append({
                "id": s.id,
                "target": s.target,
                "scan_type": s.scan_type.value if hasattr(s.scan_type, "value") else str(s.scan_type),
                "timestamp": s.timestamp.isoformat(),
                "findings_count": _extract_findings_count(sd, s.scan_type),
                "risk_score": _extract_risk_score(sd),
            })
    return json.dumps(rows, indent=2)


@mcp.tool()
def get_scan_detail(scan_id: str) -> str:
    """Get the full findings and analysis data for a single scan by ID.
    Use list_recent_scans first to find the scan_id you want.

    Args:
        scan_id: The scan record ID from list_recent_scans.
    """
    with Session(engine) as session:
        stmt = select(ScanRecord).where(
            ScanRecord.id == scan_id,
            ScanRecord.user_id == _LOCAL_USER_ID,
        )
        scan = session.exec(stmt).first()
        if scan is None:
            return json.dumps({"error": f"Scan {scan_id} not found."})
        return json.dumps(
            {
                "id": scan.id,
                "target": scan.target,
                "scan_type": scan.scan_type.value if hasattr(scan.scan_type, "value") else str(scan.scan_type),
                "timestamp": scan.timestamp.isoformat(),
                "scan_data": scan.scan_data,
            },
            indent=2,
            default=str,
        )


@mcp.tool()
async def trigger_scan(
    target_path: str,
    extensions: Optional[str] = None,
    excluded_dirs: Optional[str] = None,
) -> str:
    """Start a quantum-vulnerable code scan on target_path in the background
    and return immediately — this does not wait for the scan to finish.

    Rescans of already-scanned code usually complete in well under a second
    thanks to Rivet's incremental file cache (unchanged files are served from
    cache; only changed files are reanalyzed). A first-time scan of a large
    codebase can take longer. Do not call this tool repeatedly to check
    status — after calling it, use list_recent_scans or get_scan_detail with
    the returned scan_id to check whether the scan has completed and see
    results.

    Args:
        target_path: Absolute path to the file or directory to scan.
        extensions: Optional comma-separated file extensions to include, e.g. "py,ts,js".
        excluded_dirs: Optional comma-separated directories to exclude, e.g. "node_modules,tests".
    """
    from app.api.code_analysis import load_scanner_config, run_scan_sync
    from app.utils.input_validation import sanitize_excluded_dirs, sanitize_extensions

    try:
        scanner_config = load_scanner_config()

        selected_extensions = None
        if extensions:
            selected_extensions = sanitize_extensions(
                [e.strip() for e in extensions.split(",") if e.strip()]
            )

        excluded_dir_list = None
        if excluded_dirs:
            excluded_dir_list = sanitize_excluded_dirs(
                [d.strip() for d in excluded_dirs.split(",") if d.strip()]
            )
    except ValueError as exc:
        return json.dumps({"status": "error", "message": str(exc)})

    async def _run_in_background():
        try:
            async with get_session_context() as session:
                await asyncio.to_thread(
                    run_scan_sync,
                    target_path,
                    None,  # config: auto-detect, same as `rivet scan` default
                    scanner_config,
                    selected_extensions,
                    excluded_dir_list,
                    session,
                    _LOCAL_USER_ID,
                )
        except Exception as exc:
            logger.error(f"Background MCP scan failed for {target_path}: {exc}")

    task = asyncio.create_task(_run_in_background())
    _background_tasks.add(task)
    task.add_done_callback(_background_tasks.discard)

    return json.dumps({
        "status": "started",
        "target": target_path,
        "message": (
            "Scan running in the background. Poll list_recent_scans "
            "(filtered by target if helpful) to check for completion."
        ),
    })


def main():
    """Entry point for `rivet mcp-server`. Gates on Enterprise license, then
    serves over stdio — never writes anything but JSON-RPC to stdout."""
    from app.license_guard import check_license, require_feature

    license_result = check_license(skip=False)
    require_feature(license_result, "json_export", "rivet mcp-server")

    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
