storage.py

from __future__ import annotations

import hashlib
import json
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from schemas import CodeEvalSpec, ModelEntry, PromptArchive, PromptEntry


def ensure_directory(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def slugify(value: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9]+", "_", value.strip().lower())
    slug = slug.strip("_")
    return slug or "prompt"


def load_models(models_file: Path) -> list[ModelEntry]:
    if not models_file.exists():
        raise FileNotFoundError(f"Missing models file: {models_file}")

    payload: Any = json.loads(models_file.read_text(encoding="utf-8"))
    if isinstance(payload, dict):
        raw_models = payload.get("models", [])
    else:
        raw_models = payload

    if not isinstance(raw_models, list):
        raise ValueError("models.json must contain a list or a {'models': [...]} object.")

    return [ModelEntry.model_validate(item) for item in raw_models]


def save_models(models_file: Path, models: list[ModelEntry]) -> None:
    if models_file.parent and not models_file.parent.exists():
        models_file.parent.mkdir(parents=True, exist_ok=True)

    payload = {"models": [model.model_dump(mode="json") for model in models]}
    models_file.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def discover_prompts(prompts_dir: Path) -> list[PromptEntry]:
    if not prompts_dir.exists():
        return []

    prompts: list[PromptEntry] = []
    for path in sorted(prompts_dir.glob("*.md")):
        content = path.read_text(encoding="utf-8")
        checksum = hashlib.sha256(content.encode("utf-8")).hexdigest()
        spec_path = path.with_suffix(".spec.json")
        eval_spec = None
        eval_spec_path: str | None = None
        if spec_path.exists():
            try:
                spec_payload: Any = json.loads(spec_path.read_text(encoding="utf-8"))
                eval_spec = CodeEvalSpec.model_validate(spec_payload)
                eval_spec_path = str(spec_path)
            except Exception as exc:
                raise ValueError(f"Invalid prompt spec file: {spec_path} ({exc})") from exc
        prompts.append(
            PromptEntry(
                name=path.stem,
                path=str(path),
                content=content,
                sha256=checksum,
                eval_spec_path=eval_spec_path,
                eval_spec=eval_spec,
            )
        )
    return prompts


def archive_filename(prompt_name: str, timestamp: datetime | None = None) -> str:
    ts = timestamp or datetime.now(timezone.utc)
    return f"{slugify(prompt_name)}_{ts.strftime('%Y%m%d_%H%M%S')}.json"


def save_archive(results_dir: Path, archive: PromptArchive, timestamp: datetime | None = None) -> Path:
    ensure_directory(results_dir)
    path = results_dir / archive_filename(archive.prompt.name, timestamp=timestamp)
    payload = archive.model_dump(mode="json")
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    return path


def load_archive(path: Path) -> PromptArchive:
    payload = json.loads(path.read_text(encoding="utf-8"))
    return PromptArchive.model_validate(payload)


def load_latest_archive(results_dir: Path) -> PromptArchive | None:
    if not results_dir.exists():
        return None

    candidates = sorted(
        (p for p in results_dir.glob("*.json") if p.is_file()),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    for candidate in candidates:
        try:
            return load_archive(candidate)
        except Exception:
            continue
    return None