"""Rule Generation Engine — Create custom security rules from templates.

Allows users to generate new security rules from predefined templates
(command_injection, path_traversal, auth_check, input_validation,
secret_detection) or from existing scan findings.

Generated rules are valid Python plugin files loadable by the plugin system.
"""

from __future__ import annotations

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


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------

@dataclass
class RuleTemplate:
    """Definition of a rule template."""

    name: str
    description: str
    pattern: str
    severity: str
    cwe: str
    category: str
    language: str = "python"


# ---------------------------------------------------------------------------
# Predefined templates
# ---------------------------------------------------------------------------

BUILTIN_TEMPLATES: dict[str, RuleTemplate] = {
    "command_injection": RuleTemplate(
        name="command_injection",
        description="Detect command injection via subprocess with shell=True",
        pattern=r"subprocess\.\w+\(.*shell\s*=\s*True",
        severity="CRITICAL",
        cwe="CWE-78",
        category="injection",
        language="python",
    ),
    "path_traversal": RuleTemplate(
        name="path_traversal",
        description="Detect path traversal via unvalidated path joins",
        pattern=r"os\.path\.join\(.*\+|Path\(.*\+",
        severity="HIGH",
        cwe="CWE-22",
        category="traversal",
        language="python",
    ),
    "auth_check": RuleTemplate(
        name="auth_check",
        description="Detect missing authentication checks in tool handlers",
        pattern=r"@(mcp\.tool|app\.tool|server\.tool)",
        severity="HIGH",
        cwe="CWE-306",
        category="authentication",
        language="python",
    ),
    "input_validation": RuleTemplate(
        name="input_validation",
        description="Detect tool handlers without input validation",
        pattern=r"def\s+\w+\(.*\)\s*:",
        severity="MEDIUM",
        cwe="CWE-20",
        category="validation",
        language="python",
    ),
    "secret_detection": RuleTemplate(
        name="secret_detection",
        description="Detect hardcoded secrets and API keys",
        pattern=(
            r"(?:api_key|secret|password|token|auth)\s*="
            r'\s*["\'][^"\']{8,}["\']'
        ),
        severity="HIGH",
        cwe="CWE-798",
        category="secrets",
        language="python",
    ),
}


def list_templates() -> list[str]:
    """Return names of all available built-in templates."""
    return sorted(BUILTIN_TEMPLATES.keys())


def get_template(name: str) -> RuleTemplate | None:
    """Get a built-in template by name.

    Args:
        name: Template name (e.g. 'command_injection').

    Returns:
        RuleTemplate if found, None otherwise.
    """
    return BUILTIN_TEMPLATES.get(name)


# ---------------------------------------------------------------------------
# Rule code generation
# ---------------------------------------------------------------------------

def _sanitize_class_name(rule_id: str) -> str:
    """Convert a rule ID like MY-001 to a valid Python class name."""
    clean = re.sub(r"[^a-zA-Z0-9]", "_", rule_id)
    parts = clean.split("_")
    return "".join(p.capitalize() for p in parts if p) + "Rule"


def _sanitize_rule_id(rule_id: str) -> str:
    """Normalize rule ID to uppercase with hyphens."""
    return rule_id.upper().strip()


def _escape_python_string(s: str) -> str:
    """Escape a string for safe inclusion in Python source."""
    return s.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")


def _generate_rule_code(
    rule_id: str,
    name: str,
    description: str,
    severity: str,
    cwe: str,
    category: str,
    pattern: str,
) -> str:
    """Generate valid Python source for a plugin rule.

    Args:
        rule_id: Unique rule identifier (e.g. 'MY-001').
        name: Human-readable rule name.
        description: What the rule detects.
        severity: Severity level (CRITICAL/HIGH/MEDIUM/LOW/INFO).
        cwe: CWE reference (e.g. 'CWE-78').
        category: Rule category (e.g. 'injection').
        pattern: Regex pattern to detect.

    Returns:
        Complete Python source code as a string.
    """
    class_name = _sanitize_class_name(rule_id)
    safe_id = _sanitize_rule_id(rule_id)
    safe_desc = _escape_python_string(description)
    safe_name = _escape_python_string(name)
    safe_pattern = _escape_python_string(pattern)
    severity_upper = severity.upper()

    return textwrap.dedent(f'''\
        """Auto-generated rule: {safe_id} — {safe_name}.

        Generated by mcpaudit rule generation engine.
        Category: {category}
        CWE: {cwe}
        """

        from __future__ import annotations

        import ast
        import re

        from mcpaudit.models import Finding, Severity
        from mcpaudit.rules import Rule


        _PATTERN = re.compile(r"{safe_pattern}")


        class {class_name}(Rule):
            """Detects: {safe_desc}"""

            @property
            def rule_id(self) -> str:
                return "{safe_id}"

            @property
            def name(self) -> str:
                return "{safe_name}"

            @property
            def severity(self) -> Severity:
                return Severity.{severity_upper}

            @property
            def description(self) -> str:
                return "{safe_desc}"

            @property
            def cwe(self) -> str:
                return "{cwe}"

            def check(self, tree: ast.Module, source: str) -> list[Finding]:
                findings: list[Finding] = []
                for i, line in enumerate(source.splitlines(), start=1):
                    if _PATTERN.search(line):
                        findings.append(Finding(
                            rule_id=self.rule_id,
                            severity=self.severity,
                            title=self.name,
                            description=self.description,
                            file="",
                            line=i,
                            code_snippet=line.strip(),
                            cwe=self.cwe,
                        ))
                return findings


        def register_rules() -> list[Rule]:
            """Plugin entry point — return rule instances."""
            return [{class_name}()]
    ''')


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

