"""Rule Authoring SDK for MCP Security Audit.

Simplified API for creating custom security rules with decorators.
Supports AST-based and regex-based rules with full metadata validation.

Usage:
    from mcpaudit.rule_sdk import rule, RuleRegistry, create_rule

    @rule(id="CUSTOM-001", severity="HIGH", cwe="CWE-78",
          description="Detects shell injection via os.system")
    def detect_os_system(code, filename):
        findings = []
        for i, line in enumerate(code.split("\\n"), 1):
            if "os.system(" in line:
                findings.append({"line": i, "snippet": line.strip()})
        return findings

    registry = RuleRegistry()
    registry.register(detect_os_system)
    results = registry.scan("import os\\nos.system('ls')", "test.py")

Zero external dependencies. Part of mcpaudit v2.2.0.
"""

from __future__ import annotations

import ast
import json
import re
from dataclasses import dataclass, field
from typing import Any, Callable


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

VALID_SEVERITIES = frozenset({"CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"})
VALID_ID_PATTERN = re.compile(r"^[A-Z][A-Z0-9]*-\d{3,}$")

VALID_OWASP_CATEGORIES = frozenset({
    "A01", "A02", "A03", "A04", "A05",
    "A06", "A07", "A08", "A09", "A10",
})


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

@dataclass
class RuleFinding:
    """A single finding produced by a custom rule."""

    rule_id: str
    severity: str
    title: str
    description: str
    file: str
    line: int
    code_snippet: str = ""
    cwe: str = ""
    owasp_category: str = ""
    recommendation: str = ""

    def to_dict(self) -> dict[str, Any]:
        result: dict[str, Any] = {
            "rule_id": self.rule_id,
            "severity": self.severity,
            "title": self.title,
            "description": self.description,
            "file": self.file,
            "line": self.line,
        }
        if self.code_snippet:
            result["code_snippet"] = self.code_snippet
        if self.cwe:
            result["cwe"] = self.cwe
        if self.owasp_category:
            result["owasp_category"] = self.owasp_category
        if self.recommendation:
            result["recommendation"] = self.recommendation
        return result


@dataclass
class RuleMetadata:
    """Metadata for a custom rule."""

    id: str
    severity: str = "MEDIUM"
    cwe: str = ""
    description: str = ""
    owasp_category: str = ""
    tags: list[str] = field(default_factory=list)
    enabled: bool = True

    def validate(self) -> list[str]:
        """Validate metadata fields.

        Returns:
            List of error messages (empty if valid).
        """
        errors: list[str] = []
        if not self.id:
            errors.append("Rule id is required")
        elif not VALID_ID_PATTERN.match(self.id):
            errors.append(
                f"Rule id '{self.id}' must match pattern "
                f"'PREFIX-NNN' (e.g. CUSTOM-001)"
            )
        if self.severity not in VALID_SEVERITIES:
            errors.append(
                f"Invalid severity '{self.severity}'. "
                f"Valid: {', '.join(sorted(VALID_SEVERITIES))}"
            )
        if self.cwe and not _validate_cwe(self.cwe):
            errors.append(
                f"Invalid CWE '{self.cwe}'. "
                f"Must be 'CWE-<number>' (e.g. CWE-78)"
            )
        if self.owasp_category and self.owasp_category not in VALID_OWASP_CATEGORIES:
            errors.append(
                f"Invalid OWASP category '{self.owasp_category}'. "
                f"Valid: {', '.join(sorted(VALID_OWASP_CATEGORIES))}"
            )
        return errors

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "severity": self.severity,
            "cwe": self.cwe,
            "description": self.description,
            "owasp_category": self.owasp_category,
            "tags": self.tags,
            "enabled": self.enabled,
        }


