"""Custom Checkov policy generator (Pro+ tier).

Generates Python check classes for Checkov custom policies where built-in checks
do not cover the compliance control.
"""

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

_CATEGORY_MAP: dict[str, str] = {
    "database": "ENCRYPTION",
    "object_storage": "ENCRYPTION",
    "disk": "ENCRYPTION",
    "kms": "ENCRYPTION",
    "iam": "IAM",
    "network": "NETWORKING",
    "load_balancer": "NETWORKING",
    "api_gateway": "NETWORKING",
    "dns": "NETWORKING",
    "cdn": "NETWORKING",
    "logging": "LOGGING",
    "monitoring": "LOGGING",
    "compute": "GENERAL_SECURITY",
    "container": "GENERAL_SECURITY",
    "serverless": "GENERAL_SECURITY",
}

_OPERATOR_TEMPLATES: dict[str, str] = {
    "equals": (
        '        val = conf.get("{attr}", [None])[0]\n        return val == {expected}'
    ),
    "not_empty": (
        '        val = conf.get("{attr}", [""])[0]\n        return bool(val)'
    ),
    "exists": '        return "{attr}" in conf',
    "not_equals": (
        '        val = conf.get("{attr}", [None])[0]\n        return val != {expected}'
    ),
    "in_list": (
        '        val = conf.get("{attr}", [None])[0]\n        return val in {expected}'
    ),
    "gte": (
        '        val = int(conf.get("{attr}", [0])[0])\n'
        "        return val >= {expected}"
    ),
    "regex_match": (
        "        import re\n"
        '        val = conf.get("{attr}", [""])[0]\n'
        "        return bool(re.match({expected}, str(val)))"
    ),
}

_POLICY_TEMPLATE = '''\
"""Auto-generated Checkov custom policy by ComplyForm v{tool_version}.

Control: {control_id} ({framework})
Resource: {resource_type}
Generated: {generated_at}
"""

from __future__ import annotations

from checkov.common.models.enums import CheckCategories, CheckResult
from checkov.terraform.checks.resource.base_resource_check import BaseResourceCheck


class {class_name}(BaseResourceCheck):
    """Check: {name}"""

    def __init__(self) -> None:
        name = "{name}"
        id = "{check_id}"
        supported_resources = {supported_resources!r}
        categories = [CheckCategories.{category}]
        super().__init__(
            name=name,
            id=id,
            categories=categories,
            supported_resources=list(supported_resources),
        )

    def scan_resource_conf(self, conf: dict) -> CheckResult:  # type: ignore[type-arg]
        """Evaluate resource configuration."""
{scan_body}
        return CheckResult.PASSED

check = {class_name}()
'''

_INIT_TEMPLATE = '''\
"""Auto-generated Checkov custom policy bundle by ComplyForm v{tool_version}.

Generated: {generated_at}
Cloud: {cloud}
Policy count: {count}
"""

from __future__ import annotations

from pathlib import Path

CHECKS_DIR = Path(__file__).parent
'''


@dataclass
class CheckovCustomPolicy:
    """Single custom Checkov policy definition."""

    check_id: str
    class_name: str
    name: str
    control_id: str
    framework: str
    resource_type: str
    supported_resources: tuple[str, ...]
    category: str
    attribute_path: str
    operator: str
    expected: str | None
    failure_message: str
    cloud: str
    generated_at: str
    tool_version: str


@dataclass
class CheckovCustomBundle:
    """Collection of custom Checkov policies for a cloud."""

    cloud: str
    policies: list[CheckovCustomPolicy] = field(default_factory=list)
    id_counter: dict[str, int] = field(default_factory=dict)


def _to_class_name(check_id: str) -> str:
    """Convert CKV_COMPLYFORM_GCP_001 to CkvComplyformGcp001."""
    parts = check_id.split("_")
    return "".join(p.capitalize() for p in parts)


def _get_category(resource_category: str) -> str:
    """Map resource_category to CheckCategories enum string."""
    return _CATEGORY_MAP.get(resource_category, "GENERAL_SECURITY")


def _build_scan_body(operator: str, attr: str, expected: str | None) -> str:
    """Build the scan_resource_conf body for the given operator."""
    template = _OPERATOR_TEMPLATES.get(operator, _OPERATOR_TEMPLATES["not_empty"])
    if expected is not None and operator not in ("gte",):
        expected_val = f'"{expected}"'
    else:
        expected_val = expected or '""'
    if operator == "in_list" and expected is not None:
        # Parse comma-separated list
        items = [item.strip() for item in expected.split(",")]
        expected_val = repr(items)
    body = template.format(attr=attr, expected=expected_val)
    return (
        f"        # Operator: {operator}\n"
        f"{body}\n"
        f"        if not self:\n"
        f"            return CheckResult.FAILED\n"
    )


def _build_scan_body_v2(operator: str, attr: str, expected: str | None) -> str:
    """Build scan body with proper return logic."""
    template = _OPERATOR_TEMPLATES.get(operator, _OPERATOR_TEMPLATES["not_empty"])

    if operator == "in_list" and expected is not None:
        items = [item.strip() for item in expected.split(",")]
        expected_val = repr(items)
    elif operator == "gte" and expected is not None:
        expected_val = expected
    elif operator in ("regex_match",) and expected is not None:
        expected_val = f'r"{expected}"'
    elif expected is not None:
        expected_val = f'"{expected}"'
    else:
        expected_val = '""'

    check_line = template.format(attr=attr, expected=expected_val)
    return f"        # Operator: {operator} | Attribute: {attr}\n{check_line}\n"


