#!/usr/bin/env python3
"""
InsAIts MCP Server
==================
A real Model Context Protocol server that Claude Code connects to.
InsAIts runs as a separate process. Claude Code calls it as a tool.
Claude cannot bypass this -- it is a hard gate, not a suggestion.

HOW IT WORKS:
  1. You register this server in ~/.claude/claude_desktop_config.json
  2. Claude Code sees `insaits_check` as an available tool
  3. You instruct Claude Code (once, in settings) to call it after responses
  4. InsAIts inspects the response and returns structured feedback
  5. If not ok, Claude Code MUST address the feedback before finishing

SETUP (one time):
  pip install mcp insa-its

  Add to ~/.claude/claude_desktop_config.json:
  {
    "mcpServers": {
      "insaits": {
        "command": "python",
        "args": ["-m", "insa_its.mcp.server"],
        "env": {
          "INSAITS_PROJECT": "YUY AI",
          "INSAITS_MODEL":   "claude-opus"
        }
      }
    }
  }

  Then in Claude Code, add to your CLAUDE.md (just one line):
  "After every response, call the insaits_check tool."

  That one line is valid because it doesn't ask Claude to self-judge --
  it asks Claude to call an external tool that judges it objectively.

Author: Cristi Bogdan -- Steddy Nova SRL / YuyAI
"""

import os
import sys
import json
import re
import logging
import asyncio
import datetime
from pathlib import Path
from typing import Any, Dict, List

# -- MCP SDK --
try:
    from mcp.server import Server
    from mcp.server.models import InitializationOptions
    from mcp.server.stdio import stdio_server
    from mcp.types import (
        Tool, TextContent, CallToolResult,
        INVALID_PARAMS, INTERNAL_ERROR,
    )
    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False
    print("ERROR: mcp package not installed. Run: pip install mcp", file=sys.stderr)

# -- InsAIts detectors (adapted import path) --
try:
    from insa_its.ai_monitor import (
        AIMonitor, AnomalyType, AISeverity,
        BlankResponseDetector,
        WaitingForDescriptionDetector,
        IncompleteCodeDetector,
        TruncatedOutputDetector,
        RepetitionLoopDetector,
        ContextCollapseDetector,
    )
    MONITOR_AVAILABLE = True
except ImportError:
    MONITOR_AVAILABLE = False

logging.basicConfig(level=logging.INFO, stream=sys.stderr)
logger = logging.getLogger("insaits.mcp")

# -----------------------------------------------
# GLOBAL MONITOR INSTANCE
# -----------------------------------------------

PROJECT = os.environ.get("INSAITS_PROJECT", "unknown")
MODEL_ID = os.environ.get("INSAITS_MODEL", "claude-opus")
AUDIT_DIR = Path(os.environ.get("INSAITS_AUDIT_DIR", Path.home() / ".insaits" / "audit"))
AUDIT_DIR.mkdir(parents=True, exist_ok=True)

monitor = AIMonitor(model_id=MODEL_ID) if MONITOR_AVAILABLE else None
session_stats: Dict[str, Any] = {
    "calls": 0,
    "anomalies": 0,
    "interventions": 0,
    "started": datetime.datetime.utcnow().isoformat(),
}


def write_audit(event: Dict[str, Any]) -> None:
    """Append to tamper-evident audit log."""
    path = AUDIT_DIR / f"session_{datetime.date.today()}.jsonl"
    try:
        with open(path, "a") as f:
            f.write(json.dumps({
                "ts": datetime.datetime.utcnow().isoformat(),
                "timestamp_us": __import__("time").time_ns() // 1000,
                **event
            }) + "\n")
    except Exception as e:
        logger.warning("Audit write failed: %s", e)


# -----------------------------------------------
# TOOL DEFINITIONS
# -----------------------------------------------

TOOLS: List[Any] = []