@dataclass
class Rule:
    """A custom security rule with check function and metadata."""

    metadata: RuleMetadata
    check_fn: Callable[[str, str], list[dict[str, Any]]]

    @property
    def id(self) -> str:
        return self.metadata.id

    @property
    def severity(self) -> str:
        return self.metadata.severity

    @property
    def enabled(self) -> bool:
        return self.metadata.enabled

    def check(self, code: str, filename: str) -> list[RuleFinding]:
        """Run this rule's check function against code.

        Args:
            code: Source code to scan.
            filename: Name of the file being scanned.

        Returns:
            List of RuleFinding objects.
        """
        if not self.enabled:
            return []

        raw_results = self.check_fn(code, filename)
        findings: list[RuleFinding] = []
        for raw in raw_results:
            if not isinstance(raw, dict):
                continue
            findings.append(RuleFinding(
                rule_id=self.metadata.id,
                severity=self.metadata.severity,
                title=raw.get("title", self.metadata.description),
                description=raw.get(
                    "description", self.metadata.description,
                ),
                file=filename,
                line=raw.get("line", 0),
                code_snippet=raw.get("snippet", ""),
                cwe=self.metadata.cwe,
                owasp_category=self.metadata.owasp_category,
                recommendation=raw.get("recommendation", ""),
            ))
        return findings

    def to_dict(self) -> dict[str, Any]:
        return {
            "metadata": self.metadata.to_dict(),
            "function": self.check_fn.__name__,
        }


# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------

def _validate_cwe(cwe: str) -> bool:
    """Validate CWE format: CWE-<number>."""
    if not cwe.startswith("CWE-"):
        return False
    num_part = cwe[4:]
    return num_part.isdigit() and int(num_part) > 0


def validate_rule_metadata(metadata: RuleMetadata) -> list[str]:
    """Validate rule metadata.

    Args:
        metadata: RuleMetadata to validate.

    Returns:
        List of error messages (empty if valid).
    """
    return metadata.validate()


# ---------------------------------------------------------------------------
# Decorator
# ---------------------------------------------------------------------------

def rule(
    id: str,
    severity: str = "MEDIUM",
    cwe: str = "",
    description: str = "",
    owasp_category: str = "",
    tags: list[str] | None = None,
    enabled: bool = True,
) -> Callable:
    """Decorator to create a custom rule from a check function.

    The decorated function must accept (code: str, filename: str) and
    return a list of dicts with at least a 'line' key.

    Args:
        id: Rule identifier (e.g. "CUSTOM-001").
        severity: Severity level (CRITICAL/HIGH/MEDIUM/LOW/INFO).
        cwe: CWE reference (e.g. "CWE-78").
        description: Human-readable description.
        owasp_category: OWASP top-10 category (A01-A10).
        tags: Optional tags for filtering.
        enabled: Whether the rule is active.

    Returns:
        Decorated function with .rule attribute.
    """
    metadata = RuleMetadata(
        id=id,
        severity=severity,
        cwe=cwe,
        description=description,
        owasp_category=owasp_category,
        tags=tags or [],
        enabled=enabled,
    )

    def decorator(fn: Callable) -> Callable:
        rule_obj = Rule(metadata=metadata, check_fn=fn)
        fn.rule = rule_obj  # type: ignore[attr-defined]
        fn.metadata = metadata  # type: ignore[attr-defined]
        return fn

    return decorator


# ---------------------------------------------------------------------------
# Factory function
# ---------------------------------------------------------------------------

def create_rule(
    id: str,
    severity: str,
    check_fn: Callable[[str, str], list[dict[str, Any]]],
    cwe: str = "",
    description: str = "",
    owasp_category: str = "",
    tags: list[str] | None = None,
    enabled: bool = True,
) -> Rule:
    """Create a Rule object programmatically.

    Args:
        id: Rule identifier (e.g. "CUSTOM-001").
        severity: Severity level.
        check_fn: Function(code, filename) -> list[dict].
        cwe: CWE reference.
        description: Human-readable description.
        owasp_category: OWASP category.
        tags: Optional tags.
        enabled: Whether the rule is active.

    Returns:
        Rule object.

    Raises:
        ValueError: If metadata is invalid.
    """
    metadata = RuleMetadata(
        id=id,
        severity=severity,
        cwe=cwe,
        description=description,
        owasp_category=owasp_category,
        tags=tags or [],
        enabled=enabled,
    )
    errors = metadata.validate()
    if errors:
        raise ValueError(
            f"Invalid rule metadata: {'; '.join(errors)}"
        )
    return Rule(metadata=metadata, check_fn=check_fn)


