"""DriftGuard MCP Server — expose drift detection tools via Model Context Protocol.

MCP Tools:
    drift_scan      — Scan for drift across configured cloud providers
    drift_history   — View drift detection history
    drift_classify  — Classify a detected drift
    drift_fix       — Generate IaC fix for a detected drift
    drift_policy    — Manage drift detection policies
    drift_baseline  — View/update baselines
    drift_report    — Generate drift posture report
"""

from __future__ import annotations

import json
import logging
import os
from pathlib import Path
from typing import Any

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import TextContent, Tool

from .core.baseline import BaselineManager
from .core.classifier import DriftClassifier
from .core.differ import DriftDiffer
from .core.remediation import RemediationEngine
from .iac.terraform import TerraformStateParser
from .iac.pulumi import PulumiStateParser
from .policies.security import SecurityPolicyEngine
from .policies.cost import CostPolicyEngine
from .policies.compliance import CompliancePolicyEngine
from .store import (
    DriftClassification,
    DriftSeverity,
    DriftStore,
    RemediationStatus,
)

logger = logging.getLogger(__name__)

# Initialize server
app = Server("driftguard")

# Global store instance
_store: DriftStore | None = None


def _get_store() -> DriftStore:
    global _store
    if _store is None:
        db_path = os.environ.get("DRIFTGUARD_DB", "driftguard.db")
        _store = DriftStore(db_path)
    return _store


@app.list_tools()
async def list_tools() -> list[Tool]:
    """List all available DriftGuard tools."""
    return [
        Tool(
            name="drift_scan",
            description=(
                "Scan for infrastructure drift by comparing IaC state (Terraform/Pulumi) "
                "against live cloud resources. Returns drift summary, affected resources, "
                "severity classification, and policy violations."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "state_path": {
                        "type": "string",
                        "description": "Path to the IaC state file (e.g., terraform.tfstate)",
                    },
                    "iac_format": {
                        "type": "string",
                        "enum": ["terraform", "pulumi"],
                        "default": "terraform",
                        "description": "IaC tool format",
                    },
                    "provider": {
                        "type": "string",
                        "enum": ["azure", "aws", "gcp", "auto"],
                        "default": "auto",
                        "description": "Cloud provider to scan",
                    },
                    "run_policies": {
                        "type": "boolean",
                        "default": True,
                        "description": "Run security/cost/compliance policies",
                    },
                },
                "required": ["state_path"],
            },
        ),
        Tool(
            name="drift_history",
            description=(
                "View drift detection history with optional filters by scan ID, "
                "severity, provider, or remediation status."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "scan_id": {"type": "string", "description": "Filter by scan ID"},
                    "severity": {
                        "type": "string",
                        "enum": ["critical", "high", "medium", "low", "info"],
                    },
                    "provider": {"type": "string"},
                    "status": {
                        "type": "string",
                        "enum": ["pending", "in_progress", "remediated", "accepted", "ignored"],
                    },
                    "limit": {"type": "integer", "default": 50},
                },
            },
        ),
        Tool(
            name="drift_classify",
            description=(
                "Classify a detected drift as unauthorized, intentional, or emergency. "
                "Uses heuristics based on attribute type, change patterns, and time."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "drift_id": {
                        "type": "string",
                        "description": "Drift record ID to classify",
                    },
                    "classification": {
                        "type": "string",
                        "enum": ["unauthorized", "intentional", "emergency"],
                        "description": "Manual classification override (optional)",
                    },
                },
                "required": ["drift_id"],
            },
        ),
        Tool(
            name="drift_fix",
            description=(
                "Generate IaC remediation code (Terraform HCL or Pulumi Python) "
                "to fix a detected drift. Optionally create a GitHub PR."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "scan_id": {"type": "string", "description": "Scan ID to fix drifts for"},
                    "iac_format": {
                        "type": "string",
                        "enum": ["terraform", "pulumi"],
                        "default": "terraform",
                    },
                    "create_pr": {
                        "type": "boolean",
                        "default": False,
                        "description": "Create a GitHub PR with the fix",
                    },
                },
                "required": ["scan_id"],
            },
        ),
        Tool(
            name="drift_policy",
            description=(
                "Manage drift detection policies. List, enable, or disable "
                "security, cost, and compliance rules."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["list", "enable", "disable"],
                        "default": "list",
                    },
                    "category": {
                        "type": "string",
                        "enum": ["security", "cost", "compliance", "all"],
                        "default": "all",
                    },
                    "rule_id": {
                        "type": "string",
                        "description": "Rule ID to enable/disable",
                    },
                },
            },
        ),
        Tool(
            name="drift_baseline",
            description=(
                "View or update drift baselines. Baselines mark known deviations "
                "as accepted so they don't trigger alerts."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["list", "add", "remove"],
                        "default": "list",
                    },
                    "resource_id": {"type": "string"},
                    "attribute": {"type": "string"},
                    "accepted_value": {"type": "string"},
                    "reason": {"type": "string"},
                    "ttl_days": {"type": "integer"},
                },
            },
        ),
        Tool(
            name="drift_report",
            description=(
                "Generate a drift posture report with trends, severity breakdown, "
                "resolution rates, and recent scan results."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "days": {
                        "type": "integer",
                        "default": 30,
                        "description": "Report period in days",
                    },
                },
            },
        ),
    ]


