"""Remediation Playbooks for MCP Security Findings.

Generates step-by-step remediation guides for each security finding,
with code examples, effort estimates, and reference links.

Supports output in text, JSON, Markdown, and HTML formats.
CLI integration: mcpaudit playbook <rule-id> | --all | --from-scan <file>

Zero external dependencies. Python 3.11+ stdlib only.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from .models import Finding, ScanResult, Severity


__all__ = [
    "Playbook",
    "PlaybookLibrary",
    "RemediationStep",
    "export_playbooks_html",
    "export_playbooks_json",
    "export_playbooks_markdown",
    "export_playbooks_text",
    "generate_all_playbooks",
    "generate_playbook",
    "generate_playbook_by_id",
    "generate_playbooks",
    "generate_playbooks_from_findings",
    "generate_playbooks_from_scan_data",
    "get_covered_rule_ids",
]


# ---------------------------------------------------------------------------
# Data Models
# ---------------------------------------------------------------------------

@dataclass
class RemediationStep:
    """A single step in a remediation playbook."""

    step_number: int
    action: str
    code_before: str = ""
    code_after: str = ""
    explanation: str = ""

    def to_dict(self) -> dict[str, Any]:
        result: dict[str, Any] = {
            "step_number": self.step_number,
            "action": self.action,
            "explanation": self.explanation,
        }
        if self.code_before:
            result["code_before"] = self.code_before
        if self.code_after:
            result["code_after"] = self.code_after
        return result


@dataclass
class Playbook:
    """A complete remediation playbook for a finding rule."""

    finding_rule_id: str
    title: str
    severity: Severity
    steps: list[RemediationStep] = field(default_factory=list)
    estimated_effort: str = ""  # e.g. "15 minutes", "1 hour"
    references: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "finding_rule_id": self.finding_rule_id,
            "title": self.title,
            "severity": self.severity.value,
            "steps": [s.to_dict() for s in self.steps],
            "estimated_effort": self.estimated_effort,
            "references": self.references,
        }


# ---------------------------------------------------------------------------
# Playbook Database (22 rule IDs with specific remediation)
#
# NOTE: code_before / code_after strings are educational examples showing
# vulnerable patterns and their secure replacements. They are NEVER executed.
# ---------------------------------------------------------------------------

def _build_playbook_db() -> dict[str, dict[str, Any]]:
    """Build the playbook database mapping rule IDs to playbook data."""
    db: dict[str, dict[str, Any]] = {}

    # --- CMD-001 ---
    db["CMD-001"] = {
        "title": "Command Injection via subprocess",
        "severity": "CRITICAL",
        "estimated_effort": "30 minutes",
        "references": [
            "CWE-78: Improper Neutralization of Special Elements used in an OS Command",
            "OWASP Agentic Security: Tool Misuse",
        ],
        "steps": [
            {
                "action": "Replace shell=True with shell=False and use list arguments",
                "code_before": 'subprocess.run(f"ls {user_input}", shell=True)',
                "code_after": 'subprocess.run(["ls", user_input], shell=False)',
                "explanation": (
                    "Using shell=True passes the command through the system shell, "
                    "allowing injection of arbitrary commands via metacharacters."
                ),
            },
            {
                "action": "Use shlex.quote for any dynamic arguments",
                "code_before": 'cmd = f"grep {pattern} {filename}"',
                "code_after": 'import shlex\ncmd = f"grep {shlex.quote(pattern)} {shlex.quote(filename)}"',
                "explanation": (
                    "shlex.quote escapes shell metacharacters, preventing injection."
                ),
            },
            {
                "action": "Validate inputs against an allowlist",
                "code_before": 'subprocess.run(["git", action], check=True)',
                "code_after": (
                    'ALLOWED = {"status", "log", "diff"}\n'
                    'if action not in ALLOWED:\n'
                    '    raise ValueError(f"Invalid: {action}")\n'
                    'subprocess.run(["git", action], check=True)'
                ),
                "explanation": "Allowlisting restricts inputs to a known-safe set.",
            },
        ],
    }

    # --- CMD-002 ---
    db["CMD-002"] = {
        "title": "Command Injection via os module calls",
        "severity": "CRITICAL",
        "estimated_effort": "30 minutes",
        "references": [
            "CWE-78: OS Command Injection",
            "OWASP Agentic Security: Tool Misuse",
        ],
        "steps": [
            {
                "action": "Replace os module shell calls with subprocess.run using list args",
                "code_before": "# Vulnerable: uses shell execution",
                "code_after": 'subprocess.run(["rm", filename], check=True)',
                "explanation": (
                    "Shell-based os calls always invoke the shell. "
                    "subprocess.run with a list argument avoids shell interpretation."
                ),
            },
            {
                "action": "Use Python file I/O instead of shell commands",
                "code_before": "# Shell call to read file contents",
                "code_after": (
                    "from pathlib import Path\n"
                    "safe_path = Path(filepath).resolve()\n"
                    "if not str(safe_path).startswith('/allowed/dir'):\n"
                    "    raise ValueError('Invalid path')\n"
                    "content = safe_path.read_text()"
                ),
                "explanation": "Use Python's file I/O instead of shell commands when possible.",
            },
        ],
    }

    # --- PATH-001 ---
    db["PATH-001"] = {
        "title": "Path Traversal Vulnerability",
        "severity": "HIGH",
        "estimated_effort": "20 minutes",
        "references": [
            "CWE-22: Improper Limitation of a Pathname to a Restricted Directory",
            "OWASP Agentic Security: Excessive Agency",
        ],
        "steps": [
            {
                "action": "Resolve and validate paths against a base directory",
                "code_before": 'open(base_dir + "/" + user_path)',
                "code_after": (
                    "from pathlib import Path\n"
                    "safe = (Path(base_dir) / user_path).resolve()\n"
                    "if not str(safe).startswith(str(Path(base_dir).resolve())):\n"
                    "    raise ValueError('Path traversal detected')\n"
                    "open(safe)"
                ),
                "explanation": (
                    "Resolving the path eliminates .. components, "
                    "then checking the prefix ensures it stays within the base."
                ),
            },
            {
                "action": "Reject paths containing traversal sequences",
                "code_before": 'filepath = request.args["file"]',
                "code_after": (
                    'filepath = request.args["file"]\n'
                    'if ".." in filepath or filepath.startswith("/"):\n'
                    '    raise ValueError("Invalid path")'
                ),
                "explanation": "Quick input validation catches obvious traversal attempts.",
            },
        ],
    }

    # --- SQL-001 ---
    db["SQL-001"] = {
        "title": "SQL Injection Vulnerability",
        "severity": "CRITICAL",
        "estimated_effort": "20 minutes",
        "references": [
            "CWE-89: SQL Injection",
            "OWASP Top 10 A03:2021 Injection",
        ],
        "steps": [
            {
                "action": "Use parameterized queries instead of string formatting",
                "code_before": 'cursor.execute(f"SELECT * FROM users WHERE name = \'{name}\'")',
                "code_after": 'cursor.execute("SELECT * FROM users WHERE name = ?", (name,))',
                "explanation": (
                    "Parameterized queries separate SQL structure from data, "
                    "making injection impossible."
                ),
            },
            {
                "action": "Use an ORM for complex queries",
                "code_before": "# String-formatted SQL query",
                "code_after": "products = session.query(Product).filter_by(category=cat).all()",
                "explanation": "ORMs handle parameterization automatically.",
            },
        ],
    }

    # --- AUTH-001 ---
    db["AUTH-001"] = {
        "title": "Missing Authentication Check",
        "severity": "HIGH",
        "estimated_effort": "45 minutes",
        "references": [
            "CWE-306: Missing Authentication for Critical Function",
            "OWASP Agentic Security: Lack of Authentication",
        ],
        "steps": [
            {
                "action": "Add authentication middleware/decorator",
                "code_before": (
                    "@app.tool()\n"
                    "def delete_user(user_id: str):\n"
                    "    db.delete(user_id)"
                ),
                "code_after": (
                    "@app.tool()\n"
                    "@require_auth\n"
                    "def delete_user(user_id: str, ctx: Context):\n"
                    "    if not ctx.authenticated:\n"
                    "        raise PermissionError('Authentication required')\n"
                    "    db.delete(user_id)"
                ),
                "explanation": (
                    "Every tool that performs state changes must verify "
                    "the caller is authenticated."
                ),
            },
            {
                "action": "Implement token validation",
                "code_before": "# No authentication check",
                "code_after": (
                    "def require_auth(func):\n"
                    "    def wrapper(*args, **kwargs):\n"
                    "        token = kwargs.get('auth_token')\n"
                    "        if not validate_token(token):\n"
                    "            raise PermissionError('Invalid token')\n"
                    "        return func(*args, **kwargs)\n"
                    "    return wrapper"
                ),
                "explanation": "Token-based auth ensures only authorized callers can invoke tools.",
            },
        ],
    }

    # --- SSRF-001 ---
    db["SSRF-001"] = {
        "title": "Server-Side Request Forgery (SSRF)",
        "severity": "HIGH",
        "estimated_effort": "30 minutes",
        "references": [
            "CWE-918: Server-Side Request Forgery (SSRF)",
            "OWASP Top 10 A10:2021 SSRF",
        ],
        "steps": [
            {
                "action": "Validate URLs against an allowlist of domains",
                "code_before": "# Fetching user-provided URL without validation",
                "code_after": (
                    "from urllib.parse import urlparse\n"
                    "ALLOWED = {'api.example.com', 'cdn.example.com'}\n"
                    "parsed = urlparse(user_url)\n"
                    "if parsed.hostname not in ALLOWED:\n"
                    "    raise ValueError('Domain not allowed')"
                ),
                "explanation": "Domain allowlisting restricts fetch targets to trusted hosts.",
            },
            {
                "action": "Block private/internal IP ranges",
                "code_before": "# No IP validation",
                "code_after": (
                    "import ipaddress\n"
                    "def is_safe_url(url: str) -> bool:\n"
                    "    parsed = urlparse(url)\n"
                    "    try:\n"
                    "        ip = ipaddress.ip_address(parsed.hostname)\n"
                    "        return ip.is_global\n"
                    "    except ValueError:\n"
                    "        return True  # hostname, not IP"
                ),
                "explanation": "Block requests to private IP ranges (10.x, 172.16.x, 192.168.x, 127.x).",
            },
        ],
    }

    # --- DESER-001 ---
    db["DESER-001"] = {
        "title": "Unsafe Deserialization",
        "severity": "CRITICAL",
        "estimated_effort": "20 minutes",
        "references": [
            "CWE-502: Deserialization of Untrusted Data",
            "OWASP Top 10 A08:2021 Software and Data Integrity Failures",
        ],
        "steps": [
            {
                "action": "Replace unsafe deserialization with json.loads",
                "code_before": "# Deserializing untrusted data with unsafe method",
                "code_after": "data = json.loads(user_data)",
                "explanation": (
                    "Unsafe deserialization can execute arbitrary code. "
                    "JSON is a safe data format."
                ),
            },
            {
                "action": "Replace yaml.load with yaml.safe_load",
                "code_before": "config = yaml.load(content)",
                "code_after": "config = yaml.safe_load(content)",
                "explanation": "yaml.load can instantiate arbitrary objects. safe_load restricts to basic types.",
            },
        ],
    }

    # --- ERR-001 ---
    db["ERR-001"] = {
        "title": "Information Disclosure via Error Messages",
        "severity": "MEDIUM",
        "estimated_effort": "15 minutes",
        "references": [
            "CWE-209: Error Message Containing Sensitive Information",
            "OWASP Agentic Security: Excessive Information Disclosure",
        ],
        "steps": [
            {
                "action": "Catch exceptions and return sanitized error messages",
                "code_before": (
                    "try:\n"
                    "    result = process(data)\n"
                    "except Exception as e:\n"
                    "    return str(e)  # leaks internals"
                ),
                "code_after": (
                    "import logging\n"
                    "logger = logging.getLogger(__name__)\n"
                    "try:\n"
                    "    result = process(data)\n"
                    "except Exception as e:\n"
                    "    logger.exception('Processing failed')\n"
                    "    return 'An internal error occurred'"
                ),
                "explanation": (
                    "Log the full error for debugging but return a generic message to clients."
                ),
            },
        ],
    }

    # --- LOG-001 ---
    db["LOG-001"] = {
        "title": "Sensitive Data in Logs",
        "severity": "MEDIUM",
        "estimated_effort": "20 minutes",
        "references": [
            "CWE-532: Sensitive Information in Log File",
        ],
        "steps": [
            {
                "action": "Redact sensitive fields before logging",
                "code_before": 'logger.info(f"User: {username}, password: {password}")',
                "code_after": 'logger.info(f"User: {username}, password: ***")',
                "explanation": "Never log credentials, tokens, or PII.",
            },
            {
                "action": "Use structured logging with field filtering",
                "code_before": 'logging.info(f"Request: {request_data}")',
                "code_after": (
                    "REDACT = {'password', 'token', 'secret', 'api_key'}\n"
                    "safe = {k: '***' if k in REDACT else v for k, v in data.items()}\n"
                    "logging.info(f'Request: {safe}')"
                ),
                "explanation": "Centralized redaction prevents sensitive data leaks through logs.",
            },
        ],
    }

    # --- PERM-001 ---
    db["PERM-001"] = {
        "title": "Overly Permissive Tool Access",
        "severity": "HIGH",
        "estimated_effort": "30 minutes",
        "references": [
            "CWE-284: Improper Access Control",
            "OWASP Agentic Security: Excessive Agency",
        ],
        "steps": [
            {
                "action": "Implement role-based access control for tools",
                "code_before": (
                    "@app.tool()\n"
                    "def admin_action(command: str):\n"
                    "    return execute(command)"
                ),
                "code_after": (
                    "@app.tool()\n"
                    "def admin_action(command: str, ctx: Context):\n"
                    "    if 'admin' not in ctx.user_roles:\n"
                    "        raise PermissionError('Admin required')\n"
                    "    return execute(command)"
                ),
                "explanation": "Privileged tools must verify caller permissions.",
            },
        ],
    }

    # --- RATE-001 ---
    db["RATE-001"] = {
        "title": "Missing Rate Limiting",
        "severity": "MEDIUM",
        "estimated_effort": "30 minutes",
        "references": [
            "CWE-770: Resource Allocation Without Limits",
            "OWASP Agentic Security: Denial of Service",
        ],
        "steps": [
            {
                "action": "Add rate limiting to tool handlers",
                "code_before": (
                    "@app.tool()\n"
                    "def search(query: str):\n"
                    "    return db.search(query)"
                ),
                "code_after": (
                    "from collections import defaultdict\n"
                    "import time\n"
                    "_limits: dict[str, list] = defaultdict(list)\n\n"
                    "@app.tool()\n"
                    "def search(query: str, ctx: Context):\n"
                    "    now = time.time()\n"
                    "    window = [t for t in _limits[ctx.id] if now - t < 60]\n"
                    "    if len(window) >= 10:\n"
                    "        raise RuntimeError('Rate limit exceeded')\n"
                    "    _limits[ctx.id] = window + [now]\n"
                    "    return db.search(query)"
                ),
                "explanation": "Rate limiting prevents abuse and resource exhaustion.",
            },
        ],
    }

    # --- SEC-001 ---
    db["SEC-001"] = {
        "title": "Hardcoded Secrets",
        "severity": "CRITICAL",
        "estimated_effort": "15 minutes",
        "references": [
            "CWE-798: Hard-coded Credentials",
            "OWASP Top 10 A07:2021 Authentication Failures",
        ],
        "steps": [
            {
                "action": "Move secrets to environment variables",
                "code_before": 'API_KEY = "sk-abc123secret456"',
                "code_after": (
                    "import os\n"
                    "API_KEY = os.environ.get('API_KEY')\n"
                    "if not API_KEY:\n"
                    "    raise RuntimeError('API_KEY required')"
                ),
                "explanation": "Environment variables keep secrets out of the codebase.",
            },
            {
                "action": "Rotate the exposed secret immediately",
                "code_before": "# Secret was committed to version control",
                "code_after": (
                    "# 1. Revoke the old key\n"
                    "# 2. Generate a new key\n"
                    "# 3. Set as environment variable"
                ),
                "explanation": "Any committed secret must be considered compromised.",
            },
        ],
    }

    # --- CRYPTO-001 ---
    db["CRYPTO-001"] = {
        "title": "Weak Cryptographic Algorithm",
        "severity": "HIGH",
        "estimated_effort": "20 minutes",
        "references": [
            "CWE-327: Broken or Risky Cryptographic Algorithm",
        ],
        "steps": [
            {
                "action": "Replace MD5/SHA1 with SHA-256 or stronger",
                "code_before": "digest = hashlib.md5(data).hexdigest()",
                "code_after": "digest = hashlib.sha256(data).hexdigest()",
                "explanation": "MD5 and SHA-1 have known collision attacks. Use SHA-256+.",
            },
        ],
    }

    # --- FILE-001 ---
    db["FILE-001"] = {
        "title": "Insecure File Operations",
        "severity": "MEDIUM",
        "estimated_effort": "15 minutes",
        "references": [
            "CWE-732: Incorrect Permission Assignment",
        ],
        "steps": [
            {
                "action": "Set restrictive file permissions",
                "code_before": "open(filepath, 'w').write(sensitive_data)",
                "code_after": (
                    "import os\n"
                    "fd = os.open(filepath, os.O_WRONLY | os.O_CREAT, 0o600)\n"
                    "with os.fdopen(fd, 'w') as f:\n"
                    "    f.write(sensitive_data)"
                ),
                "explanation": "0o600 ensures only the owner can read/write the file.",
            },
        ],
    }

    # --- INFO-001 ---
    db["INFO-001"] = {
        "title": "Information Exposure through Debug Mode",
        "severity": "MEDIUM",
        "estimated_effort": "10 minutes",
        "references": [
            "CWE-215: Sensitive Information in Debugging Code",
        ],
        "steps": [
            {
                "action": "Disable debug mode in production",
                "code_before": "app.run(debug=True)",
                "code_after": (
                    "import os\n"
                    "debug = os.environ.get('DEBUG', 'false').lower() == 'true'\n"
                    "app.run(debug=debug)"
                ),
                "explanation": "Debug mode exposes stack traces. Use env vars to control it.",
            },
        ],
    }

    # --- VAL-001 ---
    db["VAL-001"] = {
        "title": "Missing Input Validation",
        "severity": "MEDIUM",
        "estimated_effort": "20 minutes",
        "references": [
            "CWE-20: Improper Input Validation",
            "OWASP Agentic Security: Tool Input Validation",
        ],
        "steps": [
            {
                "action": "Add type and range validation to tool inputs",
                "code_before": (
                    "@app.tool()\ndef process(data: str):\n    return execute(data)"
                ),
                "code_after": (
                    "@app.tool()\ndef process(data: str):\n"
                    "    if not isinstance(data, str) or len(data) > 10000:\n"
                    "        raise ValueError('Invalid input')\n"
                    "    return execute(data)"
                ),
                "explanation": "Validate all inputs for type, length, and format.",
            },
        ],
    }

    # --- RES-001 ---
    db["RES-001"] = {
        "title": "Resource Leak",
        "severity": "MEDIUM",
        "estimated_effort": "10 minutes",
        "references": [
            "CWE-404: Improper Resource Shutdown or Release",
        ],
        "steps": [
            {
                "action": "Use context managers for resource handling",
                "code_before": "f = open(path)\ndata = f.read()",
                "code_after": "with open(path) as f:\n    data = f.read()",
                "explanation": "Context managers ensure resources are released even on exceptions.",
            },
        ],
    }

    # --- MCP-001 ---
    db["MCP-001"] = {
        "title": "Tool Without Description",
        "severity": "LOW",
        "estimated_effort": "5 minutes",
        "references": [
            "MCP Specification: Tool Metadata",
            "OWASP Agentic Security: Transparency",
        ],
        "steps": [
            {
                "action": "Add a descriptive docstring to the tool function",
                "code_before": "@app.tool()\ndef fetch(url: str):\n    return get(url).text",
                "code_after": (
                    "@app.tool()\ndef fetch(url: str):\n"
                    '    """Fetch content from a URL."""\n'
                    "    return get(url).text"
                ),
                "explanation": "Tool descriptions are sent to the LLM. Clear descriptions prevent misuse.",
            },
        ],
    }

    # --- MCP-005 ---
    db["MCP-005"] = {
        "title": "No Error Handling in Tool",
        "severity": "MEDIUM",
        "estimated_effort": "15 minutes",
        "references": [
            "CWE-755: Improper Handling of Exceptional Conditions",
            "MCP Specification: Error Handling",
        ],
        "steps": [
            {
                "action": "Wrap tool logic in try/except with sanitized errors",
                "code_before": "@app.tool()\ndef query(sql: str):\n    return db.execute(sql)",
                "code_after": (
                    "@app.tool()\ndef query(sql: str):\n"
                    "    try:\n"
                    "        return db.execute(sql)\n"
                    "    except Exception:\n"
                    "        return {'error': 'Query failed'}"
                ),
                "explanation": "Catch and return sanitized error responses.",
            },
        ],
    }

    # --- CORS-001 ---
    db["CORS-001"] = {
        "title": "Permissive CORS Configuration",
        "severity": "MEDIUM",
        "estimated_effort": "10 minutes",
        "references": [
            "CWE-942: Permissive Cross-domain Policy",
        ],
        "steps": [
            {
                "action": "Restrict CORS origins to specific domains",
                "code_before": 'allow_origins=["*"]',
                "code_after": 'allow_origins=["https://app.example.com"]',
                "explanation": "Wildcard CORS allows any website to make requests.",
            },
        ],
    }

    # --- REDOS-001 ---
    db["REDOS-001"] = {
        "title": "Regular Expression Denial of Service (ReDoS)",
        "severity": "MEDIUM",
        "estimated_effort": "20 minutes",
        "references": [
            "CWE-1333: Inefficient Regular Expression Complexity",
        ],
        "steps": [
            {
                "action": "Simplify regex patterns and add input length limits",
                "code_before": 're.match(r"(a+)+$", user_input)',
                "code_after": (
                    "if len(user_input) > 1000:\n"
                    "    raise ValueError('Input too long')\n"
                    "re.match(r'a+$', user_input)"
                ),
                "explanation": "Nested quantifiers cause exponential backtracking.",
            },
        ],
    }

    # --- IDOR-001 ---
    db["IDOR-001"] = {
        "title": "Insecure Direct Object Reference",
        "severity": "HIGH",
        "estimated_effort": "30 minutes",
        "references": [
            "CWE-639: Authorization Bypass Through User-Controlled Key",
            "OWASP Top 10 A01:2021 Broken Access Control",
        ],
        "steps": [
            {
                "action": "Verify resource ownership before access",
                "code_before": (
                    "@app.tool()\n"
                    "def get_document(doc_id: str):\n"
                    "    return db.get(doc_id)"
                ),
                "code_after": (
                    "@app.tool()\n"
                    "def get_document(doc_id: str, ctx: Context):\n"
                    "    doc = db.get(doc_id)\n"
                    "    if doc.owner_id != ctx.user_id:\n"
                    "        raise PermissionError('Access denied')\n"
                    "    return doc"
                ),
                "explanation": "Always verify ownership or permission before returning resources.",
            },
        ],
    }

    # --- MCP-006 ---
    db["MCP-006"] = {
        "title": "Excessive Tool Permissions",
        "severity": "HIGH",
        "estimated_effort": "30 minutes",
        "references": [
            "CWE-250: Execution with Unnecessary Privileges",
            "OWASP Agentic Security: Excessive Agency",
        ],
        "steps": [
            {
                "action": "Split broad tools into narrowly-scoped ones",
                "code_before": (
                    "@app.tool()\n"
                    "def manage_files(action: str, path: str):\n"
                    "    # Does everything: read, write, delete"
                ),
                "code_after": (
                    "@app.tool()\n"
                    "def read_file(path: str):\n"
                    '    """Read a file."""\n'
                    "    return safe_read(path)\n\n"
                    "@app.tool()\n"
                    "def delete_file(path: str, ctx: Context):\n"
                    '    """Delete (admin only)."""\n'
                    "    require_role(ctx, 'admin')\n"
                    "    safe_delete(path)"
                ),
                "explanation": (
                    "Least privilege: split broad tools and require elevated permissions "
                    "for destructive operations."
                ),
            },
        ],
    }

    return db


_PLAYBOOK_DB: dict[str, dict[str, Any]] | None = None


def _get_db() -> dict[str, dict[str, Any]]:
    """Get the playbook database (lazy init)."""
    global _PLAYBOOK_DB
    if _PLAYBOOK_DB is None:
        _PLAYBOOK_DB = _build_playbook_db()
    return _PLAYBOOK_DB


# ---------------------------------------------------------------------------
# Playbook Generation
# ---------------------------------------------------------------------------

def generate_playbook(finding: Finding) -> Playbook | None:
    """Generate a remediation playbook for a single finding.

    Args:
        finding: The security finding to remediate.

    Returns:
        Playbook if a playbook exists for this rule ID, None otherwise.
    """
    db = _get_db()
    rule_data = db.get(finding.rule_id)
    if rule_data is None:
        return None

    steps: list[RemediationStep] = []
    for i, step_data in enumerate(rule_data.get("steps", []), 1):
        steps.append(RemediationStep(
            step_number=i,
            action=step_data["action"],
            code_before=step_data.get("code_before", ""),
            code_after=step_data.get("code_after", ""),
            explanation=step_data.get("explanation", ""),
        ))

    severity_str = rule_data.get("severity", finding.severity.value)
    try:
        severity = Severity[severity_str]
    except KeyError:
        severity = finding.severity

    return Playbook(
        finding_rule_id=finding.rule_id,
        title=rule_data.get("title", finding.title),
        severity=severity,
        steps=steps,
        estimated_effort=rule_data.get("estimated_effort", ""),
        references=rule_data.get("references", []),
    )


def generate_playbooks(scan_result: ScanResult) -> list[Playbook]:
    """Generate playbooks for all findings in a scan result.

    Args:
        scan_result: The scan result containing findings.

    Returns:
        List of playbooks (one per unique rule ID with findings).
    """
    seen_rules: set[str] = set()
    playbooks: list[Playbook] = []

    for finding in scan_result.findings:
        if finding.rule_id in seen_rules:
            continue
        seen_rules.add(finding.rule_id)
        pb = generate_playbook(finding)
        if pb is not None:
            playbooks.append(pb)

    # Sort by severity (CRITICAL first)
    severity_order = {
        Severity.CRITICAL: 0,
        Severity.HIGH: 1,
        Severity.MEDIUM: 2,
        Severity.LOW: 3,
        Severity.INFO: 4,
    }
    playbooks.sort(key=lambda p: severity_order.get(p.severity, 5))
    return playbooks


def generate_playbooks_from_findings(findings: list[Finding]) -> list[Playbook]:
    """Generate playbooks from a flat list of findings.

    Args:
        findings: List of findings.

    Returns:
        List of playbooks for unique rule IDs.
    """
    result = ScanResult(file="", findings=findings)
    return generate_playbooks(result)


def get_covered_rule_ids() -> list[str]:
    """Return the list of rule IDs that have playbooks.

    Returns:
        Sorted list of rule IDs with playbook coverage.
    """
    return sorted(_get_db().keys())


# ---------------------------------------------------------------------------
# Export: HTML
# ---------------------------------------------------------------------------

def export_playbooks_html(playbooks: list[Playbook]) -> str:
    """Export playbooks as an HTML page with collapsible sections.

    Args:
        playbooks: List of playbooks to export.

    Returns:
        Complete HTML document string.
    """
    severity_colors = {
        "CRITICAL": "#ef4444",
        "HIGH": "#f97316",
        "MEDIUM": "#eab308",
        "LOW": "#22c55e",
        "INFO": "#3b82f6",
    }

    playbook_html = ""
    for pb in playbooks:
        color = severity_colors.get(pb.severity.value, "#6b7280")
        steps_html = ""
        for step in pb.steps:
            code_diff = ""
            if step.code_before or step.code_after:
                before = _html_escape(step.code_before) if step.code_before else ""
                after = _html_escape(step.code_after) if step.code_after else ""
                code_diff = f"""
                <div class="code-diff">
                    <div class="code-before">
                        <div class="code-label">Before</div>
                        <pre><code>{before}</code></pre>
                    </div>
                    <div class="code-after">
                        <div class="code-label">After</div>
                        <pre><code>{after}</code></pre>
                    </div>
                </div>"""

            explanation_html = ""
            if step.explanation:
                explanation_html = (
                    f'<p class="explanation">{_html_escape(step.explanation)}</p>'
                )

            steps_html += f"""
            <div class="step">
                <div class="step-number">Step {step.step_number}</div>
                <div class="step-action">{_html_escape(step.action)}</div>
                {code_diff}
                {explanation_html}
            </div>"""

        refs_html = ""
        if pb.references:
            refs_items = "".join(
                f"<li>{_html_escape(ref)}</li>" for ref in pb.references
            )
            refs_html = (
                '<div class="references"><strong>References:</strong>'
                f"<ul>{refs_items}</ul></div>"
            )

        effort_html = ""
        if pb.estimated_effort:
            effort_html = (
                f'<span class="effort">'
                f"Estimated: {_html_escape(pb.estimated_effort)}</span>"
            )

        playbook_html += f"""
        <details class="playbook" open>
            <summary>
                <span class="severity-badge" style="background:{color}">{pb.severity.value}</span>
                <span class="rule-id">{_html_escape(pb.finding_rule_id)}</span>
                {_html_escape(pb.title)}
                {effort_html}
            </summary>
            <div class="playbook-body">
                {steps_html}
                {refs_html}
            </div>
        </details>"""

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>mcpaudit Remediation Playbooks</title>
<style>
* {{ margin: 0; padding: 0; box-sizing: border-box; }}
body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
       background: #0f172a; color: #e2e8f0; padding: 24px; max-width: 1200px; margin: 0 auto; }}
h1 {{ font-size: 1.5rem; margin-bottom: 8px; }}
.subtitle {{ color: #94a3b8; margin-bottom: 24px; }}
.playbook {{ background: #1e293b; border-radius: 12px; margin-bottom: 16px;
             border: 1px solid #334155; }}
.playbook summary {{ cursor: pointer; padding: 16px 20px; font-size: 1rem;
                     display: flex; align-items: center; gap: 12px; list-style: none; }}
.playbook summary::-webkit-details-marker {{ display: none; }}
.playbook[open] summary {{ border-bottom: 1px solid #334155; }}
.severity-badge {{ padding: 2px 8px; border-radius: 4px; font-size: 0.75rem;
                   font-weight: 700; color: #fff; text-transform: uppercase; }}
.rule-id {{ color: #94a3b8; font-family: monospace; font-size: 0.85rem; }}
.effort {{ margin-left: auto; color: #94a3b8; font-size: 0.8rem; }}
.playbook-body {{ padding: 20px; }}
.step {{ margin-bottom: 20px; padding-bottom: 20px; border-bottom: 1px solid #334155; }}
.step:last-child {{ border-bottom: none; margin-bottom: 0; padding-bottom: 0; }}
.step-number {{ font-size: 0.75rem; color: #3b82f6; text-transform: uppercase;
                font-weight: 700; margin-bottom: 4px; }}
.step-action {{ font-weight: 600; margin-bottom: 8px; }}
.code-diff {{ display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin: 12px 0; }}
.code-before, .code-after {{ background: #0f172a; border-radius: 8px; padding: 12px; overflow-x: auto; }}
.code-before {{ border-left: 3px solid #ef4444; }}
.code-after {{ border-left: 3px solid #22c55e; }}
.code-label {{ font-size: 0.7rem; text-transform: uppercase; color: #64748b;
               margin-bottom: 8px; font-weight: 700; }}
pre {{ font-family: 'SF Mono', 'Fira Code', monospace; font-size: 0.8rem;
       line-height: 1.5; white-space: pre-wrap; word-break: break-word; }}
.explanation {{ color: #94a3b8; font-size: 0.9rem; margin-top: 8px; line-height: 1.5; }}
.references {{ margin-top: 16px; font-size: 0.85rem; color: #94a3b8; }}
.references ul {{ margin-top: 4px; padding-left: 20px; }}
.references li {{ margin-bottom: 2px; }}
.footer {{ text-align: center; color: #475569; margin-top: 32px; font-size: 0.8rem; }}
.stats {{ display: flex; gap: 16px; margin-bottom: 24px; }}
.stat {{ background: #1e293b; border-radius: 8px; padding: 12px 20px; }}
.stat-label {{ font-size: 0.7rem; text-transform: uppercase; color: #94a3b8; }}
.stat-value {{ font-size: 1.5rem; font-weight: 700; }}
@media (max-width: 768px) {{ .code-diff {{ grid-template-columns: 1fr; }} }}
</style>
</head>
<body>
<h1>mcpaudit Remediation Playbooks</h1>
<p class="subtitle">Step-by-step remediation guides for security findings</p>
<div class="stats">
    <div class="stat">
        <div class="stat-label">Playbooks</div>
        <div class="stat-value">{len(playbooks)}</div>
    </div>
    <div class="stat">
        <div class="stat-label">Total Steps</div>
        <div class="stat-value">{sum(len(p.steps) for p in playbooks)}</div>
    </div>
</div>
{playbook_html}
<p class="footer">Generated by mcpaudit</p>
</body>
</html>"""


