/storage/materialize.py

from __future__ import annotations

import hashlib
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping

from storage.models import AssetStatus, CompileReport, MaterializationStatus
from storage.repo import StorageRepo
from storage.revisions import (
    active_cost_path,
    active_inputs_dir,
    active_model_path,
    active_prompt_path,
    active_provenance_path,
    active_traces_dir,
)


def build_materialization_fingerprint(
    *,
    model_py_sha256: str | None,
    model_urdf_sha256: str | None = None,
    sdk_fingerprint: str | None = None,
    materializer_version: str = "v1",
) -> str:
    payload = "|".join(
        [
            model_py_sha256 or "",
            model_urdf_sha256 or "",
            sdk_fingerprint or "",
            materializer_version,
        ]
    )
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def build_compile_fingerprint_inputs(
    *,
    model_path: Path,
    sdk_fingerprint: str | None = None,
) -> dict[str, str | None]:
    return {
        "model_py_sha256": sha256_file(model_path),
        "sdk_fingerprint": sdk_fingerprint,
    }


def build_compile_fingerprint_from_inputs(
    inputs: Mapping[str, Any],
    *,
    materializer_version: str = "v1",
) -> str:
    model_py_sha256 = inputs.get("model_py_sha256")
    sdk_fingerprint = inputs.get("sdk_fingerprint")
    return build_materialization_fingerprint(
        model_py_sha256=str(model_py_sha256) if isinstance(model_py_sha256, str) else None,
        sdk_fingerprint=str(sdk_fingerprint) if isinstance(sdk_fingerprint, str) else None,
        materializer_version=materializer_version,
    )


def build_model_source_snapshot(*, model_path: Path) -> dict[str, int]:
    stat = model_path.stat()
    return {
        "model_py_mtime_ns": int(stat.st_mtime_ns),
        "model_py_size_bytes": int(stat.st_size),
    }


def _compile_report_metrics(report: object) -> dict[str, Any] | None:
    if not isinstance(report, dict):
        return None
    metrics = report.get("metrics")
    return metrics if isinstance(metrics, dict) else None


def compile_report_matches_model_source_snapshot(
    report: object,
    *,
    model_path: Path,
) -> bool | None:
    metrics = _compile_report_metrics(report)
    if metrics is None:
        return None
    expected_mtime_ns = metrics.get("model_py_mtime_ns")
    expected_size_bytes = metrics.get("model_py_size_bytes")
    if not isinstance(expected_mtime_ns, int) or not isinstance(expected_size_bytes, int):
        return None
    try:
        stat = model_path.stat()
    except OSError:
        return None
    return int(stat.st_mtime_ns) == expected_mtime_ns and int(stat.st_size) == expected_size_bytes


def compile_report_visual_mesh_footprint(report: object) -> tuple[int | None, int | None]:
    metrics = _compile_report_metrics(report)
    if metrics is None:
        return None, None
    mesh_bytes = metrics.get("visual_mesh_bytes")
    mesh_files = metrics.get("visual_mesh_file_count")
    if not isinstance(mesh_bytes, int) or mesh_bytes < 0:
        mesh_bytes = None
    if not isinstance(mesh_files, int) or mesh_files < 0:
        mesh_files = None
    return mesh_bytes, mesh_files


def compile_report_elapsed_seconds(report: object) -> float | None:
    metrics = _compile_report_metrics(report)
    if metrics is None:
        return None
    value = metrics.get("compile_elapsed_seconds")
    if isinstance(value, (int, float)) and value >= 0:
        return float(value)
    return None


@dataclass(slots=True)
class MaterializationStore:
    repo: StorageRepo

    def ensure_record_dirs(self, record_id: str) -> Path:
        materialization_dir = self.repo.layout.record_materialization_dir(record_id)
        materialization_dir.mkdir(parents=True, exist_ok=True)
        return materialization_dir

    def write_compile_report(self, record_id: str, report: CompileReport) -> Path:
        path = self.repo.layout.record_materialization_compile_report_path(record_id)
        self.repo.write_json(path, report.to_dict())
        return path

    def asset_status(self, record_id: str) -> AssetStatus:
        assets_dir = self.repo.layout.record_materialization_assets_dir(record_id)
        return AssetStatus(
            record_id=record_id,
            assets_dir=assets_dir,
            meshes_present=self.repo.layout.record_materialization_asset_meshes_dir(
                record_id
            ).exists(),
            glb_present=self.repo.layout.record_materialization_asset_glb_dir(record_id).exists(),
            viewer_present=self.repo.layout.record_materialization_asset_viewer_dir(
                record_id
            ).exists(),
        )


