import logging
from datetime import UTC, datetime
from pathlib import Path
from typing import TypedDict

from ruamel.yaml import YAML
from ruamel.yaml.scalarstring import LiteralScalarString

from mitre_emb3d._graph import MITREGraph
from mitre_emb3d._models import PropertyId

from ._models import PropertyMapperOutput

_LOGGER = logging.getLogger(__name__)


class EvidenceEntry(TypedDict):
    file_name: str
    code_snippet: str


class PropertyArtifactDocument(TypedDict):
    ai_model: str
    head_commit: str
    updated_at: datetime
    property_id: str
    property_name: str
    category: str
    is_applicable: bool
    evidence: list[EvidenceEntry]


def _build_document(
    result: PropertyMapperOutput,
    mitre_graph: MITREGraph,
    head_commit: str,
    ai_model: str,
) -> PropertyArtifactDocument:
    prop = mitre_graph.get_property_from_id(result.property_id)

    evidence_entries = []
    for ev in result.evidence:
        evidence_entries.append(
            EvidenceEntry(
                file_name=ev.file_name,
                code_snippet=LiteralScalarString(ev.code_snippet),
            )
        )

    return PropertyArtifactDocument(
        ai_model=ai_model,
        head_commit=head_commit,
        updated_at=datetime.now(UTC),
        property_id=result.property_id,
        property_name=prop.name,
        category=prop.category.value,
        is_applicable=result.is_relevant,
        evidence=evidence_entries,
    )


def write_property_documents(
    results: dict[PropertyId, PropertyMapperOutput],
    mitre_graph: MITREGraph,
    output_dir: Path,
    head_commit: str,
    ai_model: str,
) -> None:
    properties_dir = output_dir / "properties"
    properties_dir.mkdir(parents=True, exist_ok=True)

    yaml = YAML()
    yaml.default_flow_style = False

    for pid, result in results.items():
        doc = _build_document(result, mitre_graph, head_commit, ai_model)
        out_path = properties_dir / f"{pid}.yaml"

        with out_path.open("w") as f:
            yaml.dump(doc, f)

        _LOGGER.info(f"Wrote {out_path}")


def read_property_documents(output_dir: Path) -> dict[PropertyId, PropertyArtifactDocument]:
    properties_dir = output_dir / "properties"
    if not properties_dir.exists():
        _LOGGER.warning(f"Properties directory {properties_dir} does not exist.")
        return {}

    yaml = YAML()
    documents: dict[PropertyId, PropertyArtifactDocument] = {}

    for prop_file in properties_dir.glob("*.yaml"):
        with prop_file.open() as f:
            doc = yaml.load(f)
            pid = doc["property_id"]
            documents[pid] = doc

    return documents
