"""Thin OID facade over the canonical CCD dual-clock collector."""
from __future__ import annotations

from dataclasses import dataclass
import json
from pathlib import Path
from typing import Any

import numpy as np

from CCD_analysis.utils.dual_clock import DualClockCollector, field_steps_from_interval
from CCD_analysis.utils.run_contract import atomic_publish_json
from .contracts import CANONICAL_ROLES, exact_branch_gates


@dataclass(frozen=True)
class CollectionClock:
    control_interval: int
    control_count: int
    field_interval: int

    @property
    def field_steps(self) -> tuple[int, ...]:
        return field_steps_from_interval(self.control_interval * self.control_count, self.field_interval)


def telemetry_snapshot(obs_slice: tuple[int, int]):
    """Build a read-only callback sampled immediately after each CFD segment."""
    start, stop = obs_slice

    def snapshot(flow_field: Any, lattice_step: int) -> dict[str, np.ndarray]:
        if not hasattr(flow_field, "current_step_observation"):
            raise TypeError("field-level telemetry requires current_step_observation")
        raw = np.asarray(flow_field.current_step_observation())[start:stop]
        if raw.size != 12:
            raise ValueError("canonical paired telemetry requires six sensor and six force channels")
        return {"sensors": raw[:6].astype(np.float32, copy=False),
                "forces": raw[6:12].astype(np.float32, copy=False)}

    return snapshot


def make_collector(flow_field, clock: CollectionClock, *, u0: float,
                   obs_slice: tuple[int, int] | None = None) -> DualClockCollector:
    callback = None if obs_slice is None else telemetry_snapshot(obs_slice)
    return DualClockCollector(flow_field, control_interval=clock.control_interval,
                              control_count=clock.control_count, u0=u0,
                              field_steps=clock.field_steps, snapshot_callback=callback)


def field_level_telemetry(collector: DualClockCollector, *, sample_id_start: int = 0,
                          lattice_dt: float = 1.0) -> dict[str, np.ndarray]:
    """Return canonical telemetry bound to the collector's exact field lattice."""
    if lattice_dt <= 0 or not np.isfinite(lattice_dt):
        raise ValueError("lattice_dt must be positive and finite")
    steps = np.asarray(collector.saved_steps, dtype=np.int64)
    if set(collector.snapshots) != {"sensors", "forces"}:
        raise ValueError("collector does not contain canonical field-level telemetry")
    if any(len(values) != len(steps) for values in collector.snapshots.values()):
        raise ValueError("field and telemetry sample counts differ")
    return {"sensors": np.stack(collector.snapshots["sensors"]),
            "forces": np.stack(collector.snapshots["forces"]),
            "lattice_steps": steps,
            "sample_ids": np.arange(sample_id_start, sample_id_start + len(steps), dtype=np.int64),
            "times": steps.astype(np.float64) * lattice_dt}


def publish_paired_manifest(branch_dirs: dict[str, str | Path], output: str | Path) -> dict[str, Any]:
    """Gate and publish three completed branch manifests as one paired collection."""
    if set(branch_dirs) != set(CANONICAL_ROLES):
        raise ValueError("branch directories must use exactly the canonical roles")
    manifests = {}
    for role in CANONICAL_ROLES:
        directory = Path(branch_dirs[role])
        if not (directory / "_SUCCESS").is_file():
            raise ValueError(f"branch {role} is incomplete")
        manifests[role] = json.loads((directory / "branch.json").read_text(encoding="utf-8"))
    paired = exact_branch_gates(manifests)
    atomic_publish_json(output, paired)
    return paired


def zero_action(n_objects: int) -> np.ndarray:
    if n_objects < 1:
        raise ValueError("n_objects must be positive")
    return np.zeros(n_objects, dtype=np.float32)
