"""Dataset generation orchestrator with sharding, checkpointing, and statistics.

Provides :class:`DatasetConfig` for configuration and :class:`DatasetGenerator`
for parallel generation of synthetic assembly training data.
"""

from __future__ import annotations

import hashlib
import json
import logging
import math
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from typing import Any

__all__ = [
    "DatasetConfig",
    "DatasetGenerator",
    "parse_simple_yaml",
]

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


@dataclass
class DatasetConfig:
    """Configuration for synthetic dataset generation."""

    name: str = "synthetic"
    num_assemblies: int = 100_000
    output_dir: str = "data/synthetic"
    shard_size: int = 1000
    complexity_distribution: dict[str, float] = field(
        default_factory=lambda: {"simple": 0.4, "medium": 0.4, "complex": 0.2}
    )
    body_count_min: int = 2
    body_count_max: int = 50
    templates: list[str] = field(default_factory=lambda: ["chain", "tree", "loop", "star", "mixed"])
    grounded_ratio: float = 0.5
    seed: int = 42
    num_workers: int = 4
    checkpoint_every: int = 5
    resume: bool = True

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> DatasetConfig:
        """Construct from a parsed config dict (e.g. YAML or OmegaConf).

        Handles both flat keys (``body_count_min``) and nested forms
        (``body_count: {min: 2, max: 50}``).
        """
        kw: dict[str, Any] = {}
        for key in (
            "name",
            "num_assemblies",
            "output_dir",
            "shard_size",
            "grounded_ratio",
            "seed",
            "num_workers",
            "checkpoint_every",
            "resume",
        ):
            if key in d:
                kw[key] = d[key]

        # Handle nested body_count dict
        if "body_count" in d and isinstance(d["body_count"], dict):
            bc = d["body_count"]
            if "min" in bc:
                kw["body_count_min"] = int(bc["min"])
            if "max" in bc:
                kw["body_count_max"] = int(bc["max"])
        else:
            if "body_count_min" in d:
                kw["body_count_min"] = int(d["body_count_min"])
            if "body_count_max" in d:
                kw["body_count_max"] = int(d["body_count_max"])

        if "complexity_distribution" in d:
            cd = d["complexity_distribution"]
            if isinstance(cd, dict):
                kw["complexity_distribution"] = {str(k): float(v) for k, v in cd.items()}
        if "templates" in d and isinstance(d["templates"], list):
            kw["templates"] = [str(t) for t in d["templates"]]

        return cls(**kw)


# ---------------------------------------------------------------------------
# Shard specification / result
# ---------------------------------------------------------------------------


@dataclass
class ShardSpec:
    """Specification for generating a single shard."""

    shard_id: int
    start_example_id: int
    count: int
    seed: int
    complexity_distribution: dict[str, float]
    body_count_min: int
    body_count_max: int
    grounded_ratio: float


@dataclass
class ShardResult:
    """Result returned from a shard worker."""

    shard_id: int
    num_examples: int
    file_path: str
    generation_time_s: float


# ---------------------------------------------------------------------------
# Seed derivation
# ---------------------------------------------------------------------------


def _derive_shard_seed(global_seed: int, shard_id: int) -> int:
    """Derive a deterministic per-shard seed from the global seed."""
    h = hashlib.sha256(f"{global_seed}:{shard_id}".encode()).hexdigest()
    return int(h[:8], 16)


# ---------------------------------------------------------------------------
# Progress display
# ---------------------------------------------------------------------------


class _PrintProgress:
    """Fallback progress display when tqdm is unavailable."""

    def __init__(self, total: int) -> None:
        self.total = total
        self.current = 0
        self.start_time = time.monotonic()

    def update(self, n: int = 1) -> None:
        self.current += n
        elapsed = time.monotonic() - self.start_time
        rate = self.current / elapsed if elapsed > 0 else 0.0
        eta = (self.total - self.current) / rate if rate > 0 else 0.0
        pct = 100.0 * self.current / self.total
        sys.stdout.write(
            f"\r[{pct:5.1f}%] {self.current}/{self.total} shards"
            f" | {rate:.1f} shards/s | ETA: {eta:.0f}s"
        )
        sys.stdout.flush()

    def close(self) -> None:
        sys.stdout.write("\n")
        sys.stdout.flush()


def _make_progress(total: int) -> _PrintProgress:
    """Create a progress tracker (tqdm if available, else print-based)."""
    try:
        from tqdm import tqdm  # type: ignore[import-untyped]

        return tqdm(total=total, desc="Generating shards", unit="shard")  # type: ignore[no-any-return,return-value]
    except ImportError:
        return _PrintProgress(total)


# ---------------------------------------------------------------------------
# Shard I/O
# ---------------------------------------------------------------------------