@dataclass
class GeneratedRule:
    """Result of rule generation."""

    rule_id: str
    name: str
    description: str
    severity: str
    cwe: str
    category: str
    pattern: str
    source_code: str
    template_used: str = ""

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "rule_id": self.rule_id,
            "name": self.name,
            "description": self.description,
            "severity": self.severity,
            "cwe": self.cwe,
            "category": self.category,
            "pattern": self.pattern,
            "template_used": self.template_used,
        }


def generate_rule_from_template(
    template_name: str,
    params: dict[str, str],
) -> GeneratedRule:
    """Generate a rule from a built-in template.

    Args:
        template_name: Name of the template (e.g. 'command_injection').
        params: Override parameters. Supported keys:
            - rule_id: Custom rule ID (required).
            - name: Custom rule name.
            - description: Custom description.
            - severity: Override severity.
            - cwe: Override CWE reference.
            - pattern: Override detection pattern.

    Returns:
        GeneratedRule with source code ready to write.

    Raises:
        ValueError: If template_name is unknown or rule_id missing.
    """
    template = get_template(template_name)
    if template is None:
        available = ", ".join(list_templates())
        raise ValueError(
            f"Unknown template '{template_name}'. "
            f"Available: {available}"
        )

    rule_id = params.get("rule_id", "").strip()
    if not rule_id:
        raise ValueError("rule_id is required in params")

    name = params.get("name", template.description)
    description = params.get("description", template.description)
    severity = params.get("severity", template.severity).upper()
    cwe = params.get("cwe", template.cwe)
    pattern = params.get("pattern", template.pattern)
    category = params.get("category", template.category)

    _validate_severity(severity)

    source_code = _generate_rule_code(
        rule_id=rule_id,
        name=name,
        description=description,
        severity=severity,
        cwe=cwe,
        category=category,
        pattern=pattern,
    )

    return GeneratedRule(
        rule_id=_sanitize_rule_id(rule_id),
        name=name,
        description=description,
        severity=severity,
        cwe=cwe,
        category=category,
        pattern=pattern,
        source_code=source_code,
        template_used=template_name,
    )


def generate_rule_from_finding(finding: dict[str, Any]) -> GeneratedRule:
    """Convert a finding (dict or JSON) into a reusable rule.

    Args:
        finding: Finding dict with keys: rule_id, severity, title,
                 description, code_snippet, cwe.

    Returns:
        GeneratedRule with source code ready to write.

    Raises:
        ValueError: If required fields are missing.
    """
    rule_id = finding.get("rule_id", "")
    if not rule_id:
        raise ValueError("Finding must have a 'rule_id' field")

    severity = finding.get("severity", "MEDIUM").upper()
    _validate_severity(severity)

    title = finding.get("title", finding.get("description", "Custom rule"))
    description = finding.get("description", title)
    cwe = finding.get("cwe", "")
    code_snippet = finding.get("code_snippet", "")

    # Build a pattern from the code snippet if available
    pattern = _pattern_from_snippet(code_snippet)
    category = _guess_category(rule_id, description)

    source_code = _generate_rule_code(
        rule_id=rule_id,
        name=title,
        description=description,
        severity=severity,
        cwe=cwe,
        category=category,
        pattern=pattern,
    )

    return GeneratedRule(
        rule_id=_sanitize_rule_id(rule_id),
        name=title,
        description=description,
        severity=severity,
        cwe=cwe,
        category=category,
        pattern=pattern,
        source_code=source_code,
        template_used="from_finding",
    )


def save_generated_rule(
    generated: GeneratedRule,
    output_dir: Path,
) -> Path:
    """Write a generated rule to disk as a .py file.

    Args:
        generated: The generated rule to save.
        output_dir: Directory to write the file into.

    Returns:
        Path to the written file.
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    filename = generated.rule_id.lower().replace("-", "_") + ".py"
    filepath = output_dir / filename
    filepath.write_text(generated.source_code, encoding="utf-8")
    return filepath


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_VALID_SEVERITIES = {"CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"}


def _validate_severity(severity: str) -> None:
    """Validate severity string."""
    if severity not in _VALID_SEVERITIES:
        raise ValueError(
            f"Invalid severity '{severity}'. "
            f"Valid: {', '.join(sorted(_VALID_SEVERITIES))}"
        )


def _pattern_from_snippet(snippet: str) -> str:
    """Derive a regex pattern from a code snippet.

    Escapes special regex chars and creates a simple matching pattern.
    """
    if not snippet.strip():
        return r".*"
    # Escape regex metacharacters, keep word boundaries useful
    escaped = re.escape(snippet.strip())
    # Replace escaped whitespace with flexible whitespace
    escaped = re.sub(r"\\ ", r"\\s+", escaped)
    return escaped


def _guess_category(rule_id: str, description: str) -> str:
    """Guess a category from rule ID prefix or description."""
    prefix = rule_id.split("-")[0].upper() if "-" in rule_id else ""
    prefix_map = {
        "CMD": "injection",
        "PATH": "traversal",
        "AUTH": "authentication",
        "SQL": "injection",
        "INP": "validation",
        "SEC": "secrets",
        "SSRF": "network",
        "XSS": "output",
    }
    if prefix in prefix_map:
        return prefix_map[prefix]

    desc_lower = description.lower()
    for keyword, cat in [
        ("inject", "injection"),
        ("travers", "traversal"),
        ("auth", "authentication"),
        ("valid", "validation"),
        ("secret", "secrets"),
        ("password", "secrets"),
    ]:
        if keyword in desc_lower:
            return cat
    return "custom"