if MCP_AVAILABLE:
    TOOLS = [
        Tool(
            name="insaits_check",
            description=(
                "InsAIts quality gate. Call this after EVERY response you generate. "
                "Pass your response text and the original task/prompt. "
                "InsAIts will inspect for: blank responses, stalling, incomplete code, "
                "truncation, repetition loops, and context drift. "
                "If issues are found, you MUST address the returned intervention "
                "before delivering your final response to the user."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "response": {
                        "type": "string",
                        "description": "Your response text to be inspected.",
                    },
                    "task": {
                        "type": "string",
                        "description": "The original task or prompt from the user.",
                    },
                    "agent_id": {
                        "type": "string",
                        "description": "Identifier for the agent/step.",
                        "default": MODEL_ID,
                    },
                },
                "required": ["response", "task"],
            },
        ),

        Tool(
            name="insaits_session_stats",
            description=(
                "Returns InsAIts monitoring statistics for the current session: "
                "total calls, anomaly count, anomaly rate, anomalies by type. "
                "Call this at the end of a session or when asked about monitoring status."
            ),
            inputSchema={
                "type": "object",
                "properties": {},
                "required": [],
            },
        ),

        Tool(
            name="insaits_preflight",
            description=(
                "Pre-flight task readiness check. Call this BEFORE starting a task "
                "to detect ambiguous instructions that will likely cause stalling. "
                "Returns: ready (bool), warnings (list), tip (str)."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "task": {
                        "type": "string",
                        "description": "The task description to check for ambiguity.",
                    },
                },
                "required": ["task"],
            },
        ),

        Tool(
            name="insaits_alert",
            description=(
                "Manually log an alert to the InsAIts audit trail. "
                "Use when you detect a problem yourself and want to surface it. "
                "This creates a human-readable entry in the session audit log."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "message": {"type": "string", "description": "Alert message."},
                    "severity": {
                        "type": "string",
                        "enum": ["INFO", "LOW", "MEDIUM", "HIGH", "CRITICAL"],
                        "default": "MEDIUM",
                    },
                },
                "required": ["message"],
            },
        ),
    ]


# -----------------------------------------------
# TOOL HANDLERS
# -----------------------------------------------

def handle_insaits_check(args: Dict[str, Any]) -> Dict[str, Any]:
    """Core quality gate -- called by Claude Code after each response."""
    response = args.get("response", "")
    task = args.get("task", "")
    agent_id = args.get("agent_id", MODEL_ID)

    session_stats["calls"] += 1

    if not MONITOR_AVAILABLE:
        return {
            "ok": True,
            "note": "InsAIts monitor not available -- insa_its.ai_monitor not found.",
        }

    result = monitor.inspect(response, task)

    write_audit({
        "event": "insaits_check",
        "agent_id": agent_id,
        "project": PROJECT,
        "ok": result.clean,
        "anomalies": [a.anomaly_type.value for a in result.anomalies],
        "task_preview": task[:80],
    })

    if result.clean:
        return {
            "ok": True,
            "message": "Response passed all InsAIts quality checks.",
            "stats": {
                "session_calls": session_stats["calls"],
                "session_anomalies": session_stats["anomalies"],
            },
        }

    # Build structured feedback for Claude Code
    session_stats["anomalies"] += len(result.anomalies)
    session_stats["interventions"] += 1

    issues = []
    for a in result.anomalies:
        issues.append({
            "type": a.anomaly_type.value,
            "severity": a.severity.value,
            "description": a.description,
            "intervention": a.intervention,
            "confidence": a.confidence,
        })

    write_audit({
        "event": "intervention",
        "agent_id": agent_id,
        "project": PROJECT,
        "issues": [i["type"] for i in issues],
        "severities": [i["severity"] for i in issues],
    })

    return {
        "ok": False,
        "issues": issues,
        "instruction": (
            "Your response has quality issues (listed above). "
            "Fix ALL listed issues and retry. "
            "Do not deliver the flawed response to the user. "
            "Address each intervention, then call insaits_check again to verify."
        ),
    }


def handle_insaits_preflight(args: Dict[str, Any]) -> Dict[str, Any]:
    """Pre-flight ambiguity check."""
    task = args.get("task", "")

    warnings: List[str] = []
    checks = [
        (r"\b(something|anything|whatever|however)\b",
         "Vague language detected -- Claude may stall asking for clarification."),
        (r"\b(maybe|perhaps|possibly|might)\b",
         "Uncertain language -- be more definitive about what you want."),
        (r"\bsome\b.{0,20}\bfiles?\b",
         "Unspecified files -- name the exact files to edit."),
        (r"\blike before\b|\bsame as\b|\bsimilar to\b",
         "Relies on prior context ('like before') that may not be in Claude's window."),
        (r"^.{0,40}$",
         "Task is very short -- Claude may need more context."),
        (r"^(?!.*(\.py|\.ts|\.js|\.md|class |def |function |implement|build|add|fix|refactor)).{0,200}$",
         "No specific file, function, or action mentioned -- add specifics."),
    ]

    for pattern, warning in checks:
        if re.search(pattern, task, re.IGNORECASE):
            warnings.append(warning)

    ready = len(warnings) == 0
    write_audit({
        "event": "preflight",
        "task": task[:80],
        "ready": ready,
        "warnings": warnings,
    })

    return {
        "ready": ready,
        "warnings": warnings,
        "tip": (
            "Task looks specific and clear -- proceed."
            if ready else
            "Add: exact file paths, function/class names, expected behaviour, "
            "and 'proceed with reasonable assumptions if unclear' to prevent stalling."
        ),
    }


