"""DentalComplianceOS MCP Server — exposes dental compliance tools via Model Context Protocol.

Tools:
- dental_compliance_scan: Scan practice against a dental compliance framework
- dental_compliance_status: Current compliance posture across all frameworks
- dental_sterilization_log: Log a sterilization cycle
- dental_inspection_prep: Generate inspection preparation checklist
- dental_compliance_fix: Auto-generate corrective action plans for violations
- dental_compliance_report: Generate audit-ready compliance report
- dental_radiation_log: Log radiation exposure, equipment QA, lead apron checks
"""

from __future__ import annotations

import datetime
import json
import uuid
from pathlib import Path
from typing import Any

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

from .frameworks import (
    FRAMEWORKS,
    FRAMEWORK_NAMES,
    get_all_checks,
    get_checks_for_practice_type,
    get_checks_for_state,
)
from .scanners import SCANNERS
from .store import DentalComplianceStore
from .types import (
    ComplianceFinding,
    ComplianceStatus,
    Severity,
    ScanResult,
    SterilizationRecord,
    RadiationRecord,
    PracticeType,
)


_PRACTICE_TYPES = [e.value for e in PracticeType]
_FRAMEWORK_KEYS = list(FRAMEWORKS.keys())
_STATES = ["CA", "TX", "FL", "NY", "PA", "IL", "OH", "GA", "NC", "MI"]


def _run_scan(
    target: str,
    framework: str,
    store: DentalComplianceStore,
    practice_type: str = "solo_gp",
    state: str = "",
) -> dict[str, Any]:
    """Core scan logic shared by MCP tool and CLI."""
    target_path = Path(target).resolve()
    if not target_path.exists():
        return {"error": f"Target path does not exist: {target}"}

    if framework not in FRAMEWORKS:
        return {
            "error": f"Unknown framework: {framework}. Available: {', '.join(FRAMEWORKS.keys())}"
        }

    # Get checks filtered by practice type and state
    checks = get_all_checks(framework)
    if practice_type and practice_type != "all":
        checks = [
            c for c in checks
            if "all" in c.get("practice_types", ["all"])
            or practice_type in c.get("practice_types", ["all"])
        ]
    if state and framework == "state_boards":
        checks = [
            c for c in checks
            if not c.get("states") or state.upper() in c.get("states", [])
        ]

    scan_id = f"dental-scan-{uuid.uuid4().hex[:12]}"
    timestamp = datetime.datetime.utcnow().isoformat()

    all_findings: list[ComplianceFinding] = []

    # Group checks by scanner type
    checks_by_scanner: dict[str, list[dict[str, Any]]] = {}
    for check in checks:
        scanner_type = check.get("scanner", "practice")
        checks_by_scanner.setdefault(scanner_type, []).append(check)

    for scanner_type, scanner_checks in checks_by_scanner.items():
        scanner_cls = SCANNERS.get(scanner_type)
        if not scanner_cls:
            # Default to practice scanner
            scanner_cls = SCANNERS.get("practice")
        if not scanner_cls:
            continue

        scanner = scanner_cls(target_path)
        raw_findings = scanner.scan(scanner_checks)
        all_findings.extend(raw_findings)

    # Calculate score
    total = len(checks)
    non_compliant_ids = {
        f.check_id for f in all_findings
        if f.status in (ComplianceStatus.NON_COMPLIANT, ComplianceStatus.UNKNOWN)
    }
    failed = len(non_compliant_ids)
    passed = total - failed
    score = (passed / total * 100) if total > 0 else 0.0

    result = ScanResult(
        scan_id=scan_id,
        framework=framework,
        practice_type=practice_type,
        state=state,
        timestamp=timestamp,
        total_checks=total,
        compliant=passed,
        non_compliant=failed,
        partial=len([f for f in all_findings if f.status == ComplianceStatus.PARTIAL]),
        unknown=len([f for f in all_findings if f.status == ComplianceStatus.UNKNOWN]),
        score=score,
        findings=all_findings,
    )

    store.save_scan(result)
    return result.to_dict()