# ---------------------------------------------------------------------------
# Export: Markdown
# ---------------------------------------------------------------------------

def export_playbooks_markdown(playbooks: list[Playbook]) -> str:
    """Export playbooks as a Markdown document.

    Args:
        playbooks: List of playbooks to export.

    Returns:
        Markdown string.
    """
    lines: list[str] = []
    lines.append("# mcpaudit Remediation Playbooks\n")
    lines.append(
        f"**{len(playbooks)} playbooks** | "
        f"**{sum(len(p.steps) for p in playbooks)} total steps**\n"
    )

    for pb in playbooks:
        sev_label = {
            Severity.CRITICAL: "[CRITICAL]",
            Severity.HIGH: "[HIGH]",
            Severity.MEDIUM: "[MEDIUM]",
            Severity.LOW: "[LOW]",
            Severity.INFO: "[INFO]",
        }.get(pb.severity, "")

        lines.append(f"## {sev_label} {pb.finding_rule_id}: {pb.title}\n")

        if pb.estimated_effort:
            lines.append(f"**Estimated effort:** {pb.estimated_effort}\n")

        for step in pb.steps:
            lines.append(f"### Step {step.step_number}: {step.action}\n")

            if step.code_before:
                lines.append("**Before:**")
                lines.append(f"```python\n{step.code_before}\n```\n")

            if step.code_after:
                lines.append("**After:**")
                lines.append(f"```python\n{step.code_after}\n```\n")

            if step.explanation:
                lines.append(f"> {step.explanation}\n")

        if pb.references:
            lines.append("**References:**\n")
            for ref in pb.references:
                lines.append(f"- {ref}")
            lines.append("")

        lines.append("---\n")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Export: JSON