def _save_shard(
    shard_id: int,
    examples: list[dict[str, Any]],
    shards_dir: Path,
) -> Path:
    """Save a shard to disk (.pt if torch available, else .json)."""
    shards_dir.mkdir(parents=True, exist_ok=True)
    try:
        import torch  # type: ignore[import-untyped]

        path = shards_dir / f"shard_{shard_id:05d}.pt"
        torch.save(examples, path)
    except ImportError:
        path = shards_dir / f"shard_{shard_id:05d}.json"
        with open(path, "w") as f:
            json.dump(examples, f)
    return path


def _load_shard(path: Path) -> list[dict[str, Any]]:
    """Load a shard from disk (.pt or .json)."""
    if path.suffix == ".pt":
        import torch  # type: ignore[import-untyped]

        result: list[dict[str, Any]] = torch.load(path, weights_only=False)
        return result
    with open(path) as f:
        result = json.load(f)
    return result


def _shard_format() -> str:
    """Return the shard file extension based on available libraries."""
    try:
        import torch  # type: ignore[import-untyped]  # noqa: F401

        return ".pt"
    except ImportError:
        return ".json"


# ---------------------------------------------------------------------------
# Shard worker (module-level for pickling)
# ---------------------------------------------------------------------------


def _generate_shard_worker(spec: ShardSpec, output_dir: str) -> ShardResult:
    """Generate a single shard — top-level function for ProcessPoolExecutor."""
    from solver.datagen.generator import SyntheticAssemblyGenerator

    t0 = time.monotonic()
    gen = SyntheticAssemblyGenerator(seed=spec.seed)
    rng = np.random.default_rng(spec.seed + 1)

    tiers = list(spec.complexity_distribution.keys())
    probs_list = [spec.complexity_distribution[t] for t in tiers]
    total = sum(probs_list)
    probs = [p / total for p in probs_list]

    examples: list[dict[str, Any]] = []
    for i in range(spec.count):
        tier_idx = int(rng.choice(len(tiers), p=probs))
        tier = tiers[tier_idx]
        try:
            batch = gen.generate_training_batch(
                batch_size=1,
                complexity_tier=tier,  # type: ignore[arg-type]
                grounded_ratio=spec.grounded_ratio,
            )
            ex = batch[0]
            ex["example_id"] = spec.start_example_id + i
            ex["complexity_tier"] = tier
            examples.append(ex)
        except Exception:
            logger.warning(
                "Shard %d, example %d failed — skipping",
                spec.shard_id,
                i,
                exc_info=True,
            )

    shards_dir = Path(output_dir) / "shards"
    path = _save_shard(spec.shard_id, examples, shards_dir)

    elapsed = time.monotonic() - t0
    return ShardResult(
        shard_id=spec.shard_id,
        num_examples=len(examples),
        file_path=str(path),
        generation_time_s=elapsed,
    )


# ---------------------------------------------------------------------------
# Minimal YAML parser
# ---------------------------------------------------------------------------


def _parse_scalar(value: str) -> int | float | bool | str:
    """Parse a YAML scalar value."""
    # Strip inline comments (space + #)
    if "  #" in value:
        value = value[: value.index("  #")].strip()
    elif "	#" in value:
        value = value[: value.index("	#")].strip()
    v = value.strip()
    if v.lower() in ("true", "yes"):
        return True
    if v.lower() in ("false", "no"):
        return False
    try:
        return int(v)
    except ValueError:
        pass
    try:
        return float(v)
    except ValueError:
        pass
    return v.strip("'\"")


def parse_simple_yaml(path: str) -> dict[str, Any]:
    """Parse a simple YAML file (flat scalars, one-level dicts, lists).

    This is **not** a full YAML parser.  It handles the structure of
    ``configs/dataset/synthetic.yaml``.
    """
    result: dict[str, Any] = {}
    current_key: str | None = None

    with open(path) as f:
        for raw_line in f:
            line = raw_line.rstrip()

            # Skip blank lines and full-line comments
            if not line or line.lstrip().startswith("#"):
                continue

            indent = len(line) - len(line.lstrip())

            if indent == 0 and ":" in line:
                key, _, value = line.partition(":")
                key = key.strip()
                value = value.strip()
                if value:
                    result[key] = _parse_scalar(value)
                    current_key = None
                else:
                    current_key = key
                    result[key] = {}
                continue

            if indent > 0 and line.lstrip().startswith("- "):
                item = line.lstrip()[2:].strip()
                if current_key is not None:
                    if isinstance(result.get(current_key), dict) and not result[current_key]:
                        result[current_key] = []
                    if isinstance(result.get(current_key), list):
                        result[current_key].append(_parse_scalar(item))
                continue

            if indent > 0 and ":" in line and current_key is not None:
                k, _, v = line.partition(":")
                k = k.strip()
                v = v.strip()
                if v:
                    # Strip inline comments
                    if "  #" in v:
                        v = v[: v.index("  #")].strip()
                    if not isinstance(result.get(current_key), dict):
                        result[current_key] = {}
                    result[current_key][k] = _parse_scalar(v)
                continue

    return result


