"""Policy engine — organizational policy enforcement for MCP audits.

Allows organizations to define security policies in YAML that
are evaluated against scan results. Policies can block, warn,
or ignore specific findings, set minimum trust scores,
maximum severity thresholds, and required security grades.

Example policy.yaml:
    policies:
      - name: "no-command-injection"
        rule: CMD-001
        severity: critical
        action: block
      - name: "require-auth"
        rule: AUTH-001
        action: warn
    min_trust_score: 70
    max_severity: medium
    required_grades: ["A", "B", "C"]
"""

from __future__ import annotations

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

from .models import Finding, ScanSummary, Severity


# ---------------------------------------------------------------------------
# Policy models
# ---------------------------------------------------------------------------

class PolicyAction(Enum):
    """Action to take when a policy rule matches."""

    BLOCK = "block"
    WARN = "warn"
    IGNORE = "ignore"


@dataclass
class PolicyRule:
    """A single policy rule."""

    name: str
    rule_id: str       # e.g. "CMD-001" or "*" for wildcard
    action: PolicyAction
    severity: str = ""  # optional severity filter


@dataclass
class Policy:
    """Organizational security policy."""

    rules: list[PolicyRule] = field(default_factory=list)
    min_trust_score: int = 0
    max_severity: Severity | None = None
    required_grades: list[str] = field(default_factory=list)


@dataclass
class PolicyEvaluation:
    """Result of evaluating a single policy rule."""

    rule_name: str
    rule_id: str
    action: PolicyAction
    matched_findings: list[Finding] = field(default_factory=list)
    reason: str = ""


@dataclass
class PolicyResult:
    """Result of evaluating a full policy against scan results."""

    passed: bool
    evaluations: list[PolicyEvaluation] = field(default_factory=list)
    reasons: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "passed": self.passed,
            "evaluations": [
                {
                    "rule_name": e.rule_name,
                    "rule_id": e.rule_id,
                    "action": e.action.value,
                    "matched_count": len(e.matched_findings),
                    "reason": e.reason,
                }
                for e in self.evaluations
            ],
            "reasons": self.reasons,
            "warnings": self.warnings,
        }


# ---------------------------------------------------------------------------
# YAML parsing (stdlib only — no PyYAML dependency)
# ---------------------------------------------------------------------------

def _parse_yaml_simple(text: str) -> dict[str, Any]:
    """Parse a simple YAML subset using stdlib only.

    Supports:
    - Top-level keys with scalar values
    - Lists of dicts (policies)
    - Simple lists of strings

    This avoids adding PyYAML as a dependency.

    Args:
        text: YAML content string.

    Returns:
        Parsed dict.

    Raises:
        ValueError: If the YAML cannot be parsed.
    """
    # Reject obviously malformed YAML (unbalanced brackets, etc.)
    bracket_count = text.count("[") - text.count("]")
    if bracket_count != 0:
        raise ValueError("Unbalanced brackets in YAML")

    result: dict[str, Any] = {}
    lines = text.splitlines()
    i = 0
    current_key = ""
    current_list: list[Any] = []

    while i < len(lines):
        line = lines[i]
        stripped = line.strip()

        # Skip empty lines and comments
        if not stripped or stripped.startswith("#"):
            i += 1
            continue

        # Top-level key
        if not line[0].isspace() and ":" in stripped:
            # Save previous list if any
            if current_key and current_list:
                result[current_key] = current_list
                current_list = []

            key, _, value = stripped.partition(":")
            key = key.strip()
            value = value.strip()
            current_key = key

            if value and not value.startswith("["):
                # Scalar value
                result[key] = _parse_scalar(value)
                current_key = ""
            elif value.startswith("["):
                # Inline list: [A, B, C]
                inner = value.strip("[]")
                result[key] = [
                    _parse_scalar(v.strip().strip('"').strip("'"))
                    for v in inner.split(",")
                    if v.strip()
                ]
                current_key = ""
            i += 1
            continue

        # List item (starts with -)
        if stripped.startswith("- ") and current_key:
            # Check if it's a dict item or scalar
            rest = stripped[2:].strip()
            if ":" in rest and not rest.startswith('"'):
                # Dict item — collect all key: value pairs at this indent
                item: dict[str, Any] = {}
                first_key, _, first_val = rest.partition(":")
                item[first_key.strip()] = _parse_scalar(first_val.strip())
                i += 1
                while i < len(lines):
                    next_line = lines[i]
                    next_stripped = next_line.strip()
                    if not next_stripped or next_stripped.startswith("#"):
                        i += 1
                        continue
                    indent = len(next_line) - len(next_line.lstrip())
                    if indent <= 2 or next_stripped.startswith("- "):
                        break
                    if ":" in next_stripped:
                        k, _, v = next_stripped.partition(":")
                        item[k.strip()] = _parse_scalar(v.strip())
                    i += 1
                current_list.append(item)
            else:
                # Simple scalar list item
                current_list.append(_parse_scalar(rest.strip('"').strip("'")))
                i += 1
            continue

        i += 1

    # Save last list
    if current_key and current_list:
        result[current_key] = current_list

    return result