# ---------------------------------------------------------------------------

def export_playbooks_json(playbooks: list[Playbook]) -> str:
    """Export playbooks as JSON.

    Args:
        playbooks: List of playbooks to export.

    Returns:
        JSON string.
    """
    return json.dumps(
        [pb.to_dict() for pb in playbooks],
        indent=2,
        ensure_ascii=False,
    )


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def _html_escape(text: str) -> str:
    """Escape HTML special characters."""
    return (
        text
        .replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
        .replace("'", "&#39;")
    )


# ---------------------------------------------------------------------------
# Extended API (rule-ID based, for CLI)
# ---------------------------------------------------------------------------

def generate_playbook_by_id(rule_id: str) -> Playbook | None:
    """Generate a playbook by rule ID, without needing a Finding.

    Args:
        rule_id: The rule identifier (e.g. 'CMD-001').

    Returns:
        Playbook if the rule ID has playbook coverage, None otherwise.
    """
    db = _get_db()
    rid = rule_id.upper()
    rule_data = db.get(rid)
    if rule_data is None:
        return None

    severity_str = rule_data.get("severity", "MEDIUM")
    try:
        severity = Severity[severity_str]
    except KeyError:
        severity = Severity.MEDIUM

    steps: list[RemediationStep] = []
    for i, step_data in enumerate(rule_data.get("steps", []), 1):
        steps.append(RemediationStep(
            step_number=i,
            action=step_data["action"],
            code_before=step_data.get("code_before", ""),
            code_after=step_data.get("code_after", ""),
            explanation=step_data.get("explanation", ""),
        ))

    return Playbook(
        finding_rule_id=rid,
        title=rule_data.get("title", rid),
        severity=severity,
        steps=steps,
        estimated_effort=rule_data.get("estimated_effort", ""),
        references=rule_data.get("references", []),
    )


