"""One partition contract at scheduling, support and journal boundaries."""

from dataclasses import replace

import pytest

from linkerhand_calibration.core.domain.capture_plan import (
    ADAPTIVE, CapturePlan, configure_training, evidence_digest, select_task_training)
from linkerhand_calibration.core.domain.sampling import command_nodes
from linkerhand_calibration.profiles import load_bundled_hand_profile
from linkerhand_calibration.runtime.engine import CalibrationEngine
from linkerhand_calibration.runtime.execution import SessionExecution
from linkerhand_calibration.runtime.session import CalibrationPhase
from linkerhand_calibration.runtime.task_quality import evaluate_task_input_support
from linkerhand_calibration.runtime.training import resolve_capture_plan


@pytest.fixture
def profile():
    return configure_training(load_bundled_hand_profile("o30_right_18"), ADAPTIVE)


def decision(profile, task, cycles, action):
    from linkerhand_calibration.core.fitting.training_quality import TRAINING_DECISION_VERSION
    row = dict(kind="training_decision", version=TRAINING_DECISION_VERSION,
        task_name=task.key, assessed_cycles=list(cycles), decision=action,
        input_plan_sha256=CapturePlan.from_profile(profile).sha256, failures=["noisy"] if action == "fail" else [])
    return {**row, "decision_sha256": evidence_digest(row)}


def header(profile):
    return dict(kind="session_start", capture_plan=CapturePlan.from_profile(profile).as_dict())


def support_rows(profile, task, cycles):
    return [dict(kind="joint_sample", task_name=task.key, joint=task.joints[0],
        view=profile.measurement.measurements[task.joints[0]].view, sample_phase="steady",
        cycle=cycle, direction=direction, command_u8=value, image_stamp_ns=index * 10 + node + 1)
        for index, (cycle, direction) in enumerate((c, d) for c in cycles for d in ("increasing", "decreasing"))
        for node, value in enumerate((0., 127., 255.))]


def test_optional_round_keeps_holdout_identity_grid_and_motion_segments(profile):
    units = CalibrationEngine(profile).scan_units()
    for task in profile.motion.tasks:
        original_grid = command_nodes(profile, task, holdout=True)
        profile = select_task_training(profile, task.key, (0, 1))
        assert command_nodes(profile, task, holdout=True) == original_grid
    shortened = CalibrationEngine(profile).scan_units()
    assert len(units) == 144 and len(shortened) == 108
    assert shortened == tuple(unit for unit in units if unit.cycle != 2)
    assert all(task.holdout == 3 for task in CapturePlan.from_profile(profile).tasks)


def test_secondary_camera_uses_its_primary_joint_partition():
    from linkerhand_calibration.core.domain.capture_plan import task_for_joint, training_cycles
    profile = load_bundled_hand_profile("g20_right_19")
    for primary, secondary in profile.measurement.cross_view_sources.items():
        assert task_for_joint(profile, secondary) == task_for_joint(profile, primary)
        assert training_cycles(profile, joint=secondary) == (0, 1, 2)


@pytest.mark.parametrize("mode,cycles", [("interleaved", (0, 1, 2, 3)), ("separate", (4, 5))])
def test_command_support_uses_the_correct_phase_partition(profile, mode, cycles):
    profile = replace(profile, acquisition=replace(profile.acquisition, command_capture_mode=mode))
    task = profile.motion.tasks[0]
    rows = support_rows(profile, task, cycles)
    assert evaluate_task_input_support(profile, task, rows).passed
    # Validation endpoints cannot extend training support.
    reduced = [row for row in rows if row["cycle"] == cycles[-1] or row["command_u8"] < 255.]
    assert not evaluate_task_input_support(profile, task, reduced).passed