def handle_session_stats(_args: Dict[str, Any]) -> Dict[str, Any]:
    """Return current session statistics."""
    base = session_stats.copy()
    if MONITOR_AVAILABLE:
        sdk_stats = monitor.get_session_stats()
        base["anomaly_rate"] = sdk_stats["anomaly_rate"]
        base["anomalies_by_type"] = sdk_stats["by_type"]
    base["audit_dir"] = str(AUDIT_DIR)
    return base


def handle_insaits_alert(args: Dict[str, Any]) -> Dict[str, Any]:
    """Log a manual alert to the audit trail."""
    message = args.get("message", "")
    severity = args.get("severity", "MEDIUM")
    write_audit({
        "event": "manual_alert",
        "severity": severity,
        "message": message,
    })
    return {"logged": True, "severity": severity, "message": message}


# -----------------------------------------------
# MCP SERVER
# -----------------------------------------------

async def run_server() -> None:
    """Start the MCP server (stdio transport)."""
    if not MCP_AVAILABLE:
        sys.exit(1)

    app = Server("insaits")

    @app.list_tools()
    async def list_tools() -> List[Any]:
        return TOOLS

    @app.call_tool()
    async def call_tool(name: str, arguments: Dict[str, Any]) -> List[Any]:
        try:
            if name == "insaits_check":
                result = handle_insaits_check(arguments)
            elif name == "insaits_preflight":
                result = handle_insaits_preflight(arguments)
            elif name == "insaits_session_stats":
                result = handle_session_stats(arguments)
            elif name == "insaits_alert":
                result = handle_insaits_alert(arguments)
            else:
                result = {"error": f"Unknown tool: {name}"}

            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        except Exception as e:
            logger.exception("Tool %s failed: %s", name, e)
            return [TextContent(type="text", text=json.dumps({"error": str(e)}))]

    logger.info("InsAIts MCP server starting -- project: %s, model: %s", PROJECT, MODEL_ID)
    logger.info("Audit log: %s", AUDIT_DIR)

    async with stdio_server() as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="insaits",
                server_version="3.0.1",
                capabilities=app.get_capabilities(
                    notification_options=None,
                    experimental_capabilities={},
                ),
            ),
        )


# -----------------------------------------------
# FALLBACK: run in test/demo mode without MCP
# -----------------------------------------------

def demo_mode() -> None:
    """
    Run without MCP to test the handlers directly.
    Useful during development before wiring into Claude Code.
    """
    logger.info("InsAIts MCP Server -- Demo Mode (no MCP connection)")
    logger.info("=" * 60)

    tests = [
        ("insaits_check", {
            "response": "",
            "task": "Implement the SkillRegistry class",
        }),
        ("insaits_check", {
            "response": "Could you provide more details about what you need?",
            "task": "Add Bayesian scoring to assess.py",
        }),
        ("insaits_check", {
            "response": "class SkillRegistry:\n    def lookup(self):\n        pass  # TODO",
            "task": "Implement SkillRegistry",
        }),
        ("insaits_preflight", {
            "task": "fix the thing from before",
        }),
        ("insaits_preflight", {
            "task": "Implement SkillRegistry.lookup() in yuyai/skills/registry.py using FAISS",
        }),
        ("insaits_session_stats", {}),
    ]

    handlers = {
        "insaits_check": handle_insaits_check,
        "insaits_preflight": handle_insaits_preflight,
        "insaits_session_stats": handle_session_stats,
    }

    for tool_name, args in tests:
        logger.info("\n%s", "-" * 60)
        logger.info("Tool: %s", tool_name)
        logger.info("Args: %s", json.dumps(args, indent=2))

        handler = handlers.get(tool_name)
        if handler:
            result = handler(args)
            is_ok = result.get("ok", result.get("ready", True))
            status = "OK" if is_ok else "ISSUE"
            logger.info("Result: %s", status)
            logger.info("%s", json.dumps(result, indent=2))


if __name__ == "__main__":
    if "--demo" in sys.argv:
        demo_mode()
    else:
        asyncio.run(run_server())
