"""Custom compliance framework storage and schema.

Users author frameworks as YAML files and the backend loads them alongside
the built-in Firestore frameworks so the existing ``compliance_engine``
evaluator can score findings without any special-case logic.

Storage resolution order:
    1. ``$DSC_CUSTOM_FRAMEWORK_DIR`` env override
    2. ``<workspace>/.devseccode/custom-frameworks/`` (workspace-local)
    3. ``~/.devseccode/custom-frameworks/`` (user-global fallback)

The workspace path is reported by the IDE via the remote_ws channel
(``remote_ws.get_workspace_path()``).
"""

from __future__ import annotations

import logging
import os
import re
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import yaml

logger = logging.getLogger(__name__)

_DEFAULT_DIRNAME = "custom-frameworks"
_FRAMEWORK_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.\- ]{1,62}$")
_CONTROL_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.\-()/ ]{0,63}$")
_CWE_RE = re.compile(r"^CWE-\d{1,6}$")
_ALLOWED_SEVERITY = {"critical", "high", "medium", "low", "info"}
_ALLOWED_DETECTION = {"scanner", "attestation"}


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------


@dataclass(slots=True)
class CustomControl:
    control_id: str
    control_name: str
    description: str = ""
    category: str = ""
    severity_weight: str = "medium"
    detection_method: str = "scanner"
    mapped_cwes: list[str] = field(default_factory=list)
    evidence_guidance: str = ""
    attestation_required: bool = False
    code_relevant: bool = True

    def to_dict(self) -> dict[str, Any]:
        return {
            "control_id": self.control_id,
            "control_name": self.control_name,
            "description": self.description,
            "category": self.category,
            "severity_weight": self.severity_weight,
            "detection_method": self.detection_method,
            "mapped_cwes": list(self.mapped_cwes),
            "evidence_guidance": self.evidence_guidance,
            "attestation_required": self.attestation_required,
            "code_relevant": self.code_relevant,
        }


@dataclass(slots=True)
class CustomCategory:
    id: str
    name: str
    control_count: int = 0


@dataclass(slots=True)
class CustomFramework:
    framework_id: str
    display_name: str
    full_name: str = ""
    version: str = "1.0"
    description: str = ""
    default_attestation_expiry_days: int = 90
    categories: list[CustomCategory] = field(default_factory=list)
    controls: list[CustomControl] = field(default_factory=list)
    created_at: str = ""
    updated_at: str = ""
    author: str = ""

    @property
    def code_relevant_controls(self) -> int:
        return sum(1 for c in self.controls if c.code_relevant)

    def to_framework_dict(self) -> dict[str, Any]:
        """Serialize in the shape consumed by ``compliance_engine``'s cache."""
        return {
            "framework_id": self.framework_id,
            "display_name": self.display_name,
            "full_name": self.full_name or self.display_name,
            "version": self.version,
            "total_controls": len(self.controls),
            "code_relevant_controls": self.code_relevant_controls,
            "description": self.description,
            "default_attestation_expiry_days": self.default_attestation_expiry_days,
            "categories": [
                {"id": c.id, "name": c.name, "control_count": c.control_count}
                for c in self.categories
            ],
            "custom": True,
        }

    def to_control_dicts(self) -> list[dict[str, Any]]:
        """Control rows in the shape expected by ``compliance_engine``."""
        return [
            {
                "framework": self.framework_id,
                "control_id": c.control_id,
                "control_name": c.control_name,
                "description": c.description,
                "category": c.category,
                "severity_weight": c.severity_weight,
                "detection_method": c.detection_method,
                "mapped_cwes": list(c.mapped_cwes),
                "evidence_guidance": c.evidence_guidance,
                "attestation_required": c.attestation_required,
                "code_relevant": c.code_relevant,
                "custom": True,
            }
            for c in self.controls
        ]

    def to_wire_dict(self) -> dict[str, Any]:
        """Full serialization for the builder UI."""
        return {
            "frameworkId": self.framework_id,
            "displayName": self.display_name,
            "fullName": self.full_name,
            "version": self.version,
            "description": self.description,
            "defaultAttestationExpiryDays": self.default_attestation_expiry_days,
            "categories": [
                {"id": c.id, "name": c.name, "controlCount": c.control_count}
                for c in self.categories
            ],
            "controls": [
                {
                    "controlId": c.control_id,
                    "controlName": c.control_name,
                    "description": c.description,
                    "category": c.category,
                    "severityWeight": c.severity_weight,
                    "detectionMethod": c.detection_method,
                    "mappedCwes": list(c.mapped_cwes),
                    "evidenceGuidance": c.evidence_guidance,
                    "attestationRequired": c.attestation_required,
                    "codeRelevant": c.code_relevant,
                }
                for c in self.controls
            ],
            "createdAt": self.created_at,
            "updatedAt": self.updated_at,
            "author": self.author,
            "custom": True,
        }


# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------


class FrameworkValidationError(ValueError):
    def __init__(self, errors: list[str]):
        super().__init__("; ".join(errors))
        self.errors = errors


def _coerce_str(value: Any, default: str = "") -> str:
    if value is None:
        return default
    return str(value)


def validate_payload(payload: dict[str, Any]) -> list[str]:
    """Return a list of human-readable validation errors (empty = valid)."""
    errors: list[str] = []

    fw_id = _coerce_str(payload.get("frameworkId") or payload.get("framework_id"))
    if not _FRAMEWORK_ID_RE.match(fw_id):
        errors.append(
            "frameworkId must be 2-63 chars: letters, digits, spaces, or . _ - and must start with a letter or digit"
        )

    if not _coerce_str(payload.get("displayName") or payload.get("display_name")).strip():
        errors.append("displayName is required")

    controls = payload.get("controls") or []
    if not isinstance(controls, list):
        errors.append("controls must be a list")
        return errors

    seen_ids: set[str] = set()
    for idx, raw in enumerate(controls):
        prefix = f"controls[{idx}]"
        if not isinstance(raw, dict):
            errors.append(f"{prefix} must be an object")
            continue
        ctrl_id = _coerce_str(raw.get("controlId") or raw.get("control_id")).strip()
        if not _CONTROL_ID_RE.match(ctrl_id):
            errors.append(f"{prefix}.controlId invalid (got '{ctrl_id}')")
        elif ctrl_id in seen_ids:
            errors.append(f"{prefix}.controlId duplicate: {ctrl_id}")
        else:
            seen_ids.add(ctrl_id)

        if not _coerce_str(raw.get("controlName") or raw.get("control_name")).strip():
            errors.append(f"{prefix}.controlName is required")

        sev = _coerce_str(raw.get("severityWeight") or raw.get("severity_weight") or "medium").lower()
        if sev not in _ALLOWED_SEVERITY:
            errors.append(f"{prefix}.severityWeight must be one of {sorted(_ALLOWED_SEVERITY)}")

        detection = _coerce_str(
            raw.get("detectionMethod") or raw.get("detection_method") or "scanner"
        ).lower()
        if detection not in _ALLOWED_DETECTION:
            errors.append(f"{prefix}.detectionMethod must be one of {sorted(_ALLOWED_DETECTION)}")

        mapped_raw = raw.get("mappedCwes") or raw.get("mapped_cwes") or []
        if not isinstance(mapped_raw, list):
            errors.append(f"{prefix}.mappedCwes must be a list")
        else:
            for cwe in mapped_raw:
                if not isinstance(cwe, str) or not _CWE_RE.match(cwe):
                    errors.append(f"{prefix}.mappedCwes contains invalid CWE '{cwe}'")
            # Scanner-detected controls with no CWEs won't match any finding, but
            # blocking save would prevent iterative authoring. The evaluator will
            # just report them as "not_covered" until CWEs are added.

    return errors