# ---------------------------------------------------------------------------
# Dataset generator orchestrator
# ---------------------------------------------------------------------------


class DatasetGenerator:
    """Orchestrates parallel dataset generation with sharding and checkpointing."""

    def __init__(self, config: DatasetConfig) -> None:
        self.config = config
        self.output_path = Path(config.output_dir)
        self.shards_dir = self.output_path / "shards"
        self.checkpoint_file = self.output_path / ".checkpoint.json"
        self.index_file = self.output_path / "index.json"
        self.stats_file = self.output_path / "stats.json"

    # -- public API --

    def run(self) -> None:
        """Generate the full dataset."""
        self.output_path.mkdir(parents=True, exist_ok=True)
        self.shards_dir.mkdir(parents=True, exist_ok=True)

        shards = self._plan_shards()
        total_shards = len(shards)

        # Resume: find already-completed shards
        completed: set[int] = set()
        if self.config.resume:
            completed = self._find_completed_shards()

        pending = [s for s in shards if s.shard_id not in completed]

        if not pending:
            logger.info("All %d shards already complete.", total_shards)
        else:
            logger.info(
                "Generating %d shards (%d already complete).",
                len(pending),
                len(completed),
            )
            progress = _make_progress(len(pending))
            workers = max(1, self.config.num_workers)
            checkpoint_counter = 0

            with ProcessPoolExecutor(max_workers=workers) as pool:
                futures = {
                    pool.submit(
                        _generate_shard_worker,
                        spec,
                        str(self.output_path),
                    ): spec.shard_id
                    for spec in pending
                }
                for future in as_completed(futures):
                    shard_id = futures[future]
                    try:
                        result = future.result()
                        completed.add(result.shard_id)
                        logger.debug(
                            "Shard %d: %d examples in %.1fs",
                            result.shard_id,
                            result.num_examples,
                            result.generation_time_s,
                        )
                    except Exception:
                        logger.error("Shard %d failed", shard_id, exc_info=True)
                    progress.update(1)
                    checkpoint_counter += 1
                    if checkpoint_counter >= self.config.checkpoint_every:
                        self._update_checkpoint(completed, total_shards)
                        checkpoint_counter = 0

            progress.close()

        # Finalize
        self._build_index()
        stats = self._compute_statistics()
        self._write_statistics(stats)
        self._print_summary(stats)

        # Remove checkpoint (generation complete)
        if self.checkpoint_file.exists():
            self.checkpoint_file.unlink()

    # -- internal helpers --

    def _plan_shards(self) -> list[ShardSpec]:
        """Divide num_assemblies into shards."""
        n = self.config.num_assemblies
        size = self.config.shard_size
        num_shards = math.ceil(n / size)
        shards: list[ShardSpec] = []
        for i in range(num_shards):
            start = i * size
            count = min(size, n - start)
            shards.append(
                ShardSpec(
                    shard_id=i,
                    start_example_id=start,
                    count=count,
                    seed=_derive_shard_seed(self.config.seed, i),
                    complexity_distribution=dict(self.config.complexity_distribution),
                    body_count_min=self.config.body_count_min,
                    body_count_max=self.config.body_count_max,
                    grounded_ratio=self.config.grounded_ratio,
                )
            )
        return shards

    def _find_completed_shards(self) -> set[int]:
        """Scan shards directory for existing shard files."""
        completed: set[int] = set()
        if not self.shards_dir.exists():
            return completed

        for p in self.shards_dir.iterdir():
            if p.stem.startswith("shard_"):
                try:
                    shard_id = int(p.stem.split("_")[1])
                    # Verify file is non-empty
                    if p.stat().st_size > 0:
                        completed.add(shard_id)
                except (ValueError, IndexError):
                    pass
        return completed

    def _update_checkpoint(self, completed: set[int], total_shards: int) -> None:
        """Write checkpoint file."""
        data = {
            "completed_shards": sorted(completed),
            "total_shards": total_shards,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        }
        with open(self.checkpoint_file, "w") as f:
            json.dump(data, f)

    def _build_index(self) -> None:
        """Build index.json mapping shard files to assembly ID ranges."""
        shards_info: dict[str, dict[str, int]] = {}
        total_assemblies = 0

        for p in sorted(self.shards_dir.iterdir()):
            if not p.stem.startswith("shard_"):
                continue
            try:
                shard_id = int(p.stem.split("_")[1])
            except (ValueError, IndexError):
                continue
            examples = _load_shard(p)
            count = len(examples)
            start_id = shard_id * self.config.shard_size
            shards_info[p.name] = {"start_id": start_id, "count": count}
            total_assemblies += count

        fmt = _shard_format().lstrip(".")
        index = {
            "format_version": 1,
            "total_assemblies": total_assemblies,
            "total_shards": len(shards_info),
            "shard_format": fmt,
            "shards": shards_info,
        }
        with open(self.index_file, "w") as f:
            json.dump(index, f, indent=2)

    def _compute_statistics(self) -> dict[str, Any]:
        """Aggregate statistics across all shards."""
        classification_counts: dict[str, int] = {}
        body_count_hist: dict[int, int] = {}
        joint_type_counts: dict[str, int] = {}
        dof_values: list[int] = []
        degeneracy_values: list[int] = []
        rigid_count = 0
        minimally_rigid_count = 0
        total = 0

        for p in sorted(self.shards_dir.iterdir()):
            if not p.stem.startswith("shard_"):
                continue
            examples = _load_shard(p)
            for ex in examples:
                total += 1
                cls = str(ex.get("assembly_classification", "unknown"))
                classification_counts[cls] = classification_counts.get(cls, 0) + 1
                nb = int(ex.get("n_bodies", 0))
                body_count_hist[nb] = body_count_hist.get(nb, 0) + 1
                for j in ex.get("joints", []):
                    jt = str(j.get("type", "unknown"))
                    joint_type_counts[jt] = joint_type_counts.get(jt, 0) + 1
                dof_values.append(int(ex.get("internal_dof", 0)))
                degeneracy_values.append(int(ex.get("geometric_degeneracies", 0)))
                if ex.get("is_rigid"):
                    rigid_count += 1
                if ex.get("is_minimally_rigid"):
                    minimally_rigid_count += 1

        dof_arr = np.array(dof_values) if dof_values else np.zeros(1)
        deg_arr = np.array(degeneracy_values) if degeneracy_values else np.zeros(1)

        return {
            "total_examples": total,
            "classification_distribution": dict(sorted(classification_counts.items())),
            "body_count_histogram": dict(sorted(body_count_hist.items())),
            "joint_type_distribution": dict(sorted(joint_type_counts.items())),
            "dof_statistics": {
                "mean": float(dof_arr.mean()),
                "std": float(dof_arr.std()),
                "min": int(dof_arr.min()),
                "max": int(dof_arr.max()),
                "median": float(np.median(dof_arr)),
            },
            "geometric_degeneracy": {
                "assemblies_with_degeneracy": int(np.sum(deg_arr > 0)),
                "fraction_with_degeneracy": float(np.mean(deg_arr > 0)),
                "mean_degeneracies": float(deg_arr.mean()),
            },
            "rigidity": {
                "rigid_count": rigid_count,
                "rigid_fraction": (rigid_count / total if total > 0 else 0.0),
                "minimally_rigid_count": minimally_rigid_count,
                "minimally_rigid_fraction": (minimally_rigid_count / total if total > 0 else 0.0),
            },
        }

    def _write_statistics(self, stats: dict[str, Any]) -> None:
        """Write stats.json."""
        with open(self.stats_file, "w") as f:
            json.dump(stats, f, indent=2)

    def _print_summary(self, stats: dict[str, Any]) -> None:
        """Print a human-readable summary to stdout."""
        print("\n=== Dataset Generation Summary ===")
        print(f"Total examples: {stats['total_examples']}")
        print(f"Output directory: {self.output_path}")
        print()
        print("Classification distribution:")
        for cls, count in stats["classification_distribution"].items():
            frac = count / max(stats["total_examples"], 1) * 100
            print(f"  {cls}: {count} ({frac:.1f}%)")
        print()
        print("Joint type distribution:")
        for jt, count in stats["joint_type_distribution"].items():
            print(f"  {jt}: {count}")
        print()
        dof = stats["dof_statistics"]
        print(
            f"DOF: mean={dof['mean']:.1f}, std={dof['std']:.1f}, range=[{dof['min']}, {dof['max']}]"
        )
        rig = stats["rigidity"]
        print(
            f"Rigidity: {rig['rigid_count']}/{stats['total_examples']} "
            f"({rig['rigid_fraction'] * 100:.1f}%) rigid, "
            f"{rig['minimally_rigid_count']} minimally rigid"
        )
        deg = stats["geometric_degeneracy"]
        print(
            f"Degeneracy: {deg['assemblies_with_degeneracy']} assemblies "
            f"({deg['fraction_with_degeneracy'] * 100:.1f}%)"
        )
