"""MCP tool wrapper for Codex GraphRAG."""

from __future__ import annotations

import json
import logging
import time
from pathlib import Path
from typing import Any, Dict, Optional

from src.utils.tool_wrapper import safe_tool

from .contracts import VALID_ACTIONS, error_payload, parse_event_json, validate_discrepancy_event
from .core import build_index, search_index, recommend_workflow, review_workflow, ingest_discrepancy_event
from .closure_engine import run_closure_loop, get_closure_metrics, trace_discrepancy, propose_fix
from .trust_policy import explain_trust_event, get_retention_status, get_trust_metrics

logger = logging.getLogger(__name__)

_DEFAULT_INDEX = "data/codex_graphrag/index.json"
_DEFAULT_DOCS = "docs/Codex"
_STALE_CHECK_TTL_SECONDS = 60
_stale_check_cache: Dict[str, tuple[float, bool]] = {}


def _index_is_stale(index_path: str, docs_root: str) -> bool:
    """Return True if docs changed after index was last written."""
    index = Path(index_path)
    docs = Path(docs_root)
    if not index.exists():
        return True
    if not docs.exists():
        return False

    cache_key = f"{index.resolve()}::{docs.resolve()}"
    now = time.monotonic()
    cached = _stale_check_cache.get(cache_key)
    if cached and (now - cached[0]) < _STALE_CHECK_TTL_SECONDS:
        return cached[1]

    idx_mtime = index.stat().st_mtime
    latest_doc_mtime = idx_mtime
    try:
        for p in docs.rglob("*.json"):
            try:
                latest_doc_mtime = max(latest_doc_mtime, p.stat().st_mtime)
            except OSError:
                continue
    except OSError:
        return False

    stale = latest_doc_mtime > idx_mtime
    _stale_check_cache[cache_key] = (now, stale)
    return stale


def _ensure_index(
    index_path: str,
    docs_root: str,
    auto_build: bool = True,
    refresh_if_stale: bool = True,
) -> None:
    index_exists = Path(index_path).exists()
    if index_exists and not refresh_if_stale:
        return
    if index_exists and refresh_if_stale and not _index_is_stale(index_path=index_path, docs_root=docs_root):
        return
    if auto_build:
        if index_exists:
            logger.info("Codex GraphRAG index stale; rebuilding automatically")
        build_index(docs_root=docs_root, index_path=index_path)