def framework_from_payload(payload: dict[str, Any]) -> CustomFramework:
    """Coerce a JSON/YAML payload into a CustomFramework (raises on bad input)."""
    errs = validate_payload(payload)
    if errs:
        raise FrameworkValidationError(errs)

    fw_id = _coerce_str(payload.get("frameworkId") or payload.get("framework_id")).strip()
    categories_raw = payload.get("categories") or []
    categories: list[CustomCategory] = []
    for c in categories_raw:
        if not isinstance(c, dict):
            continue
        categories.append(
            CustomCategory(
                id=_coerce_str(c.get("id")).strip() or _coerce_str(c.get("name")).strip().lower().replace(" ", "-"),
                name=_coerce_str(c.get("name")).strip(),
                control_count=int(c.get("controlCount") or c.get("control_count") or 0),
            )
        )

    controls: list[CustomControl] = []
    for c in payload.get("controls") or []:
        controls.append(
            CustomControl(
                control_id=_coerce_str(c.get("controlId") or c.get("control_id")).strip(),
                control_name=_coerce_str(c.get("controlName") or c.get("control_name")).strip(),
                description=_coerce_str(c.get("description")),
                category=_coerce_str(c.get("category")).strip(),
                severity_weight=_coerce_str(
                    c.get("severityWeight") or c.get("severity_weight") or "medium"
                ).lower(),
                detection_method=_coerce_str(
                    c.get("detectionMethod") or c.get("detection_method") or "scanner"
                ).lower(),
                mapped_cwes=[
                    str(x)
                    for x in (c.get("mappedCwes") or c.get("mapped_cwes") or [])
                    if isinstance(x, str)
                ],
                evidence_guidance=_coerce_str(
                    c.get("evidenceGuidance") or c.get("evidence_guidance")
                ),
                attestation_required=bool(
                    c.get("attestationRequired") or c.get("attestation_required") or False
                ),
                code_relevant=bool(c.get("codeRelevant") or c.get("code_relevant") or True),
            )
        )

    # Auto-populate category control_count if missing
    if categories:
        counts: dict[str, int] = {}
        for ctrl in controls:
            if ctrl.category:
                counts[ctrl.category] = counts.get(ctrl.category, 0) + 1
        for cat in categories:
            if cat.control_count == 0:
                cat.control_count = counts.get(cat.id, 0)

    now = datetime.now(timezone.utc).isoformat()
    return CustomFramework(
        framework_id=fw_id,
        display_name=_coerce_str(payload.get("displayName") or payload.get("display_name")).strip(),
        full_name=_coerce_str(payload.get("fullName") or payload.get("full_name")),
        version=_coerce_str(payload.get("version") or "1.0"),
        description=_coerce_str(payload.get("description")),
        default_attestation_expiry_days=int(
            payload.get("defaultAttestationExpiryDays")
            or payload.get("default_attestation_expiry_days")
            or 90
        ),
        categories=categories,
        controls=controls,
        created_at=_coerce_str(payload.get("createdAt") or payload.get("created_at")) or now,
        updated_at=now,
        author=_coerce_str(payload.get("author")),
    )


# ---------------------------------------------------------------------------
# Storage directory resolution
# ---------------------------------------------------------------------------


def _storage_dir() -> Path:
    env = os.environ.get("DSC_CUSTOM_FRAMEWORK_DIR")
    if env:
        path = Path(env).expanduser()
        path.mkdir(parents=True, exist_ok=True)
        return path

    try:
        from dsc_api.services.remote_ws import get_workspace_path
        ws = get_workspace_path()
    except Exception:
        ws = ""

    if ws:
        path = Path(ws) / ".devseccode" / _DEFAULT_DIRNAME
    else:
        path = Path.home() / ".devseccode" / _DEFAULT_DIRNAME
    path.mkdir(parents=True, exist_ok=True)
    return path


def _slug(framework_id: str) -> str:
    """Filesystem-safe slug for a framework id."""
    slug = re.sub(r"[^A-Za-z0-9._-]+", "-", framework_id).strip("-").lower()
    return slug or f"framework-{uuid.uuid4().hex[:8]}"