def generate_all_playbooks() -> list[Playbook]:
    """Generate playbooks for all rule IDs with coverage.

    Returns:
        Sorted list of all available playbooks.
    """
    db = _get_db()
    playbooks: list[Playbook] = []
    for rid in sorted(db.keys()):
        pb = generate_playbook_by_id(rid)
        if pb is not None:
            playbooks.append(pb)
    return playbooks


def generate_playbooks_from_scan_data(
    scan_data: dict[str, Any],
) -> list[Playbook]:
    """Generate playbooks for findings in a scan result dict.

    Parses the scan JSON data, extracts unique rule IDs, and returns
    playbooks for those that have coverage.

    Args:
        scan_data: Scan result dict (from mcpaudit JSON output).

    Returns:
        List of playbooks for findings in the scan.
    """
    rule_ids: set[str] = set()

    # Standard format: {"results": [{"findings": [...]}]}
    if "results" in scan_data:
        for r in scan_data["results"]:
            if isinstance(r, dict) and "findings" in r:
                for f in r["findings"]:
                    if isinstance(f, dict) and "rule_id" in f:
                        rule_ids.add(f["rule_id"])

    # Flat format: {"findings": [...]}
    if "findings" in scan_data and isinstance(scan_data["findings"], list):
        for f in scan_data["findings"]:
            if isinstance(f, dict) and "rule_id" in f:
                rule_ids.add(f["rule_id"])

    playbooks: list[Playbook] = []
    for rid in sorted(rule_ids):
        pb = generate_playbook_by_id(rid)
        if pb is not None:
            playbooks.append(pb)
    return playbooks


