"""OPA/Rego policy generator (Team+ tier).

Generates Rego policies from compliance findings that evaluate Terraform plan JSON.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING

from complyform import __version__
from complyform.core.generators.base import GenerationResult, PolicyGenerator

if TYPE_CHECKING:
    from pathlib import Path

_REGO_TEMPLATE = """\
# Auto-generated by ComplyForm v{version}
# Framework: {framework} | Control: {control_id} | Resource: {resource_type}
# Generated: {timestamp}
# Usage: conftest test tfplan.json --policy {policy_dir}

package complyform.{cloud}.{resource_type}

import rego.v1

{rules}
"""

_DENY_RULE_TEMPLATE = """\
# Control: {control_id} ({framework})
# Checkov: {checkov_id}
# Native IDs: {native_ids}
# Severity: {severity}
deny contains msg if {{
    some rc in input.resource_changes
    rc.type == "{tf_resource_type}"
    rc.change.after != null
    not _check_{rule_index}(rc.change.after)
    msg := "{control_id}: {description}"
}}

_check_{rule_index}(after) if {{
    after.{check_field} != null
    after.{check_field} != ""
    after.{check_field} != false
}}
"""


@dataclass
class RegoRule:
    """Single deny rule in a Rego file."""

    control_id: str
    framework: str
    checkov_id: str | None
    native_ids: dict[str, object] | None
    severity: str
    description: str
    resource_type: str
    check_field: str


@dataclass
class RegoFile:
    """Collection of rules for a single resource type."""

    cloud: str
    resource_type: str
    rules: list[RegoRule] = field(default_factory=list)


class OpaRegoGenerator(PolicyGenerator):
    """Generate OPA/Rego policies from compliance findings."""

    FORMAT_ID = "opa"
    TIER_FEATURE = "opa_rego"

    _SUPPORTED_TIERS = frozenset(
        {"team", "agency", "enterprise", "enterprise_unlimited"}
    )

    def supports_tier(self, tier: str) -> bool:
        return tier in self._SUPPORTED_TIERS

    def format_name(self) -> str:
        return "OPA/Rego"

    def generate(
        self,
        findings: list[dict[str, object]],
        mappings: dict[str, object],
        cloud: str,
        tier: str,
        output_dir: Path,
    ) -> GenerationResult:
        if not findings:
            return GenerationResult(
                format=self.FORMAT_ID,
                files_written=[],
                artifact_count=0,
                output_dir=output_dir,
            )

        # Group findings by resource_type
        by_type: dict[str, list[dict[str, object]]] = {}
        for f in findings:
            rt = str(f.get("resource_type", ""))
            if rt:
                by_type.setdefault(rt, []).append(f)

        rego_dir = output_dir / "opa" / cloud
        rego_dir.mkdir(parents=True, exist_ok=True)

        files_written: list[Path] = []
        now = datetime.now(UTC).isoformat()

        for resource_type, type_findings in sorted(by_type.items()):
            rules_text_parts: list[str] = []
            seen_controls: set[str] = set()

            for idx, finding in enumerate(type_findings):
                control_id = str(finding.get("control_id", ""))
                if control_id in seen_controls:
                    continue
                seen_controls.add(control_id)

                framework = str(finding.get("framework", ""))
                checkov_id = finding.get("checkov_id") or "N/A"
                native_ids = finding.get("native_ids")
                native_ids_str = str(native_ids) if native_ids else "N/A"
                severity = str(finding.get("severity", "MEDIUM"))
                description = str(finding.get("finding", control_id)).replace(
                    '"', '\\"'
                )

                # Derive check field from resource_type
                check_field = _derive_check_field(resource_type, control_id)

                rule_text = _DENY_RULE_TEMPLATE.format(
                    control_id=control_id,
                    framework=framework,
                    checkov_id=checkov_id,
                    native_ids=native_ids_str,
                    severity=severity,
                    tf_resource_type=resource_type,
                    rule_index=idx,
                    description=description,
                    check_field=check_field,
                )
                rules_text_parts.append(rule_text)

            # Use first finding for header metadata
            first = type_findings[0]
            file_content = _REGO_TEMPLATE.format(
                version=__version__,
                framework=str(first.get("framework", "")),
                control_id=str(first.get("control_id", "")),
                resource_type=resource_type,
                timestamp=now,
                policy_dir=str(rego_dir),
                cloud=cloud,
                rules="\n".join(rules_text_parts),
            )

            file_path = rego_dir / f"{resource_type}.rego"
            file_path.write_text(file_content, encoding="utf-8")
            files_written.append(file_path)

        return GenerationResult(
            format=self.FORMAT_ID,
            files_written=files_written,
            artifact_count=len(files_written),
            output_dir=output_dir,
        )


def _derive_check_field(resource_type: str, control_id: str) -> str:
    """Derive a reasonable check field from resource type context."""
    # Common patterns for compliance checks
    rt_lower = resource_type.lower()
    if "kms" in rt_lower or "crypto" in rt_lower:
        return "kms_key_id"
    if "log" in rt_lower or "trail" in rt_lower:
        return "enabled"
    if "sql" in rt_lower or "db" in rt_lower or "database" in rt_lower:
        return "encryption_configuration"
    if "storage" in rt_lower or "bucket" in rt_lower or "s3" in rt_lower:
        return "encryption"
    if "iam" in rt_lower:
        return "policy"
    if "network" in rt_lower or "firewall" in rt_lower or "security_group" in rt_lower:
        return "direction"
    return "configuration"
