from __future__ import annotations

import hashlib, json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

MODALITIES = {
    "featurescript": ("featurescript_rp", ".txt"),
    "step": ("step_abc", ".step"),
    "stl": ("stl_abc", ".stl"),
    "image": ("multiview_images_abc", ".png"),
    "annotation": ("text_annotations", ".txt"),
}


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


@dataclass
class Sample:
    sample_id: str
    files: dict[str, str] = field(default_factory=dict)
    hashes: dict[str, str] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)
    diagnostics: list[dict[str, Any]] = field(default_factory=list)

    def as_dict(self) -> dict[str, Any]: return asdict(self)


def _files(root: Path, directory: str, suffix: str) -> dict[str, Path]:
    base = root / directory; result: dict[str, Path] = {}
    if not base.exists(): return result
    for path in base.rglob(f"*{suffix}"):
        if path.is_file() and path.stem.isdigit(): result[path.stem] = path
    return result


def _jsonl_records(path: Path) -> list[dict[str, Any]]:
    if not path.exists(): return []
    records = []
    for index, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        try:
            value = json.loads(line); value["_line"] = index; records.append(value)
        except json.JSONDecodeError:
            records.append({"_line": index, "_invalid": True})
    return records


def _assistant_script(record: dict[str, Any]) -> str | None:
    for message in record.get("messages") or []:
        if message.get("role") == "assistant" and isinstance(message.get("content"), str) and "FeatureScript" in message["content"]:
            return message["content"]
    return None


def scan_dataset(root: Path, *, include_hashes: bool = True) -> list[Sample]:
    if not root.is_dir(): raise FileNotFoundError(f"CADFS input directory does not exist: {root}")
    by_modality = {name: _files(root, directory, suffix) for name, (directory, suffix) in MODALITIES.items()}
    ids = sorted(set().union(*(set(values) for values in by_modality.values())))
    text_records = _jsonl_records(root / "CADFS_text_test.jsonl")
    image_records = _jsonl_records(root / "CADFS_image_test.jsonl")
    text_by_hash = {hashlib.sha256(script.encode()).hexdigest(): record for record in text_records if (script := _assistant_script(record))}
    samples: list[Sample] = []
    for index, sample_id in enumerate(ids):
        sample = Sample(sample_id)
        for name, values in by_modality.items():
            path = values.get(sample_id)
            if path is None: sample.diagnostics.append({"code": "missing_modality", "modality": name}); continue
            sample.files[name] = str(path)
            if include_hashes: sample.hashes[name] = sha256(path)
        fs_path = by_modality["featurescript"].get(sample_id)
        if fs_path:
            script_hash = hashlib.sha256(fs_path.read_text(encoding="utf-8").encode()).hexdigest()
            record = text_by_hash.get(script_hash)
            if record:
                sample.metadata["text_jsonl"] = {"line": record.get("_line"), "cad_file_id": record.get("cad_file_id")}
            elif index < len(text_records):
                fallback = text_records[index]; sample.metadata["text_jsonl"] = {"line": fallback.get("_line"), "cad_file_id": fallback.get("cad_file_id"), "alignment": "fallback"}
                sample.diagnostics.append({"code": "alignment_fallback", "modality": "text_jsonl"})
        if index < len(image_records):
            record = image_records[index]
            sample.metadata["image_jsonl"] = {"line": record.get("_line"), "cad_file_id": record.get("cad_file_id")}
        samples.append(sample)
    return samples