import json
from pathlib import Path

import pytest

from reactor_sim import constants
from reactor_sim.commands import ReactorCommand
from reactor_sim.failures import HealthMonitor
from reactor_sim.reactor import Reactor
from reactor_sim.simulation import ReactorSimulation


def test_reactor_initial_state_is_cold():
    reactor = Reactor.default()
    state = reactor.initial_state()
    assert state.core.fuel_temperature == constants.ENVIRONMENT_TEMPERATURE
    assert state.primary_loop.mass_flow_rate == 0.0
    assert state.total_electrical_output() == 0.0


def test_state_save_and_load_roundtrip(tmp_path: Path):
    reactor = Reactor.default()
    reactor.control.manual_control = True
    sim = ReactorSimulation(reactor, timestep=5.0, duration=15.0)
    sim.log()
    save_path = tmp_path / "plant_state.json"
    assert sim.last_state is not None
    reactor.save_state(str(save_path), sim.last_state)
    assert save_path.exists()
    restored_reactor = Reactor.default()
    restored_state = restored_reactor.load_state(str(save_path))
    assert restored_state.core.fuel_temperature == pytest.approx(
        sim.last_state.core.fuel_temperature
    )
    assert restored_reactor.control.rod_fraction == reactor.control.rod_fraction
    assert restored_reactor.control.manual_control == reactor.control.manual_control
    assert len(restored_state.primary_pumps) == 2
    assert len(restored_state.secondary_pumps) == 2


def test_health_monitor_flags_core_failure():
    reactor = Reactor.default()
    state = reactor.initial_state()
    state.core.fuel_temperature = constants.MAX_CORE_TEMPERATURE
    failures = reactor.health_monitor.evaluate(
        state, [True, True], [True, True], [True, True, True], state.generators, dt=200.0
    )
    assert "core" in failures
    reactor._handle_failure("core")
    assert reactor.shutdown is True


def test_maintenance_recovers_component_health():
    monitor = HealthMonitor()
    pump = monitor.component("secondary_pump_1")
    pump.integrity = 0.3
    pump.fail()
    restored = monitor.maintain("secondary_pump_1", amount=0.5)
    assert restored is True
    assert pump.integrity == pytest.approx(0.8)
    assert pump.failed is False


def test_secondary_pump_loss_triggers_scram_and_no_steam():
    reactor = Reactor.default()
    state = reactor.initial_state()
    # Make sure some power is present to trigger heat-sink logic.
    state.core.power_output_mw = 500.0
    reactor.secondary_pump_active = False
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.1
    reactor.step(state, dt=1.0)
    assert reactor.shutdown is True
    assert all(t.electrical_output_mw == 0.0 for t in state.turbines)


def test_cold_shutdown_stays_subcritical():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.95
    reactor.primary_pump_active = False
    reactor.secondary_pump_active = False
    initial_power = state.core.power_output_mw
    for _ in range(10):
        reactor.step(state, dt=1.0)
    assert state.core.power_output_mw <= initial_power + 0.5
    assert reactor.shutdown is True


def test_toggle_maintenance_progresses_until_restored():
    reactor = Reactor.default()
    reactor.primary_pump_units = [False, False]
    reactor.primary_pump_active = False
    pump = reactor.health_monitor.component("primary_pump_1")
    pump.integrity = 0.2

    def provider(t: float, _state):
        if t == 0:
            return ReactorCommand.maintain("primary_pump_1")
        return None

    sim = ReactorSimulation(reactor, timestep=1.0, duration=50.0, command_provider=provider)
    sim.log()
    assert pump.integrity >= 0.99
    assert "primary_pump_1" not in reactor.maintenance_active


def test_primary_pump_unit_toggle_updates_active_flag():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.primary_pump_active = True
    reactor.primary_pump_units = [True, True]

    reactor.step(state, dt=1.0, command=ReactorCommand(primary_pumps={1: False}))
    assert reactor.primary_pump_units == [False, True]
    assert reactor.primary_pump_active is True

    reactor.step(state, dt=1.0, command=ReactorCommand(primary_pumps={2: False}))
    assert reactor.primary_pump_units == [False, False]
    assert reactor.primary_pump_active is False


def test_secondary_pump_unit_toggle_can_restart_pump():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.secondary_pump_active = False
    reactor.secondary_pump_units = [False, False]

    reactor.step(state, dt=1.0, command=ReactorCommand(secondary_pumps={1: True}))

    assert reactor.secondary_pump_units == [True, False]
    assert reactor.secondary_pump_active is True
    assert state.secondary_loop.mass_flow_rate > 0.0


def test_primary_pumps_spool_up_over_seconds():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.secondary_pump_units = [False, False]
    # Enable both pumps and command full flow; spool should take multiple steps.
    target_flow = reactor.primary_pump.flow_rate(1.0) * len(reactor.primary_pump_units)
    reactor.step(
        state,
        dt=1.0,
        command=ReactorCommand(primary_pumps={1: True, 2: True}, generator_units={1: True}, coolant_demand=1.0),
    )
    first_flow = state.primary_loop.mass_flow_rate
    assert 0.0 < first_flow < target_flow

    for _ in range(10):
        reactor.step(state, dt=1.0, command=ReactorCommand(coolant_demand=1.0))

    assert state.primary_loop.mass_flow_rate == pytest.approx(target_flow, rel=0.15)


def test_full_rod_withdrawal_reaches_gigawatt_power():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.shutdown = False
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.0
    reactor.primary_pump_active = True
    reactor.secondary_pump_active = True

    early_power = 0.0
    for step in range(60):
        reactor.step(state, dt=1.0)
        if step == 10:
            early_power = state.core.power_output_mw
    assert state.core.power_output_mw > max(2_000.0, early_power * 2)
    assert state.core.fuel_temperature > 600.0


def test_partially_inserted_rods_hold_near_three_gw():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.shutdown = False
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.4
    reactor.primary_pump_active = True
    reactor.secondary_pump_active = True

    for _ in range(120):
        reactor.step(state, dt=1.0)

    assert 2_000.0 < state.core.power_output_mw < 4_000.0
    assert 500.0 < state.core.fuel_temperature < 800.0


def test_generator_spools_and_powers_pumps():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.shutdown = False
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.95  # keep power low; focus on aux power
    reactor.turbine_unit_active = [False, False, False]
    reactor.secondary_pump_units = [False, False]

    for step in range(12):
        cmd = ReactorCommand(generator_units={1: True}, primary_pumps={1: True}) if step == 0 else None
        reactor.step(state, dt=1.0, command=cmd)

    assert state.generators and state.generators[0].running is True
    assert state.generators[0].power_output_mw > 0.0
    assert state.primary_loop.mass_flow_rate > 0.0


def test_meltdown_triggers_shutdown():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.shutdown = False
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.0
    reactor.primary_pump_active = True
    reactor.secondary_pump_active = True
    state.core.fuel_temperature = constants.CORE_MELTDOWN_TEMPERATURE + 50.0

    reactor.step(state, dt=1.0)

    assert reactor.shutdown is True
    assert reactor.meltdown is True