# ---------------------------------------------------------------------------
# PlaybookLibrary
# ---------------------------------------------------------------------------

class PlaybookLibrary:
    """Container for playbooks with lookup and filtering."""

    def __init__(self, playbooks: list[Playbook] | None = None) -> None:
        if playbooks is None:
            playbooks = generate_all_playbooks()
        self._playbooks = {pb.finding_rule_id: pb for pb in playbooks}

    @property
    def count(self) -> int:
        return len(self._playbooks)

    def get(self, rule_id: str) -> Playbook | None:
        """Get playbook by rule ID."""
        return self._playbooks.get(rule_id.upper())

    def all(self) -> list[Playbook]:
        """Return all playbooks sorted by rule ID."""
        return sorted(self._playbooks.values(), key=lambda p: p.finding_rule_id)

    def by_severity(self, severity: str) -> list[Playbook]:
        """Filter playbooks by severity."""
        sev_str = severity.upper()
        try:
            sev = Severity[sev_str]
        except KeyError:
            return []
        return [p for p in self.all() if p.severity == sev]

    def total_estimated_effort(self) -> str:
        """Summarize total estimated effort across all playbooks."""
        total_min = 0
        for pb in self._playbooks.values():
            effort = pb.estimated_effort.lower()
            if "hour" in effort:
                try:
                    hours = int(effort.split()[0])
                    total_min += hours * 60
                except (ValueError, IndexError):
                    pass
            elif "minute" in effort:
                try:
                    mins = int(effort.split()[0])
                    total_min += mins
                except (ValueError, IndexError):
                    pass
        if total_min >= 60:
            hours = total_min // 60
            mins = total_min % 60
            if mins:
                return f"{hours}h {mins}m"
            return f"{hours}h"
        return f"{total_min}m"