# ---------------------------------------------------------------------------
# Rule builders (AST-based and regex-based)
# ---------------------------------------------------------------------------

def create_regex_rule(
    id: str,
    severity: str,
    pattern: str,
    description: str = "",
    cwe: str = "",
    owasp_category: str = "",
    recommendation: str = "",
) -> Rule:
    """Create a regex-based rule.

    Args:
        id: Rule identifier.
        severity: Severity level.
        pattern: Regex pattern to match.
        description: Description.
        cwe: CWE reference.
        owasp_category: OWASP category.
        recommendation: Fix recommendation.

    Returns:
        Rule that scans code line-by-line with the regex.

    Raises:
        ValueError: If pattern is invalid or metadata fails.
    """
    try:
        compiled = re.compile(pattern)
    except re.error as exc:
        raise ValueError(f"Invalid regex pattern: {exc}")

    def check_fn(code: str, filename: str) -> list[dict[str, Any]]:
        results: list[dict[str, Any]] = []
        for i, line in enumerate(code.split("\n"), 1):
            if compiled.search(line):
                results.append({
                    "line": i,
                    "snippet": line.strip(),
                    "recommendation": recommendation,
                })
        return results

    check_fn.__name__ = f"regex_check_{id}"
    return create_rule(
        id=id,
        severity=severity,
        check_fn=check_fn,
        cwe=cwe,
        description=description,
        owasp_category=owasp_category,
    )


def create_ast_rule(
    id: str,
    severity: str,
    node_type: str,
    attr_check: Callable[[ast.AST], bool] | None = None,
    description: str = "",
    cwe: str = "",
    owasp_category: str = "",
    recommendation: str = "",
) -> Rule:
    """Create an AST-based rule for Python code.

    Args:
        id: Rule identifier.
        severity: Severity level.
        node_type: AST node type name (e.g. "Call", "Import").
        attr_check: Optional function to further filter nodes.
        description: Description.
        cwe: CWE reference.
        owasp_category: OWASP category.
        recommendation: Fix recommendation.

    Returns:
        Rule that parses Python AST and finds matching nodes.

    Raises:
        ValueError: If metadata fails validation.
    """

    def check_fn(code: str, filename: str) -> list[dict[str, Any]]:
        results: list[dict[str, Any]] = []
        try:
            tree = ast.parse(code, filename=filename)
        except SyntaxError:
            return results

        target_cls = getattr(ast, node_type, None)
        if target_cls is None:
            return results

        for node in ast.walk(tree):
            if not isinstance(node, target_cls):
                continue
            if attr_check is not None and not attr_check(node):
                continue
            lineno = getattr(node, "lineno", 0)
            snippet = _extract_snippet(code, lineno)
            results.append({
                "line": lineno,
                "snippet": snippet,
                "recommendation": recommendation,
            })
        return results

    check_fn.__name__ = f"ast_check_{id}"
    return create_rule(
        id=id,
        severity=severity,
        check_fn=check_fn,
        cwe=cwe,
        description=description,
        owasp_category=owasp_category,
    )


def _extract_snippet(code: str, lineno: int) -> str:
    """Extract a code snippet for a given line number."""
    lines = code.split("\n")
    if 1 <= lineno <= len(lines):
        return lines[lineno - 1].strip()
    return ""


# ---------------------------------------------------------------------------
# Rule Registry
# ---------------------------------------------------------------------------