@pytest.mark.parametrize("action,expected", [("freeze", 3), ("add_training", 2), ("fail", 2)])
def test_session_advances_only_after_bounded_training_decision(profile, monkeypatch, action, expected):
    from linkerhand_calibration.runtime.engine import SweepQuality
    driver = SessionExecution(profile)
    task = profile.motion.tasks[0]
    round_cycle = 2 if action == "fail" else 1
    index = max(i for i, u in enumerate(driver.session._units) if u.task_key == task.key and u.cycle == round_cycle)
    driver.session._unit_index = index
    driver.session.phase = CalibrationPhase.EVALUATE
    driver.completed_units.update(u.identity for u in driver.session._units[:index])
    monkeypatch.setattr("linkerhand_calibration.runtime.execution.evaluate_capture_unit",
        lambda *a, **k: SweepQuality(True, (), (), {}))
    driver.evaluate([], training_decision=decision(profile, task, range(round_cycle + 1), action))
    if action == "fail":
        assert driver.session.phase == CalibrationPhase.FAILED
    else:
        assert driver.session.current_unit.cycle == expected
        assert driver.session.status().task.cycle_count == (3 if action == "freeze" else 4)
        if action == "freeze":
            assert driver.session.status().task.cycle == 3


def test_journal_rejects_future_holdout_and_changed_decision(profile):
    task = profile.motion.tasks[0]
    row = decision(profile, task, (0, 1), "freeze")
    current = resolve_capture_plan(profile, [header(profile), row])
    assert current.quality.task_training_cycles[task.key] == (0, 1)
    with pytest.raises(ValueError, match="holdout_before_training_freeze"):
        resolve_capture_plan(profile, [header(profile), support_rows(profile, task, (3,))[0], row])
    with pytest.raises(ValueError, match="training_decision_invalid"):
        resolve_capture_plan(profile, [header(profile), {**row, "assessed_cycles": [0, 1, 2]}])
    with pytest.raises(ValueError, match="training_observation_after_freeze"):
        resolve_capture_plan(profile, [header(profile), row, support_rows(profile, task, (0,))[0]])
    with pytest.raises(ValueError, match="matching_plan_header"):
        resolve_capture_plan(profile, [])
    with pytest.raises(ValueError, match="third_training_without_decision"):
        resolve_capture_plan(profile, [header(profile), support_rows(profile, task, (2,))[0]])


def test_fixed_and_adaptive_checkpoints_are_isolated(profile):
    fixed = configure_training(profile, "fixed")
    engine = CalibrationEngine(profile)
    from linkerhand_calibration.core.geometry.pnp import POSE_TRACKING_POLICY_VERSION
    from linkerhand_calibration.runtime.engine import ACQUISITION_POLICY_VERSION
    row = {**header(profile), "profile_id": profile.key.profile_id,
        "acquisition_policy_version": ACQUISITION_POLICY_VERSION,
        "pose_tracking_policy_version": POSE_TRACKING_POLICY_VERSION,
        "capture_schedule_version": engine.capture_schedule_version,
        "calibration_scope": None}
    from linkerhand_calibration.core.urdf.partial_scope import profile_scope
    row["calibration_scope"] = profile_scope(profile)
    assert engine.resume_compatible(row)
    assert not CalibrationEngine(fixed).resume_compatible(row)
    row["capture_schedule_version"] = CalibrationEngine(fixed).capture_schedule_version
    assert not engine.resume_compatible(row)


def test_stage_timing_accounts_for_elapsed_time_once():
    from linkerhand_calibration.runtime.timing_diagnostics import StageTiming
    timing = StageTiming(0.)
    timing.mark_start(1.)
    timing.enter("geometry_solving", 2.)
    timing.enter("geometry_solving", 5.)
    timing.enter("steady_sampling", 6.)
    report = timing.as_dict(7.)
    assert report["stage_seconds"] == {"startup": 2., "geometry_solving": 4., "steady_sampling": 1.}
    assert sum(report["stage_seconds"].values()) == report["elapsed_seconds"]
    timing.enter('terminal', 9.)
    assert timing.as_dict(20.)['start_to_terminal_seconds'] == 8.


def test_adaptive_resume_restarts_suffix_after_incomplete_dependency(profile):
    from linkerhand_calibration.runtime.joint_resume import JointResume
    tasks = profile.motion.tasks[:3]
    for task in tasks:
        profile = select_task_training(profile, task.key, (0, 1))
    engine = CalibrationEngine(profile)
    resume = JointResume(profile)
    resume.units = {u.identity for u in engine.scan_units() if u.task_key in {t.key for t in tasks}}
    resume.units.remove(next(u.identity for u in engine.scan_units() if u.task_key == tasks[1].key))
    resume.restrict_adaptive_prefix(engine)
    assert {key[0] for key in resume.units} == {tasks[0].key}
    assert resume.profile.quality.task_training_cycles == {tasks[0].key: (0, 1)}