"""GDPR and NIS2 compliance documentation generator for Prism.

Generates markdown compliance documents based on the project spec's
privacy configuration and model-level GDPR annotations.
"""

from __future__ import annotations

from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from prisme.generators.base import GeneratedFile, GeneratorBase
from prisme.spec.privacy import DataClassification
from prisme.spec.stack import FileStrategy
from prisme.utils.case_conversion import to_snake_case
from prisme.utils.template_engine import TemplateRenderer


class ComplianceDocsGenerator(GeneratorBase):
    """Generates GDPR and NIS2 compliance documentation.

    Produces markdown files documenting data classification, GDPR
    compliance status, and NIS2 assessment based on the spec.

    Only generates when privacy.enabled=True and
    privacy.generate_compliance_docs=True.
    """

    REQUIRED_TEMPLATES = [
        "docs/compliance/GDPR-COMPLIANCE.md.jinja2",
        "docs/compliance/DATA-CLASSIFICATION.md.jinja2",
    ]

    NIS2_TEMPLATES = [
        "docs/compliance/NIS2-COMPLIANCE.md.jinja2",
    ]

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)

        privacy_config = self._get_privacy_config()
        if (
            not privacy_config
            or not privacy_config.enabled
            or not privacy_config.generate_compliance_docs
        ):
            self.skip_generation = True
            return

        self.skip_generation = False
        self.privacy_config = privacy_config
        self.renderer = TemplateRenderer()

        required = list(self.REQUIRED_TEMPLATES)
        if privacy_config.nis2_enabled:
            required.extend(self.NIS2_TEMPLATES)
        self.renderer.validate_templates_exist(required)

        self.docs_path = Path("docs") / "compliance"

    def _get_privacy_config(self) -> Any:
        """Get privacy config from project spec, or None."""
        if self.project_spec and hasattr(self.project_spec, "privacy"):
            return self.project_spec.privacy
        return None

    def _build_model_field_info(self, model: Any) -> list[dict[str, Any]]:
        """Build field info list for a model, including privacy data."""
        fields = []
        for field in model.fields:
            info: dict[str, Any] = {
                "name": field.name,
                "type": field.type.value,
                "classification": DataClassification.INTERNAL.value,
                "is_pii": False,
                "encrypt_at_rest": False,
                "redact_in_logs": False,
                "audit_access": False,
            }
            if field.privacy:
                info["classification"] = field.privacy.classification.value
                info["is_pii"] = field.privacy.is_pii
                info["encrypt_at_rest"] = field.privacy.encrypt_at_rest
                info["redact_in_logs"] = field.privacy.redact_in_logs
                info["audit_access"] = field.privacy.audit_access
            fields.append(info)
        return fields

    def _build_all_models_context(self) -> list[dict[str, Any]]:
        """Build context for all models with field-level classification."""
        models = []
        for model in self.spec.models:
            is_personal = bool(model.gdpr and model.gdpr.is_personal_data)
            legal_basis = model.gdpr.legal_basis.value if model.gdpr else "not_applicable"
            all_fields = self._build_model_field_info(model)

            models.append(
                {
                    "name": model.name,
                    "snake_name": to_snake_case(model.name),
                    "is_personal_data": is_personal,
                    "legal_basis": legal_basis,
                    "all_fields": all_fields,
                }
            )
        return models

    def _build_personal_data_models_context(self) -> list[dict[str, Any]]:
        """Build context for models marked as personal data."""
        models = []
        for model in self.spec.models:
            if not model.gdpr or not model.gdpr.is_personal_data:
                continue

            pii_fields = [f.name for f in model.fields if f.privacy and f.privacy.is_pii]
            fields_with_privacy = [
                {
                    "name": f.name,
                    "classification": f.privacy.classification.value if f.privacy else "internal",
                    "is_pii": f.privacy.is_pii if f.privacy else False,
                    "encrypt_at_rest": f.privacy.encrypt_at_rest if f.privacy else False,
                    "redact_in_logs": f.privacy.redact_in_logs if f.privacy else False,
                }
                for f in model.fields
                if f.privacy
            ]

            retention_info = None
            if model.gdpr.retention:
                retention_info = {
                    "retention_days": model.gdpr.retention.retention_days,
                    "auto_delete": model.gdpr.retention.auto_delete,
                    "soft_delete_first": model.gdpr.retention.soft_delete_first,
                    "archive_before_delete": model.gdpr.retention.archive_before_delete,
                }

            models.append(
                {
                    "name": model.name,
                    "purpose": model.gdpr.purpose,
                    "legal_basis": model.gdpr.legal_basis.value,
                    "pii_fields": pii_fields,
                    "fields_with_privacy": fields_with_privacy,
                    "retention": retention_info,
                    "retention_days": model.gdpr.retention.retention_days
                    if model.gdpr.retention
                    else None,
                    "auto_delete": model.gdpr.retention.auto_delete
                    if model.gdpr.retention
                    else False,
                }
            )
        return models

    def _build_retention_models_context(self) -> list[dict[str, Any]]:
        """Build context for models with retention policies."""
        models = []
        for model in self.spec.models:
            if model.gdpr and model.gdpr.retention:
                r = model.gdpr.retention
                models.append(
                    {
                        "name": model.name,
                        "retention_days": r.retention_days,
                        "auto_delete": r.auto_delete,
                        "soft_delete_first": r.soft_delete_first,
                        "archive_before_delete": r.archive_before_delete,
                    }
                )
        return models

    def _count_fields(self, predicate: str) -> int:
        """Count fields across all models matching a privacy predicate."""
        count = 0
        for model in self.spec.models:
            for field in model.fields:
                if field.privacy and getattr(field.privacy, predicate, False):
                    count += 1
        return count

    def _get_export_formats(self) -> list[str]:
        """Get union of export formats across all personal data models."""
        formats: set[str] = set()
        for model in self.spec.models:
            if model.gdpr and model.gdpr.is_personal_data:
                formats.update(model.gdpr.data_subject_rights.export_formats)
        return sorted(formats) if formats else ["json", "csv"]

    def _has_any_right(self, right_name: str) -> bool:
        """Check if any personal data model has a specific right enabled."""
        for model in self.spec.models:
            if (
                model.gdpr
                and model.gdpr.is_personal_data
                and getattr(model.gdpr.data_subject_rights, right_name, False)
            ):
                return True
        return False

    def _base_context(self) -> dict[str, Any]:
        """Build the shared template context."""
        all_models = self._build_all_models_context()
        personal_data_models = self._build_personal_data_models_context()
        retention_models = self._build_retention_models_context()

        has_auth = bool(
            self.project_spec
            and hasattr(self.project_spec, "auth")
            and self.project_spec.auth.enabled
        )

        return {
            "organization_name": self.privacy_config.organization_name,
            "dpo_email": self.privacy_config.dpo_email,
            "privacy_policy_url": self.privacy_config.privacy_policy_url,
            "data_processing_region": self.privacy_config.data_processing_region,
            "generated_at": datetime.now(UTC).strftime("%Y-%m-%d %H:%M UTC"),
            "all_models": all_models,
            "personal_data_models": personal_data_models,
            "retention_models": retention_models,
            "subprocessors": self.privacy_config.subprocessors,
            "has_audit_log": self.privacy_config.generate_audit_log,
            "has_retention": len(retention_models) > 0,
            "has_export": self._has_any_right("right_to_access"),
            "has_deletion": self._has_any_right("right_to_erasure"),
            "has_restriction": self._has_any_right("right_to_restrict_processing"),
            "has_classification": True,
            "has_auth": has_auth,
            "has_mfa": False,
            "export_formats": self._get_export_formats(),
            "total_pii_fields": self._count_fields("is_pii"),
            "total_encrypted_fields": self._count_fields("encrypt_at_rest"),
            "total_audited_fields": self._count_fields("audit_access"),
            # NIS2
            "nis2_sector": self.privacy_config.nis2_sector,
            "nis2_entity_type": self.privacy_config.nis2_entity_type,
        }

    def generate_files(self) -> list[GeneratedFile]:
        """Generate all compliance documentation files."""
        if self.skip_generation:
            return []

        ctx = self._base_context()
        files = [
            self._generate_gdpr_doc(ctx),
            self._generate_classification_doc(ctx),
        ]

        if self.privacy_config.nis2_enabled:
            files.append(self._generate_nis2_doc(ctx))

        return files

    def _generate_gdpr_doc(self, ctx: dict[str, Any]) -> GeneratedFile:
        content = self.renderer.render_file(
            "docs/compliance/GDPR-COMPLIANCE.md.jinja2",
            context=ctx,
        )
        return GeneratedFile(
            path=self.docs_path / "GDPR-COMPLIANCE.md",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description="GDPR compliance documentation",
        )

    def _generate_classification_doc(self, ctx: dict[str, Any]) -> GeneratedFile:
        content = self.renderer.render_file(
            "docs/compliance/DATA-CLASSIFICATION.md.jinja2",
            context=ctx,
        )
        return GeneratedFile(
            path=self.docs_path / "DATA-CLASSIFICATION.md",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description="Data classification matrix",
        )

    def _generate_nis2_doc(self, ctx: dict[str, Any]) -> GeneratedFile:
        content = self.renderer.render_file(
            "docs/compliance/NIS2-COMPLIANCE.md.jinja2",
            context=ctx,
        )
        return GeneratedFile(
            path=self.docs_path / "NIS2-COMPLIANCE.md",
            content=content,
            strategy=FileStrategy.ALWAYS_OVERWRITE,
            description="NIS2 compliance assessment",
        )