class CheckovCustomGenerator(PolicyGenerator):
    """Generate custom Checkov policies from compliance findings."""

    FORMAT_ID = "checkov"
    TIER_FEATURE = "checkov_custom"

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

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

    def format_name(self) -> str:
        return "Custom Checkov"

    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,
            )

        bundle = CheckovCustomBundle(cloud=cloud)
        cloud_upper = cloud.upper()
        bundle.id_counter[cloud_upper] = 0
        now = datetime.now(UTC).isoformat()

        # Build mappings lookup for enforcement_tier and checkov_ids
        mapping_entries = _extract_mapping_entries(mappings)

        for finding in findings:
            control_id = str(finding.get("control_id", ""))
            resource_type = str(finding.get("resource_type", ""))

            # Lookup mapping info
            mapping_info = mapping_entries.get(f"{control_id}:{resource_type}")
            if mapping_info is None:
                mapping_info = mapping_entries.get(control_id, {})

            checkov_ids = mapping_info.get("checkov_ids") or finding.get("checkov_id")
            enforcement_tier_val = mapping_info.get("enforcement_tier", "")

            # Skip if Checkov already covers this
            if checkov_ids:
                if isinstance(checkov_ids, list) and len(checkov_ids) > 0:
                    continue
                if isinstance(checkov_ids, str) and checkov_ids:
                    continue

            # Skip if not fully automatable
            if enforcement_tier_val not in ("fully_automatable", ""):
                continue

            # Generate policy
            bundle.id_counter[cloud_upper] = bundle.id_counter.get(cloud_upper, 0) + 1
            counter = bundle.id_counter[cloud_upper]
            check_id = f"CKV_COMPLYFORM_{cloud_upper}_{counter:03d}"

            resource_category = str(mapping_info.get("resource_category", ""))
            category = _get_category(resource_category)

            # Default attribute path and operator
            attr_path = str(mapping_info.get("attribute_path", "configuration"))
            operator = str(mapping_info.get("operator", "not_empty"))
            expected_val = mapping_info.get("expected")
            expected_str = str(expected_val) if expected_val is not None else None

            framework = str(finding.get("framework", ""))
            name = f"Ensure {resource_type} complies with {framework} {control_id}"

            policy = CheckovCustomPolicy(
                check_id=check_id,
                class_name=_to_class_name(check_id),
                name=name,
                control_id=control_id,
                framework=framework,
                resource_type=resource_type,
                supported_resources=(resource_type,),
                category=category,
                attribute_path=attr_path,
                operator=operator,
                expected=expected_str,
                failure_message=f"{control_id}: {str(finding.get('finding', ''))}",
                cloud=cloud,
                generated_at=now,
                tool_version=__version__,
            )
            bundle.policies.append(policy)

        # Write files
        checkov_dir = output_dir / "checkov"
        checkov_dir.mkdir(parents=True, exist_ok=True)
        files_written: list[Path] = []

        for policy in bundle.policies:
            scan_body = _build_scan_body_v2(
                policy.operator, policy.attribute_path, policy.expected
            )
            content = _POLICY_TEMPLATE.format(
                tool_version=policy.tool_version,
                control_id=policy.control_id,
                framework=policy.framework,
                resource_type=policy.resource_type,
                generated_at=policy.generated_at,
                class_name=policy.class_name,
                name=policy.name,
                check_id=policy.check_id,
                supported_resources=policy.supported_resources,
                category=policy.category,
                scan_body=scan_body,
            )
            file_path = checkov_dir / f"{policy.check_id}.py"
            file_path.write_text(content, encoding="utf-8")
            files_written.append(file_path)

        # Write __init__.py
        init_content = _INIT_TEMPLATE.format(
            tool_version=__version__,
            generated_at=now,
            cloud=cloud,
            count=len(bundle.policies),
        )
        init_path = checkov_dir / "__init__.py"
        init_path.write_text(init_content, encoding="utf-8")
        files_written.append(init_path)

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


def _extract_mapping_entries(
    mappings: dict[str, object],
) -> dict[str, dict[str, object]]:
    """Extract mapping entries into a lookup dict."""
    result: dict[str, dict[str, object]] = {}
    raw_mappings = mappings.get("mappings")
    if not isinstance(raw_mappings, list):
        return result
    for entry in raw_mappings:
        if not isinstance(entry, dict):
            continue
        control_id = str(entry.get("control_id", ""))
        resource_types = entry.get("resource_types", [])
        check_id = entry.get("check_id", "")

        info: dict[str, object] = {
            "checkov_ids": [check_id] if check_id else [],
            "enforcement_tier": entry.get("enforcement_tier", ""),
            "resource_category": entry.get("resource_category", ""),
            "attribute_path": entry.get("attribute_path", "configuration"),
            "operator": entry.get("operator", "not_empty"),
            "expected": entry.get("expected"),
        }
        result[control_id] = info
        if isinstance(resource_types, list):
            for rt in resource_types:
                result[f"{control_id}:{rt}"] = info
    return result