class RuleRegistry:
    """Registry for managing and running custom rules."""

    def __init__(self) -> None:
        self._rules: dict[str, Rule] = {}

    @property
    def count(self) -> int:
        """Number of registered rules."""
        return len(self._rules)

    def register(self, rule_or_fn: Rule | Callable) -> list[str]:
        """Register a rule or decorated function.

        Args:
            rule_or_fn: A Rule object or a @rule-decorated function.

        Returns:
            List of validation errors (empty if ok).
        """
        if isinstance(rule_or_fn, Rule):
            rule_obj = rule_or_fn
        elif hasattr(rule_or_fn, "rule"):
            rule_obj = rule_or_fn.rule  # type: ignore
        else:
            return ["Not a valid rule or decorated function"]

        errors = rule_obj.metadata.validate()
        if errors:
            return errors

        if rule_obj.id in self._rules:
            return [f"Rule '{rule_obj.id}' already registered"]

        self._rules[rule_obj.id] = rule_obj
        return []

    def unregister(self, rule_id: str) -> bool:
        """Remove a rule by ID.

        Args:
            rule_id: Rule identifier to remove.

        Returns:
            True if the rule was found and removed.
        """
        if rule_id in self._rules:
            del self._rules[rule_id]
            return True
        return False

    def get(self, rule_id: str) -> Rule | None:
        """Get a rule by ID."""
        return self._rules.get(rule_id)

    def list_rules(self) -> list[Rule]:
        """List all registered rules."""
        return list(self._rules.values())

    def list_ids(self) -> list[str]:
        """List all registered rule IDs (sorted)."""
        return sorted(self._rules.keys())

    def filter_by_severity(self, severity: str) -> list[Rule]:
        """Get rules matching a severity level.

        Args:
            severity: Severity string to match.

        Returns:
            List of matching rules.
        """
        return [
            r for r in self._rules.values()
            if r.severity == severity
        ]

    def filter_by_tag(self, tag: str) -> list[Rule]:
        """Get rules that have a specific tag.

        Args:
            tag: Tag to match.

        Returns:
            List of matching rules.
        """
        return [
            r for r in self._rules.values()
            if tag in r.metadata.tags
        ]

    def scan(
        self, code: str, filename: str,
    ) -> list[RuleFinding]:
        """Run all registered rules against code.

        Args:
            code: Source code to scan.
            filename: Filename for context.

        Returns:
            Combined list of all findings, sorted by line number.
        """
        all_findings: list[RuleFinding] = []
        for rule_obj in self._rules.values():
            all_findings.extend(rule_obj.check(code, filename))
        all_findings.sort(key=lambda f: (f.file, f.line))
        return all_findings

    def scan_file(self, path: str) -> list[RuleFinding]:
        """Scan a file on disk with all registered rules.

        Args:
            path: File path to read and scan.

        Returns:
            List of findings.
        """
        from pathlib import Path as P
        filepath = P(path)
        if not filepath.exists():
            return []
        try:
            code = filepath.read_text(encoding="utf-8")
        except (OSError, UnicodeDecodeError):
            return []
        return self.scan(code, str(filepath))

    def clear(self) -> None:
        """Remove all registered rules."""
        self._rules.clear()

    def to_dict(self) -> dict[str, Any]:
        """Serialize registry to dict."""
        return {
            "rules": [r.to_dict() for r in self._rules.values()],
            "count": self.count,
        }


# ---------------------------------------------------------------------------
# Formatting
# ---------------------------------------------------------------------------

def format_findings_text(findings: list[RuleFinding]) -> str:
    """Format rule findings as human-readable text.

    Args:
        findings: List of findings.

    Returns:
        Formatted text.
    """
    if not findings:
        return "No findings.\n"

    lines: list[str] = [
        f"Custom Rule Findings ({len(findings)})",
        "=" * 50,
    ]
    for f in findings:
        lines.append(
            f"  [{f.severity}] {f.rule_id}: "
            f"{f.title}"
        )
        lines.append(f"    File: {f.file}:{f.line}")
        if f.code_snippet:
            lines.append(f"    Code: {f.code_snippet}")
        if f.cwe:
            lines.append(f"    CWE: {f.cwe}")
    return "\n".join(lines) + "\n"