def _parse_scalar(value: str) -> Any:
    """Parse a YAML scalar value."""
    if not value:
        return ""
    # Remove quotes
    if (value.startswith('"') and value.endswith('"')) or \
       (value.startswith("'") and value.endswith("'")):
        return value[1:-1]
    # Boolean
    if value.lower() in ("true", "yes"):
        return True
    if value.lower() in ("false", "no"):
        return False
    # Integer
    try:
        return int(value)
    except ValueError:
        pass
    return value


# ---------------------------------------------------------------------------
# Loading
# ---------------------------------------------------------------------------

def load_policy(path: Path) -> Policy:
    """Load a policy from a YAML file.

    Args:
        path: Path to the policy YAML file.

    Returns:
        Parsed Policy object.

    Raises:
        FileNotFoundError: If the file doesn't exist.
        ValueError: If the YAML is invalid or malformed.
    """
    if not path.exists():
        raise FileNotFoundError(f"Policy file not found: {path}")

    content = path.read_text(encoding="utf-8")

    try:
        data = _parse_yaml_simple(content)
    except Exception as e:
        raise ValueError(f"Invalid policy YAML: {e}") from e

    if not isinstance(data, dict):
        raise ValueError("Policy file must contain a YAML mapping")

    rules: list[PolicyRule] = []
    for item in data.get("policies", []):
        if not isinstance(item, dict):
            continue
        action_str = str(item.get("action", "warn")).lower()
        try:
            action = PolicyAction(action_str)
        except ValueError:
            action = PolicyAction.WARN

        rules.append(PolicyRule(
            name=str(item.get("name", "")),
            rule_id=str(item.get("rule", "")),
            action=action,
            severity=str(item.get("severity", "")),
        ))

    max_severity = None
    max_sev_str = data.get("max_severity", "")
    if max_sev_str:
        try:
            max_severity = Severity.from_string(str(max_sev_str))
        except ValueError:
            pass

    required_grades = data.get("required_grades", [])
    if isinstance(required_grades, str):
        required_grades = [required_grades]

    return Policy(
        rules=rules,
        min_trust_score=int(data.get("min_trust_score", 0)),
        max_severity=max_severity,
        required_grades=[str(g) for g in required_grades],
    )


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

def _collect_findings(summary: ScanSummary) -> list[Finding]:
    """Collect all findings from a scan summary."""
    findings: list[Finding] = []
    for result in summary.results:
        findings.extend(result.findings)
    return findings


def _matches_rule(finding: Finding, policy_rule: PolicyRule) -> bool:
    """Check if a finding matches a policy rule.

    Args:
        finding: The scan finding.
        policy_rule: The policy rule to match against.

    Returns:
        True if the finding matches.
    """
    # Wildcard matches everything
    if policy_rule.rule_id == "*":
        return True
    return finding.rule_id == policy_rule.rule_id