@app.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
    """Dispatch tool calls to handlers."""
    try:
        if name == "drift_scan":
            result = _handle_scan(arguments)
        elif name == "drift_history":
            result = _handle_history(arguments)
        elif name == "drift_classify":
            result = _handle_classify(arguments)
        elif name == "drift_fix":
            result = _handle_fix(arguments)
        elif name == "drift_policy":
            result = _handle_policy(arguments)
        elif name == "drift_baseline":
            result = _handle_baseline(arguments)
        elif name == "drift_report":
            result = _handle_report(arguments)
        else:
            result = {"error": f"Unknown tool: {name}"}

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


def _handle_scan(args: dict[str, Any]) -> dict[str, Any]:
    """Handle drift_scan tool call."""
    store = _get_store()
    state_path = args["state_path"]
    iac = args.get("iac_format", "terraform")
    provider = args.get("provider", "auto")
    run_policies = args.get("run_policies", True)

    # Parse state
    path = Path(state_path)
    if not path.exists():
        return {"error": f"State file not found: {state_path}"}

    if iac == "terraform":
        parser = TerraformStateParser()
        resources = parser.parse_file(path)
        desired = {addr: res.attributes for addr, res in resources.items()}
        metadata = {
            addr: {"type": res.type, "id": res.resource_id, "name": res.display_name, "provider": res.cloud_provider}
            for addr, res in resources.items()
        }
    elif iac == "pulumi":
        parser_p = PulumiStateParser()
        resources = parser_p.parse_file(path)
        desired = {urn: {**res.attributes, **res.outputs} for urn, res in resources.items()}
        metadata = {
            urn: {"type": res.terraform_type, "id": res.resource_id, "name": res.display_name, "provider": res.cloud_provider}
            for urn, res in resources.items()
        }
    else:
        return {"error": f"Unsupported IaC format: {iac}"}

    # Compare (actual = desired for demo; in production, fetch from cloud)
    actual = dict(desired)
    differ = DriftDiffer(store=store)
    result = differ.compare(desired, actual, metadata, provider=provider, iac_source=iac, state_path=state_path)

    output: dict[str, Any] = {
        "summary": result.to_summary(),
        "drifts": [
            {
                "resource": d.resource_address,
                "type": d.resource_type,
                "diff_type": d.diff_type,
                "severity": d.severity.value,
                "changes_count": len(d.attribute_diffs),
            }
            for d in result.resource_diffs
        ],
    }

    # Policy violations
    if run_policies:
        violations = []
        sec = SecurityPolicyEngine()
        cost = CostPolicyEngine()
        comp = CompliancePolicyEngine()
        for addr, attrs in desired.items():
            res_type = metadata.get(addr, {}).get("type", "")
            violations.extend(sec.evaluate(addr, res_type, attrs))
            violations.extend(cost.evaluate(addr, res_type, attrs))
            violations.extend(comp.evaluate(addr, res_type, attrs))
        output["policy_violations"] = [v.to_dict() for v in violations]
        output["policy_violation_count"] = len(violations)

    return output


def _handle_history(args: dict[str, Any]) -> dict[str, Any]:
    """Handle drift_history tool call."""
    store = _get_store()
    sev = DriftSeverity(args["severity"]) if args.get("severity") else None
    status = RemediationStatus(args["status"]) if args.get("status") else None
    drifts = store.get_drifts(
        scan_id=args.get("scan_id"),
        provider=args.get("provider"),
        severity=sev,
        status=status,
        limit=args.get("limit", 50),
    )
    return {"count": len(drifts), "drifts": drifts}