def _generate_corrective_actions(
    findings: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """Generate corrective action plans for non-compliant findings."""
    actions: list[dict[str, Any]] = []
    for f in findings:
        if f.get("status") in ("non_compliant", "unknown"):
            severity = f.get("severity", "medium")
            # Set timeline based on severity
            if severity == "critical":
                timeline = "Immediate — within 24-48 hours"
                priority = 1
            elif severity == "high":
                timeline = "Urgent — within 1-2 weeks"
                priority = 2
            elif severity == "medium":
                timeline = "Standard — within 30 days"
                priority = 3
            else:
                timeline = "Low priority — within 90 days"
                priority = 4

            actions.append({
                "check_id": f.get("check_id", ""),
                "title": f.get("title", ""),
                "severity": severity,
                "priority": priority,
                "corrective_action": f.get("corrective_action", f.get("detail", "")),
                "timeline": timeline,
                "status": "pending",
            })

    actions.sort(key=lambda a: a["priority"])
    return actions


def _generate_inspection_checklist(
    state: str, practice_type: str
) -> dict[str, Any]:
    """Generate a comprehensive inspection preparation checklist."""
    checklist: dict[str, list[dict[str, str]]] = {}

    for fw_key, fw_name in FRAMEWORK_NAMES.items():
        checks = get_all_checks(fw_key)
        # Filter by practice type
        if practice_type and practice_type != "all":
            checks = [
                c for c in checks
                if "all" in c.get("practice_types", ["all"])
                or practice_type in c.get("practice_types", ["all"])
            ]
        # Filter by state for state_boards
        if fw_key == "state_boards" and state:
            checks = [
                c for c in checks
                if not c.get("states") or state.upper() in c.get("states", [])
            ]

        items: list[dict[str, str]] = []
        for c in checks:
            items.append({
                "check_id": c["id"],
                "item": c.get("check", ""),
                "regulation": c.get("regulation_reference", ""),
                "severity": c.get("severity", "medium"),
                "evidence_needed": c.get("remediation", ""),
            })
        if items:
            checklist[fw_name] = items

    return {
        "state": state,
        "practice_type": practice_type,
        "generated_at": datetime.datetime.utcnow().isoformat(),
        "total_items": sum(len(v) for v in checklist.values()),
        "frameworks": checklist,
    }


def _generate_report(
    store: DentalComplianceStore,
    framework: str | None = None,
    fmt: str = "markdown",
) -> str:
    """Generate an audit-ready compliance report."""
    now = datetime.datetime.utcnow().isoformat()[:19]

    if fmt == "json":
        report_data: dict[str, Any] = {
            "title": "DentalComplianceOS Audit Report",
            "generated_at": now,
            "frameworks": {},
        }
        frameworks_to_report = [framework] if framework else list(FRAMEWORKS.keys())
        for fw in frameworks_to_report:
            scan = store.get_latest_scan(fw)
            findings = store.get_non_compliant_findings(fw)
            report_data["frameworks"][fw] = {
                "name": FRAMEWORK_NAMES.get(fw, fw),
                "latest_scan": scan,
                "non_compliant_findings": findings,
            }
        # Add sterilization and radiation logs
        report_data["sterilization_log"] = store.get_sterilization_log(days=90)
        report_data["radiation_log"] = store.get_radiation_log(days=90)
        report_data["expiring_training"] = store.get_expiring_training(days=60)
        return json.dumps(report_data, indent=2)

    # Markdown format
    lines: list[str] = []
    lines.append("# DentalComplianceOS Audit Report")
    lines.append(f"\nGenerated: {now}")
    lines.append("")

    frameworks_to_report = [framework] if framework else list(FRAMEWORKS.keys())

    # Executive Summary
    lines.append("## Executive Summary\n")
    total_checks = 0
    total_compliant = 0
    for fw in frameworks_to_report:
        scan = store.get_latest_scan(fw)
        if scan:
            total_checks += scan.get("total_checks", 0)
            total_compliant += scan.get("compliant", 0)
            score = scan.get("score", 0)
            lines.append(
                f"- **{FRAMEWORK_NAMES.get(fw, fw)}**: {score:.1f}% compliant "
                f"({scan.get('compliant', 0)}/{scan.get('total_checks', 0)} checks passed)"
            )
        else:
            lines.append(f"- **{FRAMEWORK_NAMES.get(fw, fw)}**: No scan data")

    if total_checks > 0:
        overall = total_compliant / total_checks * 100
        lines.append(f"\n**Overall Compliance Score: {overall:.1f}%**")
    lines.append("")

    # Detailed Findings
    lines.append("## Findings by Framework\n")
    for fw in frameworks_to_report:
        fw_name = FRAMEWORK_NAMES.get(fw, fw)
        lines.append(f"### {fw_name}\n")
        findings = store.get_non_compliant_findings(fw)
        if not findings:
            lines.append("No non-compliant findings.\n")
            continue

        # Group by severity
        by_sev: dict[str, list[dict]] = {}
        for f in findings:
            sev = f.get("severity", "medium")
            by_sev.setdefault(sev, []).append(f)

        for sev in ["critical", "high", "medium", "low", "info"]:
            items = by_sev.get(sev, [])
            if not items:
                continue
            lines.append(f"#### {sev.upper()} ({len(items)})\n")
            for item in items:
                lines.append(f"- **[{item.get('check_id', '')}]** {item.get('title', '')}")
                if item.get("corrective_action"):
                    lines.append(f"  - Corrective Action: {item['corrective_action']}")
            lines.append("")

    # Sterilization Summary
    lines.append("## Sterilization Log (Last 90 Days)\n")
    sterilization = store.get_sterilization_log(days=90)
    if sterilization:
        lines.append(f"Total cycles logged: {len(sterilization)}")
        failures = [s for s in sterilization if s.get("biological_indicator") == "fail"]
        if failures:
            lines.append(f"\n**WARNING: {len(failures)} failed biological indicator(s)**")
        lines.append("")
    else:
        lines.append("No sterilization logs recorded.\n")

    # Radiation Summary
    lines.append("## Radiation Log (Last 90 Days)\n")
    radiation = store.get_radiation_log(days=90)
    if radiation:
        lines.append(f"Total entries: {len(radiation)}")
        by_type: dict[str, int] = {}
        for r in radiation:
            t = r.get("log_type", "unknown")
            by_type[t] = by_type.get(t, 0) + 1
        for t, count in by_type.items():
            lines.append(f"- {t}: {count}")
        lines.append("")
    else:
        lines.append("No radiation logs recorded.\n")

    # Training
    lines.append("## Expiring Training (Next 60 Days)\n")
    expiring = store.get_expiring_training(days=60)
    if expiring:
        for t in expiring:
            lines.append(
                f"- **{t.get('employee', 'Unknown')}**: {t.get('training_type', '')} "
                f"expires {t.get('expiration_date', '')[:10]}"
            )
        lines.append("")
    else:
        lines.append("No training expiring in the next 60 days.\n")

    return "\n".join(lines)


def create_server(db_path: str | None = None) -> Server:
    """Create and configure the MCP server."""
    server = Server("dentalcomplianceos")
    store = DentalComplianceStore(db_path)

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        return [
            Tool(
                name="dental_compliance_scan",
                description=(
                    "Scan a dental practice's documentation against a compliance framework. "
                    "Evaluates OSHA dental standards, CDC infection control, HIPAA for dental, "
                    "radiation safety, state dental board regulations, and sterilization protocols. "
                    f"Frameworks: {', '.join(_FRAMEWORK_KEYS)}"
                ),
                inputSchema={
                    "type": "object",
                    "properties": {
                        "target": {
                            "type": "string",
                            "description": "Path to the practice documentation directory",
                        },
                        "framework": {
                            "type": "string",
                            "enum": _FRAMEWORK_KEYS,
                            "description": "Dental compliance framework to scan against",
                        },
                        "practice_type": {
                            "type": "string",
                            "enum": _PRACTICE_TYPES,
                            "description": "Type of dental practice (default: solo_gp)",
                            "default": "solo_gp",
                        },
                        "state": {
                            "type": "string",
                            "enum": _STATES,
                            "description": "State for state-specific regulations (2-letter code)",
                        },
                    },
                    "required": ["target", "framework"],
                },
            ),
            Tool(
                name="dental_compliance_status",
                description=(
                    "Get current compliance posture across all dental compliance frameworks. "
                    "Shows latest scan results, scores, and number of non-compliant findings."
                ),
                inputSchema={
                    "type": "object",
                    "properties": {
                        "framework": {
                            "type": "string",
                            "enum": _FRAMEWORK_KEYS,
                            "description": "Specific framework to check (optional, shows all if omitted)",
                        },
                    },
                },
            ),
            Tool(
                name="dental_sterilization_log",
                description=(
                    "Log a sterilization cycle (autoclave run). Records cycle parameters, "
                    "biological indicator (spore test) results, and chemical indicator results."
                ),
                inputSchema={
                    "type": "object",
                    "properties": {
                        "autoclave_id": {
                            "type": "string",
                            "description": "Autoclave identifier (e.g. 'Statim-1', 'Midmark-M11')",
                        },
                        "cycle_type": {
                            "type": "string",
                            "enum": ["gravity", "pre_vacuum", "flash"],
                            "description": "Type of sterilization cycle",
                        },
                        "temperature": {
                            "type": "number",
                            "description": "Cycle temperature in Fahrenheit (e.g. 250, 270)",
                        },
                        "pressure": {
                            "type": "number",
                            "description": "Cycle pressure in PSI (e.g. 15, 30)",
                        },
                        "duration_minutes": {
                            "type": "number",
                            "description": "Cycle duration in minutes",
                        },
                        "biological_indicator": {
                            "type": "string",
                            "enum": ["pass", "fail", "pending", "not_tested"],
                            "description": "Biological indicator (spore test) result",
                        },
                        "chemical_indicator": {
                            "type": "string",
                            "enum": ["pass", "fail"],
                            "description": "Chemical indicator result",
                        },
                        "load_contents": {
                            "type": "string",
                            "description": "Description of instruments in the load",
                        },
                        "operator": {
                            "type": "string",
                            "description": "Name of person running the cycle",
                        },
                        "notes": {
                            "type": "string",
                            "description": "Additional notes",
                        },
                    },
                    "required": [
                        "autoclave_id",
                        "cycle_type",
                        "temperature",
                        "pressure",
                        "duration_minutes",
                        "biological_indicator",
                        "chemical_indicator",
                    ],
                },
            ),
            Tool(
                name="dental_inspection_prep",
                description=(
                    "Generate a comprehensive inspection preparation checklist based on "
                    "state and practice type. Covers all frameworks with specific evidence needed."
                ),
                inputSchema={
                    "type": "object",
                    "properties": {
                        "state": {
                            "type": "string",
                            "enum": _STATES,
                            "description": "State (2-letter code) for state-specific requirements",
                        },
                        "practice_type": {
                            "type": "string",
                            "enum": _PRACTICE_TYPES,
                            "description": "Type of dental practice",
                            "default": "solo_gp",
                        },
                    },
                    "required": ["state"],
                },
            ),
            Tool(
                name="dental_compliance_fix",
                description=(
                    "Auto-generate corrective action plans for non-compliant findings. "
                    "Prioritizes by severity with specific timelines and remediation steps."
                ),
                inputSchema={
                    "type": "object",
                    "properties": {
                        "framework": {
                            "type": "string",
                            "enum": _FRAMEWORK_KEYS,
                            "description": "Framework to generate fixes for (optional, all if omitted)",
                        },
                    },
                },
            ),
            Tool(
                name="dental_compliance_report",
                description=(
                    "Generate an audit-ready compliance report. Includes executive summary, "
                    "findings by framework, sterilization logs, radiation logs, and training status."
                ),
                inputSchema={
                    "type": "object",
                    "properties": {
                        "framework": {
                            "type": "string",
                            "enum": _FRAMEWORK_KEYS,
                            "description": "Specific framework to report on (optional, all if omitted)",
                        },
                        "format": {
                            "type": "string",
                            "enum": ["markdown", "json"],
                            "description": "Output format (default: markdown)",
                            "default": "markdown",
                        },
                    },
                },
            ),
            Tool(
                name="dental_radiation_log",
                description=(
                    "Log radiation-related events: patient exposure, equipment QA tests, "
                    "lead apron inspections, and equipment checks."
                ),
                inputSchema={
                    "type": "object",
                    "properties": {
                        "equipment_id": {
                            "type": "string",
                            "description": "Equipment identifier (e.g. 'Dexis-Op1', 'Pano-1', 'CBCT-1')",
                        },
                        "log_type": {
                            "type": "string",
                            "enum": ["exposure", "qa_test", "equipment_check", "lead_apron_check"],
                            "description": "Type of radiation log entry",
                        },
                        "details": {
                            "type": "string",
                            "description": "Description of the event (exposure details, test performed, findings)",
                        },
                        "result": {
                            "type": "string",
                            "description": "Result or outcome (e.g. 'pass', 'fail', 'within limits')",
                        },
                        "operator": {
                            "type": "string",
                            "description": "Name of person performing the action",
                        },
                    },
                    "required": ["equipment_id", "log_type", "details"],
                },
            ),
        ]

    @server.call_tool()
    async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
        if name == "dental_compliance_scan":
            result = _run_scan(
                target=arguments["target"],
                framework=arguments["framework"],
                store=store,
                practice_type=arguments.get("practice_type", "solo_gp"),
                state=arguments.get("state", ""),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "dental_compliance_status":
            framework = arguments.get("framework")
            if framework:
                scan = store.get_latest_scan(framework)
                if not scan:
                    return [TextContent(
                        type="text",
                        text=f"No scan results for {FRAMEWORK_NAMES.get(framework, framework)}. "
                             "Run dental_compliance_scan first.",
                    )]
                findings = store.get_non_compliant_findings(framework)
                status = {
                    "framework": FRAMEWORK_NAMES.get(framework, framework),
                    "last_scan": scan.get("timestamp", ""),
                    "score": scan.get("score", 0),
                    "total_checks": scan.get("total_checks", 0),
                    "compliant": scan.get("compliant", 0),
                    "non_compliant": scan.get("non_compliant", 0),
                    "findings_by_severity": {
                        "critical": len([f for f in findings if f["severity"] == "critical"]),
                        "high": len([f for f in findings if f["severity"] == "high"]),
                        "medium": len([f for f in findings if f["severity"] == "medium"]),
                        "low": len([f for f in findings if f["severity"] == "low"]),
                    },
                }
                return [TextContent(type="text", text=json.dumps(status, indent=2))]
            else:
                summaries = store.get_all_scan_summaries()
                if not summaries:
                    return [TextContent(
                        type="text",
                        text="No scan results yet. Run dental_compliance_scan first.",
                    )]
                result = {
                    "dental_compliance_posture": [
                        {
                            "framework": FRAMEWORK_NAMES.get(s["framework"], s["framework"]),
                            "last_scan": s.get("last_scan", ""),
                            "score": s.get("score", 0),
                            "total_checks": s.get("total_checks", 0),
                            "compliant": s.get("compliant", 0),
                            "non_compliant": s.get("non_compliant", 0),
                        }
                        for s in summaries
                    ]
                }
                return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "dental_sterilization_log":
            record = SterilizationRecord(
                autoclave_id=arguments["autoclave_id"],
                cycle_type=arguments["cycle_type"],
                temperature=arguments["temperature"],
                pressure=arguments["pressure"],
                duration_minutes=arguments["duration_minutes"],
                biological_indicator=arguments["biological_indicator"],
                chemical_indicator=arguments["chemical_indicator"],
                load_contents=arguments.get("load_contents", ""),
                operator=arguments.get("operator", ""),
                notes=arguments.get("notes", ""),
            )
            log_id = store.log_sterilization(record)

            # Check for concerning results
            warnings: list[str] = []
            if arguments["biological_indicator"] == "fail":
                warnings.append(
                    "CRITICAL: Biological indicator FAILED. Remove autoclave from service immediately. "
                    "Recall all items processed since last negative BI. Do not use until 3 consecutive "
                    "negative BIs obtained."
                )
            if arguments["chemical_indicator"] == "fail":
                warnings.append(
                    "WARNING: Chemical indicator failed. Do not use instruments from this load. "
                    "Re-clean and re-sterilize. Check autoclave function."
                )

            result = {
                "log_id": log_id,
                "status": "logged",
                "autoclave_id": arguments["autoclave_id"],
                "cycle_type": arguments["cycle_type"],
                "biological_indicator": arguments["biological_indicator"],
                "chemical_indicator": arguments["chemical_indicator"],
                "warnings": warnings,
            }
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "dental_inspection_prep":
            state = arguments["state"]
            practice_type = arguments.get("practice_type", "solo_gp")
            checklist = _generate_inspection_checklist(state, practice_type)
            return [TextContent(type="text", text=json.dumps(checklist, indent=2))]

        elif name == "dental_compliance_fix":
            framework = arguments.get("framework")
            findings = store.get_non_compliant_findings(framework)
            if not findings:
                msg = "No non-compliant findings"
                if framework:
                    msg += f" for {FRAMEWORK_NAMES.get(framework, framework)}"
                return [TextContent(type="text", text=msg + ". Run dental_compliance_scan first.")]

            actions = _generate_corrective_actions(findings)
            result = {
                "total_findings": len(findings),
                "corrective_actions": len(actions),
                "actions": actions,
            }
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "dental_compliance_report":
            framework = arguments.get("framework")
            fmt = arguments.get("format", "markdown")
            report = _generate_report(store, framework, fmt)
            return [TextContent(type="text", text=report)]

        elif name == "dental_radiation_log":
            record = RadiationRecord(
                equipment_id=arguments["equipment_id"],
                log_type=arguments["log_type"],
                details=arguments["details"],
                result=arguments.get("result", ""),
                operator=arguments.get("operator", ""),
            )
            log_id = store.log_radiation(record)

            # Check for concerning results
            warnings: list[str] = []
            if arguments["log_type"] == "lead_apron_check" and arguments.get("result") == "fail":
                warnings.append(
                    "WARNING: Lead apron inspection FAILED. Remove from service immediately. "
                    "Replace with new lead apron. Do not use for patient protection until replaced."
                )
            if arguments["log_type"] == "qa_test" and arguments.get("result") == "fail":
                warnings.append(
                    "WARNING: QA test FAILED. Remove equipment from clinical use. "
                    "Schedule service/calibration before returning to use."
                )

            result = {
                "log_id": log_id,
                "status": "logged",
                "equipment_id": arguments["equipment_id"],
                "log_type": arguments["log_type"],
                "warnings": warnings,
            }
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        return [TextContent(type="text", text=f"Unknown tool: {name}")]

    return server


async def _run() -> None:
    """Run the MCP server over stdio."""
    server = create_server()
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


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