def format_findings_json(findings: list[RuleFinding]) -> str:
    """Format rule findings as JSON.

    Args:
        findings: List of findings.

    Returns:
        JSON string.
    """
    data = {
        "findings": [f.to_dict() for f in findings],
        "total": len(findings),
    }
    return json.dumps(data, indent=2, ensure_ascii=False)


def format_registry_text(registry: RuleRegistry) -> str:
    """Format registry as text summary.

    Args:
        registry: RuleRegistry to describe.

    Returns:
        Text summary.
    """
    rules = registry.list_rules()
    if not rules:
        return "No custom rules registered.\n"

    lines: list[str] = [
        f"Custom Rules ({registry.count})",
        "=" * 50,
    ]
    for r in rules:
        status = "enabled" if r.enabled else "disabled"
        lines.append(
            f"  {r.id:<16} [{r.severity:<8}] "
            f"{r.metadata.description} ({status})"
        )
        if r.metadata.cwe:
            lines.append(f"    CWE: {r.metadata.cwe}")
    return "\n".join(lines) + "\n"


# ---------------------------------------------------------------------------
# Example rules using the SDK
# ---------------------------------------------------------------------------

@rule(
    id="SDK-001",
    severity="HIGH",
    cwe="CWE-78",
    description="Detects os.system() calls (shell injection risk)",
    owasp_category="A03",
    tags=["injection", "shell"],
)
def detect_os_system(code: str, filename: str) -> list[dict[str, Any]]:
    """Detect os.system() calls."""
    results: list[dict[str, Any]] = []
    for i, line in enumerate(code.split("\n"), 1):
        if "os.system(" in line:
            results.append({
                "line": i,
                "snippet": line.strip(),
                "title": "Shell command via os.system()",
                "recommendation": "Use subprocess.run() with shell=False",
            })
    return results


@rule(
    id="SDK-002",
    severity="CRITICAL",
    cwe="CWE-502",
    description="Detects unsafe deserialization calls (security risk)",
    owasp_category="A08",
    tags=["deserialization"],
)
def detect_unsafe_deserialize(
    code: str, filename: str,
) -> list[dict[str, Any]]:
    """Detect unsafe deserialization usage."""
    # Detect calls to modules known for unsafe deserialization
    _UNSAFE_PATTERNS = [
        "pickle.loads(",
        "pickle.load(",
        "yaml.load(",
        "marshal.loads(",
    ]
    results: list[dict[str, Any]] = []
    for i, line in enumerate(code.split("\n"), 1):
        for pat in _UNSAFE_PATTERNS:
            if pat in line:
                results.append({
                    "line": i,
                    "snippet": line.strip(),
                    "title": "Unsafe deserialization detected",
                    "recommendation": "Use safe alternatives (json, yaml.safe_load)",
                })
                break
    return results


@rule(
    id="SDK-003",
    severity="MEDIUM",
    cwe="CWE-327",
    description="Detects use of MD5/SHA1 for security purposes",
    tags=["crypto", "weak-hash"],
)
def detect_weak_hash(code: str, filename: str) -> list[dict[str, Any]]:
    """Detect weak hash algorithms (MD5, SHA1)."""
    weak_re = re.compile(r"(?:hashlib\.(?:md5|sha1)|MD5|SHA1)\b")
    results: list[dict[str, Any]] = []
    for i, line in enumerate(code.split("\n"), 1):
        if weak_re.search(line):
            results.append({
                "line": i,
                "snippet": line.strip(),
                "title": "Weak hash algorithm detected",
                "recommendation": "Use SHA-256 or stronger",
            })
    return results


# Pre-built example rules list
EXAMPLE_RULES: list[Callable] = [
    detect_os_system,
    detect_unsafe_deserialize,
    detect_weak_hash,
]


def get_example_registry() -> RuleRegistry:
    """Create a registry pre-loaded with example rules.

    Returns:
        RuleRegistry with 3 example rules.
    """
    reg = RuleRegistry()
    for fn in EXAMPLE_RULES:
        reg.register(fn)
    return reg