# ---------------------------------------------------------------------------
# Public CRUD API
# ---------------------------------------------------------------------------


def list_frameworks() -> list[CustomFramework]:
    results: list[CustomFramework] = []
    for yml in sorted(_storage_dir().glob("*.yml")):
        try:
            results.append(_load_file(yml))
        except Exception as exc:
            logger.warning("Skipping malformed custom framework %s: %s", yml.name, exc)
    return results


def get_framework(framework_id: str) -> CustomFramework | None:
    for fw in list_frameworks():
        if fw.framework_id.upper() == framework_id.upper():
            return fw
    return None


def save_framework(payload: dict[str, Any]) -> CustomFramework:
    framework = framework_from_payload(payload)
    target = _storage_dir() / f"{_slug(framework.framework_id)}.yml"
    with target.open("w", encoding="utf-8") as fh:
        yaml.safe_dump(
            _to_yaml_dict(framework),
            fh,
            sort_keys=False,
            default_flow_style=False,
            allow_unicode=True,
        )
    logger.info("Saved custom framework %s at %s", framework.framework_id, target)
    return framework


def delete_framework(framework_id: str) -> bool:
    target = _storage_dir() / f"{_slug(framework_id)}.yml"
    if target.exists():
        target.unlink()
        logger.info("Deleted custom framework %s", framework_id)
        return True
    return False


def _load_file(path: Path) -> CustomFramework:
    with path.open("r", encoding="utf-8") as fh:
        raw = yaml.safe_load(fh) or {}
    return framework_from_payload(raw)


def _to_yaml_dict(fw: CustomFramework) -> dict[str, Any]:
    return {
        "frameworkId": fw.framework_id,
        "displayName": fw.display_name,
        "fullName": fw.full_name,
        "version": fw.version,
        "description": fw.description,
        "defaultAttestationExpiryDays": fw.default_attestation_expiry_days,
        "author": fw.author,
        "createdAt": fw.created_at,
        "updatedAt": fw.updated_at,
        "categories": [
            {"id": c.id, "name": c.name, "controlCount": c.control_count}
            for c in fw.categories
        ],
        "controls": [
            {
                "controlId": c.control_id,
                "controlName": c.control_name,
                "description": c.description,
                "category": c.category,
                "severityWeight": c.severity_weight,
                "detectionMethod": c.detection_method,
                "mappedCwes": list(c.mapped_cwes),
                "evidenceGuidance": c.evidence_guidance,
                "attestationRequired": c.attestation_required,
                "codeRelevant": c.code_relevant,
            }
            for c in fw.controls
        ],
    }


# ---------------------------------------------------------------------------
# Cache merge hook (called from compliance_engine)
# ---------------------------------------------------------------------------


def merge_into_caches(
    frameworks_cache: list[dict[str, Any]],
    controls_cache: dict[str, list[dict[str, Any]]],
) -> None:
    """Inject custom frameworks into the compliance engine's in-memory caches.

    Called from ``compliance_engine._refresh_cache_if_needed`` after the
    Firestore fetch so custom frameworks are treated identically by every
    downstream path (list, evaluate, baseline, evidence).
    """
    for fw in list_frameworks():
        fw_dict = fw.to_framework_dict()
        # Replace if a framework with the same id already exists (case-insensitive)
        existing_idx = next(
            (
                i
                for i, existing in enumerate(frameworks_cache)
                if existing.get("framework_id", "").upper() == fw.framework_id.upper()
            ),
            None,
        )
        if existing_idx is not None:
            frameworks_cache[existing_idx] = fw_dict
        else:
            frameworks_cache.append(fw_dict)
        controls_cache[fw.framework_id] = fw.to_control_dicts()
        # Also index under uppercase (engine uses upper-case keys in some paths)
        controls_cache[fw.framework_id.upper()] = fw.to_control_dicts()