def canonical_record_paths(repo: StorageRepo, record_id: str) -> dict[str, Path]:
    return {
        "record_json": repo.layout.record_metadata_path(record_id),
        "prompt_txt": active_prompt_path(repo, record_id),
        "model_py": active_model_path(repo, record_id),
        "provenance_json": active_provenance_path(repo, record_id),
        "cost_json": active_cost_path(repo, record_id),
        "inputs_dir": active_inputs_dir(repo, record_id),
        "traces_dir": active_traces_dir(repo, record_id),
    }


def materialization_paths(repo: StorageRepo, record_id: str) -> dict[str, Path]:
    return {
        "root": repo.layout.record_materialization_dir(record_id),
        "model_urdf": repo.layout.record_materialization_urdf_path(record_id),
        "compile_report_json": repo.layout.record_materialization_compile_report_path(record_id),
        "assets_dir": repo.layout.record_materialization_assets_dir(record_id),
        "meshes_dir": repo.layout.record_materialization_asset_meshes_dir(record_id),
        "glb_dir": repo.layout.record_materialization_asset_glb_dir(record_id),
        "viewer_dir": repo.layout.record_materialization_asset_viewer_dir(record_id),
    }


def ensure_record_artifacts_exist(
    repo: StorageRepo,
    record_id: str,
    *,
    required: tuple[str, ...],
) -> None:
    paths = canonical_record_paths(repo, record_id)
    missing = [name for name in required if not paths[name].exists()]
    if not missing:
        return

    missing_labels = ", ".join(missing)
    raise FileNotFoundError(
        f"Record {record_id} is missing canonical artifact(s): {missing_labels}"
    )


def _has_nonempty_dir(path: Path) -> bool:
    return path.exists() and path.is_dir() and any(path.iterdir())


def summarize_visual_mesh_footprint(mesh_root: Path) -> tuple[int, int]:
    if not mesh_root.exists() or not mesh_root.is_dir():
        return 0, 0

    total_bytes = 0
    mesh_files = 0
    for dirpath, dirnames, filenames in os.walk(mesh_root):
        rel_parts = Path(dirpath).relative_to(mesh_root).parts
        if rel_parts and rel_parts[0] == "collision":
            dirnames[:] = []
            continue
        for filename in filenames:
            if not filename.lower().endswith(".obj"):
                continue
            path = Path(dirpath) / filename
            try:
                total_bytes += path.stat().st_size
                mesh_files += 1
            except OSError:
                continue
    return total_bytes, mesh_files


def build_materialization_summary(repo: StorageRepo, record_id: str) -> dict[str, Any]:
    paths = materialization_paths(repo, record_id)
    mesh_bytes, mesh_file_count = summarize_visual_mesh_footprint(paths["meshes_dir"])

    has_materialized_assets = mesh_file_count > 0
    if not has_materialized_assets:
        has_materialized_assets = _has_nonempty_dir(paths["glb_dir"]) or _has_nonempty_dir(
            paths["viewer_dir"]
        )

    if has_materialized_assets:
        materialization_status: MaterializationStatus = "available"
    elif paths["model_urdf"].exists() and not urdf_references_external_meshes(paths["model_urdf"]):
        materialization_status = "available"
    else:
        materialization_status = "missing"

    return {
        "materialization_status": materialization_status,
        "visual_mesh_bytes": mesh_bytes,
        "visual_mesh_file_count": mesh_file_count,
    }


def record_has_materialized_assets(repo: StorageRepo, record_id: str) -> bool:
    return any(
        _has_nonempty_dir(path)
        for path in (
            repo.layout.record_materialization_asset_meshes_dir(record_id),
            repo.layout.record_materialization_asset_glb_dir(record_id),
            repo.layout.record_materialization_asset_viewer_dir(record_id),
        )
    )


def urdf_references_external_meshes(urdf_path: Path) -> bool:
    if not urdf_path.exists() or not urdf_path.is_file():
        return False
    try:
        text = urdf_path.read_text(encoding="utf-8", errors="ignore")
    except OSError:
        return False
    lowered = text.lower()
    return "<mesh" in lowered and "filename=" in lowered


def infer_materialization_status(
    repo: StorageRepo,
    record_id: str,
) -> MaterializationStatus:
    if record_has_materialized_assets(repo, record_id):
        return "available"

    urdf_path = repo.layout.record_materialization_urdf_path(record_id)
    if urdf_path.exists() and not urdf_references_external_meshes(urdf_path):
        return "available"
    return "missing"