"""Schema and serialization for AWS controls cache files."""
from __future__ import annotations

import json
from dataclasses import dataclass, asdict
from datetime import datetime
from typing import Literal


Severity = Literal["critical", "high", "medium", "low", "informational"]
ControlSource = Literal["fsbp", "cis", "aws-bp"]


@dataclass
class Control:
    control_id: str
    title: str
    severity: Severity
    resource_types: list[str]
    requirement: str
    source_url: str

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class ControlsFile:
    source: ControlSource
    fetched_at: datetime
    controls: list[Control]

    def _by_resource_type(self) -> dict[str, list[str]]:
        """Map terraform resource type -> list of control_ids.

        Full control records live in `controls`. Looking up details by ID
        from `controls` keeps the index small.
        """
        index: dict[str, list[str]] = {}
        for c in self.controls:
            for rt in c.resource_types:
                index.setdefault(rt, []).append(c.control_id)
        for rt in index:
            index[rt].sort()
        return dict(sorted(index.items()))

    def to_dict(self) -> dict:
        return {
            "source": self.source,
            "fetched_at": self.fetched_at.isoformat(),
            "controls": [c.to_dict() for c in self.controls],
            "by_resource_type": self._by_resource_type(),
        }

    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent, sort_keys=False)