# ---------------------------------------------------------------------------
# Export: Text (terminal output)
# ---------------------------------------------------------------------------

def export_playbooks_text(playbooks: list[Playbook]) -> str:
    """Export playbooks as human-readable terminal text.

    Args:
        playbooks: List of playbooks to export.

    Returns:
        Formatted text string.
    """
    if not playbooks:
        return "No playbooks available."

    lines: list[str] = []
    lines.append(f"\n{'=' * 60}")
    lines.append("  mcpaudit Remediation Playbooks")
    lines.append(f"{'=' * 60}")
    lines.append(f"  Playbooks: {len(playbooks)}")
    lines.append(f"  Total steps: {sum(len(p.steps) for p in playbooks)}")
    lines.append(f"{'=' * 60}")

    for pb in playbooks:
        lines.append(f"\n  [{pb.severity.value}] {pb.finding_rule_id}: {pb.title}")
        if pb.estimated_effort:
            lines.append(f"  Estimated effort: {pb.estimated_effort}")
        lines.append(f"  {'-' * 50}")

        for step in pb.steps:
            lines.append(f"  Step {step.step_number}: {step.action}")
            if step.code_before:
                lines.append(f"    Before: {step.code_before.splitlines()[0]}")
            if step.code_after:
                lines.append(f"    After:  {step.code_after.splitlines()[0]}")
            if step.explanation:
                lines.append(f"    Why: {step.explanation}")
            lines.append("")

        if pb.references:
            lines.append("  References:")
            for ref in pb.references:
                lines.append(f"    - {ref}")

    lines.append(f"\n{'=' * 60}\n")
    return "\n".join(lines)