def register_codex_graphrag_tools(mcp, settings):
    """Register unified Codex GraphRAG MCP tool."""

    @mcp.tool()
    @safe_tool("codex_graphrag")
    def codex_graphrag(
        action: str,
        query: Optional[str] = None,
        goal: Optional[str] = None,
        docs_root: str = _DEFAULT_DOCS,
        index_path: str = _DEFAULT_INDEX,
        k: int = 10,
        expand_depth: int = 1,
        category: str = "",
        domain: str = "",
        max_steps: int = 10,
        review_mode: str = "balanced",
        persona: str = "balanced",
        response_mode: str = "plain",
        include_evidence: bool = True,
        include_guardrails: bool = False,
        event_json: str = "",
        runtime_events_path: str = "data/codex_graphrag/runtime_discrepancy_events.json",
        out_path: str = "",
        auto_build: bool = True,
        allow_cortex: bool = False,
        tenant_id: str = "default",
        ingest_mode: str = "dual",
        dedupe_window_seconds: int = 900,
        strict_contract: bool = True,
        policy_mode: str = "shadow",
        event_key: str = "",
        metrics_path: str = "data/codex_graphrag/closure_metrics.json",
    ) -> Dict[str, Any]:
        """Unified GraphRAG action-dispatch tool.

        Actions:
        - build: Build index from docs/Codex JSON sidecars
        - search: Retrieve relevant nodes for a query
        - workflow: Recommend ordered workflow steps for a goal
        - review: Generate structured review report + AI review prompt for a goal
        - ingest_discrepancy: Ingest a production discrepancy event into runtime trust memory
        - closure_loop: Run full Sprinkler closure cycle (trace -> propose -> verify)
        - closure_metrics: Get aggregated closure loop metrics
        - propose_fix: Three-tier fix generation (trust replay -> relationship RAG -> Cortex AI)
        - trust_metrics: Aggregate trust/policy metrics over runtime events
        - trust_explain: Explain trust scoring/policy decisions for an event_key
        - retention_status: Summarize hot/warm/archive retention tiers

        Args:
            action: One of build|search|workflow|review|ingest_discrepancy|closure_loop|closure_metrics|propose_fix|trust_metrics|trust_explain|retention_status
            query: Search query (required for action=search)
            goal: Goal statement (required for action=workflow)
            docs_root: Path to docs root containing JSON sidecars
            index_path: Path to persisted index JSON
            k: Top-k results (search)
            expand_depth: Graph expansion depth (search)
            category: Optional category filter (search)
            domain: Optional domain filter (search/workflow)
            max_steps: Maximum steps in recommendation (workflow)
            review_mode: Review rigor for action=review: strict|balanced|fast
            persona: Answer framing: executive|technical|balanced
            response_mode: plain|cfo_strict
            include_evidence: Include citations/evidence fields
            include_guardrails: Include anti-pattern guardrails in output
            event_json: JSON object string for action=ingest_discrepancy
            runtime_events_path: Path to runtime discrepancy events JSON file
            out_path: Optional output path to persist search/workflow result
            auto_build: Auto-build index if missing for search/workflow
            allow_cortex: Explicitly allow Tier-3 Cortex fix fallback for propose/closure
            tenant_id: Tenant scope for discrepancy events
            ingest_mode: dual|snowflake_only|file_only for discrepancy ingestion
            dedupe_window_seconds: Event key dedupe time window
            strict_contract: Require source_system/source in event payload
            policy_mode: shadow|enforce policy execution mode
            event_key: Required for action=trust_explain
            metrics_path: Optional closure metrics file path
        """
        action = (action or "").strip().lower()
        valid_actions = set(VALID_ACTIONS)
        if action not in valid_actions:
            return error_payload(
                code="invalid_action",
                message=f"Invalid action '{action}'",
                action=action,
                valid_actions=sorted(valid_actions),
            )

        try:
            if action == "build":
                result = build_index(docs_root=docs_root, index_path=index_path)
            elif action == "search":
                if not query:
                    return error_payload(
                        code="missing_query",
                        message="query is required for action='search'",
                        action=action,
                        details={"required_field": "query"},
                    )
                _ensure_index(index_path=index_path, docs_root=docs_root, auto_build=auto_build)
                result = search_index(
                    index_path=index_path,
                    query=query,
                    k=k,
                    expand_depth=expand_depth,
                    category=category,
                    domain=domain,
                    persona=persona,
                    response_mode=response_mode,
                    include_evidence=include_evidence,
                    include_guardrails=include_guardrails,
                )
            elif action == "ingest_discrepancy":
                event, parse_error = parse_event_json(event_json, action=action)
                if parse_error:
                    return parse_error
                validation_error = validate_discrepancy_event(event, action=action, strict_contract=strict_contract)
                if validation_error:
                    return validation_error
                result = ingest_discrepancy_event(
                    event=event,
                    index_path=index_path,
                    docs_root=docs_root,
                    runtime_events_path=runtime_events_path,
                    auto_build=auto_build,
                    tenant_id=tenant_id,
                    ingest_mode=ingest_mode,
                    dedupe_window_seconds=dedupe_window_seconds,
                    policy_mode=policy_mode,
                )
            elif action == "closure_loop":
                event, parse_error = parse_event_json(event_json, action=action)
                if parse_error:
                    return parse_error
                validation_error = validate_discrepancy_event(event, action=action, strict_contract=strict_contract)
                if validation_error:
                    return validation_error
                result = run_closure_loop(
                    event=event,
                    index_path=index_path,
                    docs_root=docs_root,
                    runtime_events_path=runtime_events_path,
                    auto_ingest=auto_build,
                    allow_cortex=allow_cortex,
                    tenant_id=tenant_id,
                    ingest_mode=ingest_mode,
                    dedupe_window_seconds=dedupe_window_seconds,
                    policy_mode=policy_mode,
                )
            elif action == "closure_metrics":
                result = get_closure_metrics(metrics_path=metrics_path)
            elif action == "trust_metrics":
                result = get_trust_metrics(
                    runtime_events_path=runtime_events_path,
                    closure_metrics_path=metrics_path,
                )
            elif action == "retention_status":
                result = get_retention_status(runtime_events_path=runtime_events_path)
            elif action == "trust_explain":
                if not event_key:
                    return error_payload(
                        code="missing_event_key",
                        message="event_key is required for action='trust_explain'",
                        action=action,
                        details={"required_field": "event_key"},
                    )
                result = explain_trust_event(
                    event_key=event_key,
                    runtime_events_path=runtime_events_path,
                )
            elif action == "propose_fix":
                event, parse_error = parse_event_json(event_json, action=action)
                if parse_error:
                    return parse_error
                validation_error = validate_discrepancy_event(event, action=action, strict_contract=strict_contract)
                if validation_error:
                    return validation_error
                _ensure_index(index_path=index_path, docs_root=docs_root, auto_build=auto_build)
                trace = trace_discrepancy(event, index_path=index_path)
                result = propose_fix(trace, index_path=index_path, event=event, allow_cortex=allow_cortex)
            else:
                if not goal:
                    return error_payload(
                        code="missing_goal",
                        message=f"goal is required for action='{action}'",
                        action=action,
                        details={"required_field": "goal"},
                    )
                _ensure_index(index_path=index_path, docs_root=docs_root, auto_build=auto_build)
                if action == "workflow":
                    result = recommend_workflow(
                        index_path=index_path,
                        goal=goal,
                        max_steps=max_steps,
                        domain=domain,
                        persona=persona,
                        response_mode=response_mode,
                        include_evidence=include_evidence,
                        include_guardrails=include_guardrails,
                    )
                else:
                    result = review_workflow(
                        index_path=index_path,
                        goal=goal,
                        max_steps=max_steps,
                        domain=domain,
                        review_mode=review_mode,
                        persona=persona,
                        response_mode=response_mode,
                        include_evidence=include_evidence,
                        include_guardrails=include_guardrails,
                    )

            if out_path:
                out = Path(out_path)
                out.parent.mkdir(parents=True, exist_ok=True)
                out.write_text(json.dumps(result, indent=2), encoding="utf-8")
                result["saved_to"] = str(out)

            return result
        except Exception as e:
            logger.exception("codex_graphrag failed")
            return error_payload(
                code="internal_error",
                message=str(e),
                action=action,
                retriable=True,
            )

    logger.info("Registered Codex GraphRAG tool: codex_graphrag")
    return ["codex_graphrag"]