def _handle_classify(args: dict[str, Any]) -> dict[str, Any]:
    """Handle drift_classify tool call."""
    store = _get_store()
    drift_id = args["drift_id"]

    if "classification" in args:
        # Manual override
        classification = DriftClassification(args["classification"])
        store.update_classification(drift_id, classification)
        return {"drift_id": drift_id, "classification": classification.value, "source": "manual"}

    # Auto-classify would require the full ResourceDiff context
    return {"drift_id": drift_id, "message": "Auto-classification requires scan context. Use manual classification."}


def _handle_fix(args: dict[str, Any]) -> dict[str, Any]:
    """Handle drift_fix tool call."""
    store = _get_store()
    scan_id = args["scan_id"]
    iac_format = args.get("iac_format", "terraform")

    drifts = store.get_drifts(scan_id=scan_id)
    if not drifts:
        return {"error": f"No drifts found for scan {scan_id}"}

    return {
        "scan_id": scan_id,
        "drift_count": len(drifts),
        "iac_format": iac_format,
        "message": f"Found {len(drifts)} drift(s). Use the CLI for full remediation code generation.",
    }


def _handle_policy(args: dict[str, Any]) -> dict[str, Any]:
    """Handle drift_policy tool call."""
    action = args.get("action", "list")
    category = args.get("category", "all")

    engines: dict[str, Any] = {}
    if category in ("security", "all"):
        engines["security"] = SecurityPolicyEngine()
    if category in ("cost", "all"):
        engines["cost"] = CostPolicyEngine()
    if category in ("compliance", "all"):
        engines["compliance"] = CompliancePolicyEngine()

    if action == "list":
        result: dict[str, Any] = {}
        for cat, engine in engines.items():
            rules = engine.get_enabled_rules() if hasattr(engine, "get_enabled_rules") else []
            result[cat] = {"rule_count": len(rules), "rules": rules}
        return result

    rule_id = args.get("rule_id")
    if not rule_id:
        return {"error": "rule_id required for enable/disable"}

    for engine in engines.values():
        if action == "enable":
            if engine.enable_rule(rule_id):
                return {"rule_id": rule_id, "action": "enabled"}
        elif action == "disable":
            if engine.disable_rule(rule_id):
                return {"rule_id": rule_id, "action": "disabled"}

    return {"error": f"Rule {rule_id} not found"}


def _handle_baseline(args: dict[str, Any]) -> dict[str, Any]:
    """Handle drift_baseline tool call."""
    store = _get_store()
    mgr = BaselineManager(store)
    action = args.get("action", "list")

    if action == "list":
        baselines = mgr.get_all()
        return {"count": len(baselines), "baselines": baselines}

    if action == "add":
        if not all(k in args for k in ("resource_id", "attribute", "reason")):
            return {"error": "resource_id, attribute, and reason are required"}
        entry = mgr.accept_drift(
            resource_id=args["resource_id"],
            resource_type="manual",
            attribute=args["attribute"],
            accepted_value=args.get("accepted_value"),
            reason=args["reason"],
            ttl_days=args.get("ttl_days"),
        )
        return {"status": "created", "resource_id": entry.resource_id, "attribute": entry.attribute}

    if action == "remove":
        if not all(k in args for k in ("resource_id", "attribute")):
            return {"error": "resource_id and attribute are required"}
        removed = mgr.remove_baseline(args["resource_id"], args["attribute"])
        return {"status": "removed" if removed else "not_found"}

    return {"error": f"Unknown action: {action}"}


def _handle_report(args: dict[str, Any]) -> dict[str, Any]:
    """Handle drift_report tool call."""
    store = _get_store()
    days = args.get("days", 30)
    trends = store.get_drift_trends(days)
    scans = store.get_scans(limit=10)
    return {
        "trends": trends,
        "recent_scans": scans,
    }


async def run_server() -> None:
    """Run the MCP server."""
    async with stdio_server() as (read_stream, write_stream):
        await app.run(read_stream, write_stream, app.create_initialization_options())


def main() -> None:
    """Entry point for MCP server."""
    import asyncio
    asyncio.run(run_server())


if __name__ == "__main__":
    main()