def evaluate_policy(
    policy: Policy,
    summary: ScanSummary,
    trust_score: int = 100,
) -> PolicyResult:
    """Evaluate a policy against scan results.

    Args:
        policy: The organizational policy.
        summary: Scan results to evaluate.
        trust_score: Trust score from the scan (0-100).

    Returns:
        PolicyResult with pass/fail verdict and reasons.
    """
    passed = True
    evaluations: list[PolicyEvaluation] = []
    reasons: list[str] = []
    warnings: list[str] = []

    all_findings = _collect_findings(summary)

    # Evaluate each policy rule
    for rule in policy.rules:
        matched = [f for f in all_findings if _matches_rule(f, rule)]
        if not matched:
            continue

        evaluation = PolicyEvaluation(
            rule_name=rule.name,
            rule_id=rule.rule_id,
            action=rule.action,
            matched_findings=matched,
        )

        if rule.action == PolicyAction.BLOCK:
            evaluation.reason = (
                f"Policy '{rule.name}' blocks {rule.rule_id}: "
                f"{len(matched)} finding(s) matched"
            )
            reasons.append(evaluation.reason)
            passed = False
        elif rule.action == PolicyAction.WARN:
            evaluation.reason = (
                f"Policy '{rule.name}' warns on {rule.rule_id}: "
                f"{len(matched)} finding(s) matched"
            )
            warnings.append(evaluation.reason)
        # IGNORE: do nothing

        evaluations.append(evaluation)

    # Trust score gate
    if policy.min_trust_score > 0 and trust_score < policy.min_trust_score:
        reason = (
            f"Trust score {trust_score} below minimum {policy.min_trust_score}"
        )
        reasons.append(reason)
        passed = False

    # Max severity gate
    if policy.max_severity is not None:
        severity_order = list(Severity)
        max_idx = severity_order.index(policy.max_severity)
        for finding in all_findings:
            finding_idx = severity_order.index(finding.severity)
            if finding_idx < max_idx:
                reason = (
                    f"Finding {finding.rule_id} has severity "
                    f"{finding.severity.value} above max {policy.max_severity.value}"
                )
                reasons.append(reason)
                passed = False
                break

    # Grade gate
    if policy.required_grades:
        grade = summary.overall_grade
        if grade not in policy.required_grades:
            reason = (
                f"Grade {grade} not in required grades: "
                f"{', '.join(policy.required_grades)}"
            )
            reasons.append(reason)
            passed = False

    return PolicyResult(
        passed=passed,
        evaluations=evaluations,
        reasons=reasons,
        warnings=warnings,
    )


# ---------------------------------------------------------------------------
# Formatters
# ---------------------------------------------------------------------------

def format_policy_text(result: PolicyResult) -> str:
    """Format a policy result as human-readable text.

    Args:
        result: PolicyResult to format.

    Returns:
        Formatted text string.
    """
    lines: list[str] = []
    lines.append(f"\n{'=' * 60}")
    lines.append("  mcpaudit — Policy Evaluation")
    lines.append(f"{'=' * 60}")

    verdict = "PASS" if result.passed else "FAIL"
    lines.append(f"  Verdict: {verdict}")

    if result.reasons:
        lines.append(f"\n  Blocked ({len(result.reasons)}):")
        for r in result.reasons:
            lines.append(f"    [BLOCK] {r}")

    if result.warnings:
        lines.append(f"\n  Warnings ({len(result.warnings)}):")
        for w in result.warnings:
            lines.append(f"    [WARN] {w}")

    if not result.reasons and not result.warnings:
        lines.append("  All policy checks passed.")

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


def format_policy_json(result: PolicyResult) -> str:
    """Format a policy result as JSON.

    Args:
        result: PolicyResult to format.

    Returns:
        JSON string.
    """
    return json.dumps(result.to_dict(), indent=2, ensure_ascii=False)
