"""Compliance scanning for Citadel.

Provides programmatic access to CIS Kubernetes Benchmark (via kube-bench)
and vulnerability scanning (via Trivy Operator) results. Aggregates
compliance data for dashboarding and automated remediation workflows.

Integrates with:
  - kube-bench CronJob deployed in security/compliance/kube-bench-job.yaml
  - Trivy Operator deployed via security/compliance/trivy-operator-values.yaml
  - Grafana dashboard in security/compliance/compliance-dashboard.json

Design decisions:
  - ComplianceScanResult and ComplianceFailure are frozen dataclasses
    (consistent with NodeInfo, DrainResult, CertificateConfig, etc.)
  - ComplianceScanner requires a live Kubernetes client for all operations
  - Results are parsed from ConfigMaps (kube-bench) and CRDs (Trivy)
  - Severity levels align with CIS benchmark terminology: PASS, FAIL, WARN, INFO
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

from kubernetes import client, config
from kubernetes.client.exceptions import ApiException

logger = logging.getLogger(__name__)

# kube-bench job configuration
KUBE_BENCH_NAMESPACE = "citadel-compliance"
KUBE_BENCH_CRONJOB_NAME = "kube-bench-scan"
KUBE_BENCH_CONFIGMAP_PREFIX = "kube-bench-results"
KUBE_BENCH_JOB_LABEL = "app.kubernetes.io/name=kube-bench"

# Trivy Operator CRD configuration
TRIVY_GROUP = "aquasecurity.github.io"
TRIVY_VERSION = "v1alpha1"
TRIVY_VULN_PLURAL = "vulnerabilityreports"
TRIVY_AUDIT_PLURAL = "configauditreports"
TRIVY_COMPLIANCE_PLURAL = "clustercompliancereports"

# CIS severity mapping
SEVERITY_CRITICAL = "CRITICAL"
SEVERITY_HIGH = "HIGH"
SEVERITY_MEDIUM = "MEDIUM"
SEVERITY_LOW = "LOW"
SEVERITY_UNKNOWN = "UNKNOWN"

VALID_SEVERITIES = frozenset({
    SEVERITY_CRITICAL,
    SEVERITY_HIGH,
    SEVERITY_MEDIUM,
    SEVERITY_LOW,
    SEVERITY_UNKNOWN,
})


@dataclass(frozen=True)
class ComplianceFailure:
    """A single CIS benchmark check failure or warning.

    Attributes:
        check_id: CIS benchmark check identifier (e.g., "1.2.3").
        description: Human-readable description of the check.
        severity: Severity level — CRITICAL, HIGH, MEDIUM, LOW, or UNKNOWN.
        remediation: Recommended fix for the failure.
    """

    check_id: str
    description: str
    severity: str
    remediation: str

    def __post_init__(self) -> None:
        if not self.check_id:
            raise ValueError("check_id must not be empty")
        if not self.description:
            raise ValueError("description must not be empty")
        if self.severity not in VALID_SEVERITIES:
            raise ValueError(
                f"severity must be one of {sorted(VALID_SEVERITIES)}, "
                f"got '{self.severity}'"
            )


@dataclass(frozen=True)
class ComplianceScanResult:
    """Result of a CIS benchmark compliance scan.

    Immutable snapshot of the scan output — safe to cache, serialize,
    and pass between layers without defensive copying.

    Attributes:
        benchmark: Benchmark identifier (e.g., "CIS Kubernetes Benchmark v1.8").
        score: Compliance score as a percentage (0.0 to 100.0).
        total_checks: Total number of checks evaluated.
        passed: Number of checks that passed.
        failed: Number of checks that failed.
        warnings: Number of checks with warnings.
        failures: Tuple of individual failure details.
        scanned_at: ISO 8601 timestamp of when the scan was performed.
    """

    benchmark: str
    score: float
    total_checks: int
    passed: int
    failed: int
    warnings: int
    failures: tuple[ComplianceFailure, ...]
    scanned_at: str

    def __post_init__(self) -> None:
        if not self.benchmark:
            raise ValueError("benchmark must not be empty")
        if not (0.0 <= self.score <= 100.0):
            raise ValueError(
                f"score must be between 0.0 and 100.0, got {self.score}"
            )
        if self.total_checks < 0:
            raise ValueError(
                f"total_checks must be >= 0, got {self.total_checks}"
            )
        if self.passed < 0:
            raise ValueError(f"passed must be >= 0, got {self.passed}")
        if self.failed < 0:
            raise ValueError(f"failed must be >= 0, got {self.failed}")
        if self.warnings < 0:
            raise ValueError(f"warnings must be >= 0, got {self.warnings}")
        if not isinstance(self.failures, tuple):
            raise TypeError("failures must be a tuple")

    @property
    def passing_rate(self) -> float:
        """Percentage of checks that passed (excluding warnings)."""
        evaluated = self.passed + self.failed
        if evaluated == 0:
            return 0.0
        return round((self.passed / evaluated) * 100.0, 2)

    @property
    def has_critical_failures(self) -> bool:
        """True if any failure is CRITICAL severity."""
        return any(f.severity == SEVERITY_CRITICAL for f in self.failures)

    @property
    def has_high_failures(self) -> bool:
        """True if any failure is HIGH severity."""
        return any(f.severity == SEVERITY_HIGH for f in self.failures)

    @property
    def failures_by_severity(self) -> dict[str, list[ComplianceFailure]]:
        """Group failures by severity level."""
        result: dict[str, list[ComplianceFailure]] = {}
        for failure in self.failures:
            result.setdefault(failure.severity, []).append(failure)
        return result


class ComplianceScanner:
    """Queries compliance scan results from kube-bench and Trivy Operator.

    Reads from:
      - ConfigMaps written by the kube-bench CronJob (CIS benchmarks)
      - Trivy Operator CRDs (vulnerability and config audit reports)

    This class is the read layer — it does not execute scans itself.
    Scans are triggered by the CronJob schedule or by creating Jobs
    directly via ``trigger_cis_scan()``.

    Usage::

        scanner = ComplianceScanner()
        result = scanner.run_cis_benchmark()
        if result.has_critical_failures:
            for f in result.failures:
                if f.severity == "CRITICAL":
                    logger.error("CIS %s: %s", f.check_id, f.description)

        vulns = scanner.get_vulnerability_report(namespace="default")
        summary = scanner.get_compliance_summary()
    """

    def __init__(self, kubeconfig: str | None = None) -> None:
        """Initialize with Kubernetes client configuration.

        Args:
            kubeconfig: Path to kubeconfig file. If None, tries in-cluster
                config first, then falls back to default kubeconfig.
        """
        if kubeconfig:
            config.load_kube_config(config_file=kubeconfig)
        else:
            try:
                config.load_incluster_config()
            except config.ConfigException:
                config.load_kube_config()

        self._core = client.CoreV1Api()
        self._batch = client.BatchV1Api()
        self._custom = client.CustomObjectsApi()

    def run_cis_benchmark(self) -> ComplianceScanResult:
        """Retrieve the latest CIS benchmark scan result from kube-bench.

        Reads the most recent kube-bench results ConfigMap in the
        citadel-compliance namespace. The ConfigMap is created by the
        kube-bench CronJob defined in security/compliance/kube-bench-job.yaml.

        Returns:
            ComplianceScanResult with the latest scan data.

        Raises:
            RuntimeError: If no kube-bench results are found.
            ApiException: On Kubernetes API errors.
        """
        try:
            configmaps = self._core.list_namespaced_config_map(
                namespace=KUBE_BENCH_NAMESPACE,
                label_selector=KUBE_BENCH_JOB_LABEL,
            )
        except ApiException as e:
            if e.status == 404:
                raise RuntimeError(
                    f"Namespace '{KUBE_BENCH_NAMESPACE}' not found. "
                    "Deploy the compliance stack first "
                    "(scripts/11-install-compliance.sh)."
                ) from e
            raise

        if not configmaps.items:
            raise RuntimeError(
                "No kube-bench results found. Run an initial scan or wait "
                "for the CronJob to execute."
            )

        # Sort by creation timestamp, newest first
        sorted_cms = sorted(
            configmaps.items,
            key=lambda cm: cm.metadata.creation_timestamp or datetime.min.replace(
                tzinfo=timezone.utc
            ),
            reverse=True,
        )

        latest = sorted_cms[0]
        return self._parse_kube_bench_results(latest)

    def trigger_cis_scan(self) -> str:
        """Create a one-off Job from the kube-bench CronJob template.

        This triggers an immediate scan rather than waiting for the
        next CronJob schedule. The Job name is returned so callers
        can poll for completion.

        Returns:
            The name of the created Job.

        Raises:
            ApiException: On Kubernetes API errors.
        """
        try:
            cronjob = self._batch.read_namespaced_cron_job(
                name=KUBE_BENCH_CRONJOB_NAME,
                namespace=KUBE_BENCH_NAMESPACE,
            )
        except ApiException as e:
            if e.status == 404:
                raise RuntimeError(
                    f"CronJob '{KUBE_BENCH_CRONJOB_NAME}' not found in "
                    f"namespace '{KUBE_BENCH_NAMESPACE}'. Deploy the "
                    "compliance stack first (scripts/11-install-compliance.sh)."
                ) from e
            raise

        # Build a Job from the CronJob's job template
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
        job_name = f"kube-bench-manual-{timestamp}"

        job = client.V1Job(
            api_version="batch/v1",
            kind="Job",
            metadata=client.V1ObjectMeta(
                name=job_name,
                namespace=KUBE_BENCH_NAMESPACE,
                labels={
                    "app.kubernetes.io/name": "kube-bench",
                    "citadel.dev/trigger": "manual",
                },
            ),
            spec=cronjob.spec.job_template.spec,
        )

        self._batch.create_namespaced_job(
            namespace=KUBE_BENCH_NAMESPACE,
            body=job,
        )

        logger.info(
            "Triggered manual kube-bench scan: %s/%s",
            KUBE_BENCH_NAMESPACE,
            job_name,
        )
        return job_name

    def get_vulnerability_report(self, namespace: str = "default") -> list[dict]:
        """Query Trivy Operator VulnerabilityReports in a namespace.

        Returns raw vulnerability report dicts from the Trivy Operator
        CRDs. Each report corresponds to a container image scan.

        Args:
            namespace: Kubernetes namespace to query.

        Returns:
            List of vulnerability report dicts. Empty list if Trivy
            Operator CRDs are not installed or no reports exist.
        """
        try:
            result = self._custom.list_namespaced_custom_object(
                group=TRIVY_GROUP,
                version=TRIVY_VERSION,
                namespace=namespace,
                plural=TRIVY_VULN_PLURAL,
            )
            return result.get("items", [])
        except ApiException as e:
            if e.status == 404:
                logger.warning(
                    "Trivy Operator CRDs not found. Is Trivy Operator installed?"
                )
                return []
            logger.error(
                "Failed to list vulnerability reports in %s: %s",
                namespace,
                e.reason,
            )
            raise

    def get_config_audit_reports(self, namespace: str = "default") -> list[dict]:
        """Query Trivy Operator ConfigAuditReports in a namespace.

        ConfigAuditReports identify workload misconfiguration such as
        running as root, missing resource limits, or writable root
        filesystems.

        Args:
            namespace: Kubernetes namespace to query.

        Returns:
            List of config audit report dicts.
        """
        try:
            result = self._custom.list_namespaced_custom_object(
                group=TRIVY_GROUP,
                version=TRIVY_VERSION,
                namespace=namespace,
                plural=TRIVY_AUDIT_PLURAL,
            )
            return result.get("items", [])
        except ApiException as e:
            if e.status == 404:
                logger.warning(
                    "Trivy Operator CRDs not found. Is Trivy Operator installed?"
                )
                return []
            logger.error(
                "Failed to list config audit reports in %s: %s",
                namespace,
                e.reason,
            )
            raise

    def get_compliance_summary(self) -> dict[str, Any]:
        """Aggregate compliance status across CIS benchmarks and Trivy reports.

        Returns a summary dict with:
          - cis_benchmark: latest CIS scan result (or None if unavailable)
          - vulnerability_counts: {CRITICAL, HIGH, MEDIUM, LOW} tallies
          - config_audit_counts: {CRITICAL, HIGH, MEDIUM, LOW} tallies
          - overall_compliant: bool — True if no CRITICAL or HIGH issues

        Returns:
            Summary dict. Never raises — returns partial data on errors.
        """
        summary: dict[str, Any] = {
            "cis_benchmark": None,
            "vulnerability_counts": {
                "CRITICAL": 0,
                "HIGH": 0,
                "MEDIUM": 0,
                "LOW": 0,
            },
            "config_audit_counts": {
                "CRITICAL": 0,
                "HIGH": 0,
                "MEDIUM": 0,
                "LOW": 0,
            },
            "overall_compliant": True,
            "scanned_at": datetime.now(timezone.utc).isoformat(),
        }

        # CIS benchmark
        try:
            cis_result = self.run_cis_benchmark()
            summary["cis_benchmark"] = {
                "benchmark": cis_result.benchmark,
                "score": cis_result.score,
                "total_checks": cis_result.total_checks,
                "passed": cis_result.passed,
                "failed": cis_result.failed,
                "warnings": cis_result.warnings,
            }
            if cis_result.has_critical_failures or cis_result.has_high_failures:
                summary["overall_compliant"] = False
        except (RuntimeError, ApiException) as e:
            logger.warning("Could not retrieve CIS benchmark results: %s", e)

        # Vulnerability reports (scan all namespaces)
        try:
            vuln_reports = self._custom.list_custom_object_for_all_namespaces(
                group=TRIVY_GROUP,
                version=TRIVY_VERSION,
                plural=TRIVY_VULN_PLURAL,
            )
            for report in vuln_reports.get("items", []):
                vulns = (
                    report.get("report", {}).get("vulnerabilities", [])
                )
                for vuln in vulns:
                    sev = vuln.get("severity", "UNKNOWN").upper()
                    if sev in summary["vulnerability_counts"]:
                        summary["vulnerability_counts"][sev] += 1
                        if sev in ("CRITICAL", "HIGH"):
                            summary["overall_compliant"] = False
        except ApiException:
            logger.warning("Could not retrieve Trivy vulnerability reports")

        # Config audit reports (scan all namespaces)
        try:
            audit_reports = self._custom.list_custom_object_for_all_namespaces(
                group=TRIVY_GROUP,
                version=TRIVY_VERSION,
                plural=TRIVY_AUDIT_PLURAL,
            )
            for report in audit_reports.get("items", []):
                checks = report.get("report", {}).get("checks", [])
                for check in checks:
                    if not check.get("success", True):
                        sev = check.get("severity", "UNKNOWN").upper()
                        if sev in summary["config_audit_counts"]:
                            summary["config_audit_counts"][sev] += 1
                            if sev in ("CRITICAL", "HIGH"):
                                summary["overall_compliant"] = False
        except ApiException:
            logger.warning("Could not retrieve Trivy config audit reports")

        return summary

    def list_non_compliant_workloads(
        self, namespace: str = "default"
    ) -> list[dict[str, Any]]:
        """List workloads with compliance violations in a namespace.

        Combines Trivy vulnerability reports and config audit reports
        to produce a unified list of non-compliant workloads with their
        violation details.

        Args:
            namespace: Kubernetes namespace to query.

        Returns:
            List of dicts, each containing:
              - name: workload name
              - namespace: workload namespace
              - kind: workload kind (Deployment, StatefulSet, etc.)
              - vulnerabilities: count of HIGH+ vulnerabilities
              - misconfigurations: count of HIGH+ misconfigurations
              - details: list of violation descriptions
        """
        non_compliant: dict[str, dict[str, Any]] = {}

        # Vulnerability reports
        vuln_reports = self.get_vulnerability_report(namespace)
        for report in vuln_reports:
            owner = self._extract_workload_owner(report)
            key = f"{owner['namespace']}/{owner['kind']}/{owner['name']}"

            if key not in non_compliant:
                non_compliant[key] = {
                    "name": owner["name"],
                    "namespace": owner["namespace"],
                    "kind": owner["kind"],
                    "vulnerabilities": 0,
                    "misconfigurations": 0,
                    "details": [],
                }

            vulns = report.get("report", {}).get("vulnerabilities", [])
            high_plus = [
                v for v in vulns
                if v.get("severity", "").upper() in ("CRITICAL", "HIGH")
            ]
            non_compliant[key]["vulnerabilities"] += len(high_plus)
            for v in high_plus:
                non_compliant[key]["details"].append(
                    f"[{v.get('severity', 'UNKNOWN')}] {v.get('vulnerabilityID', 'N/A')}: "
                    f"{v.get('title', v.get('description', 'No description'))}"
                )

        # Config audit reports
        audit_reports = self.get_config_audit_reports(namespace)
        for report in audit_reports:
            owner = self._extract_workload_owner(report)
            key = f"{owner['namespace']}/{owner['kind']}/{owner['name']}"

            if key not in non_compliant:
                non_compliant[key] = {
                    "name": owner["name"],
                    "namespace": owner["namespace"],
                    "kind": owner["kind"],
                    "vulnerabilities": 0,
                    "misconfigurations": 0,
                    "details": [],
                }

            checks = report.get("report", {}).get("checks", [])
            failed_checks = [
                c for c in checks
                if not c.get("success", True)
                and c.get("severity", "").upper() in ("CRITICAL", "HIGH")
            ]
            non_compliant[key]["misconfigurations"] += len(failed_checks)
            for c in failed_checks:
                non_compliant[key]["details"].append(
                    f"[{c.get('severity', 'UNKNOWN')}] {c.get('checkID', 'N/A')}: "
                    f"{c.get('title', c.get('description', 'No description'))}"
                )

        # Filter to only workloads with actual violations
        return [
            wl for wl in non_compliant.values()
            if wl["vulnerabilities"] > 0 or wl["misconfigurations"] > 0
        ]

    def _parse_kube_bench_results(
        self, configmap: client.V1ConfigMap
    ) -> ComplianceScanResult:
        """Parse a kube-bench results ConfigMap into a ComplianceScanResult.

        The ConfigMap is expected to contain a 'results.json' key with
        kube-bench JSON output format.
        """
        data = configmap.data or {}
        raw_json = data.get("results.json", "{}")

        try:
            results = json.loads(raw_json)
        except json.JSONDecodeError as e:
            raise RuntimeError(
                f"Failed to parse kube-bench results from ConfigMap "
                f"'{configmap.metadata.name}': {e}"
            ) from e

        # Parse kube-bench JSON output structure
        total_pass = 0
        total_fail = 0
        total_warn = 0
        total_info = 0
        failures: list[ComplianceFailure] = []

        controls = results if isinstance(results, list) else results.get("Controls", [])

        for control in controls:
            tests = control.get("tests", [])
            for test in tests:
                for check_result in test.get("results", []):
                    status = check_result.get("status", "").upper()
                    if status == "PASS":
                        total_pass += 1
                    elif status == "FAIL":
                        total_fail += 1
                        severity = self._map_cis_severity(check_result)
                        failures.append(
                            ComplianceFailure(
                                check_id=check_result.get(
                                    "test_number", "unknown"
                                ),
                                description=check_result.get(
                                    "test_desc", "No description"
                                ),
                                severity=severity,
                                remediation=check_result.get(
                                    "remediation", "See CIS benchmark documentation"
                                ),
                            )
                        )
                    elif status == "WARN":
                        total_warn += 1
                    elif status == "INFO":
                        total_info += 1

        total_checks = total_pass + total_fail + total_warn + total_info
        score = 0.0
        if total_checks > 0:
            score = round((total_pass / total_checks) * 100.0, 2)

        # Extract benchmark version from results or use default
        benchmark = "CIS Kubernetes Benchmark v1.8"
        if isinstance(results, dict) and results.get("version"):
            benchmark = f"CIS Kubernetes Benchmark {results['version']}"

        scanned_at = datetime.now(timezone.utc).isoformat()
        if configmap.metadata.creation_timestamp:
            scanned_at = configmap.metadata.creation_timestamp.isoformat()

        return ComplianceScanResult(
            benchmark=benchmark,
            score=score,
            total_checks=total_checks,
            passed=total_pass,
            failed=total_fail,
            warnings=total_warn,
            failures=tuple(failures),
            scanned_at=scanned_at,
        )

    @staticmethod
    def _map_cis_severity(check_result: dict) -> str:
        """Map a kube-bench check result to a severity level.

        kube-bench does not natively assign severities — we infer from
        the section number and scored/unscored status. Section 1.x and
        4.x (control plane, policies) are typically higher severity.
        """
        test_number = check_result.get("test_number", "")
        scored = check_result.get("scored", True)

        if not scored:
            return SEVERITY_LOW

        # Control plane (1.x) and policies (5.x) checks are critical
        if test_number.startswith(("1.1", "1.2", "5.")):
            return SEVERITY_CRITICAL
        # API server (1.x) and kubelet (4.x) checks are high
        if test_number.startswith(("1.", "4.")):
            return SEVERITY_HIGH
        # Worker node (3.x) and etcd (2.x) are medium
        if test_number.startswith(("2.", "3.")):
            return SEVERITY_MEDIUM

        return SEVERITY_MEDIUM

    @staticmethod
    def _extract_workload_owner(report: dict) -> dict[str, str]:
        """Extract the owning workload information from a Trivy report.

        Trivy reports carry owner references in their labels, which
        we use to group reports by workload.
        """
        metadata = report.get("metadata", {})
        labels = metadata.get("labels", {})
        namespace = metadata.get("namespace", "unknown")

        # Trivy Operator labels the reports with owner info
        name = labels.get(
            "trivy-operator.resource.name",
            metadata.get("name", "unknown"),
        )
        kind = labels.get(
            "trivy-operator.resource.kind",
            "Unknown",
        )

        return {
            "name": name,
            "namespace": namespace,
            "kind": kind,
        }
