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]

    cmd = ReactorCommand(secondary_pumps={1: True}, generator_units={1: True})
    for _ in range(5):
        reactor.step(state, dt=1.0, command=cmd)
        cmd = ReactorCommand(coolant_demand=0.75)

    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)
    cmd = ReactorCommand(primary_pumps={1: True, 2: True}, generator_units={1: True}, coolant_demand=1.0)
    reactor.step(state, dt=1.0, command=cmd)
    reactor.step(state, dt=1.0, command=ReactorCommand(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
    reactor.primary_pump_units = [True, True]
    reactor.secondary_pump_units = [True, True]
    reactor.generator_auto = True
    reactor.step(state, dt=1.0, command=ReactorCommand(generator_units={1: 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.set_manual_mode(False)
    reactor.primary_pump_active = True
    reactor.secondary_pump_active = True
    reactor.primary_pump_units = [True, True]
    reactor.secondary_pump_units = [True, True]
    reactor.generator_auto = True
    reactor.step(state, dt=1.0, command=ReactorCommand(generator_units={1: True}))

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

    assert 2_500.0 < state.core.power_output_mw < 3_500.0
    assert 0.05 < reactor.control.rod_fraction < 0.9


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_generator_manual_mode_allows_single_unit_and_stop():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.shutdown = False
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.95
    reactor.generator_auto = False
    reactor.primary_pump_units = [True, False]
    reactor.secondary_pump_units = [False, False]

    reactor.step(state, dt=1.0, command=ReactorCommand(generator_units={1: True}, primary_pumps={1: True}))
    assert state.generators[0].starting or state.generators[0].running

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

    assert state.generators[0].running is True
    reactor.step(state, dt=1.0, command=ReactorCommand(generator_units={1: False}))
    for _ in range(5):
        reactor.step(state, dt=1.0)
    assert state.generators[0].running is False
    assert state.generators[1].running is False


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


def test_auto_control_resets_shutdown_and_moves_rods():
    reactor = Reactor.default()
    state = reactor.initial_state()
    reactor.shutdown = True
    reactor.control.manual_control = True
    reactor.control.rod_fraction = 0.95

    reactor.step(state, dt=1.0, command=ReactorCommand(rod_manual=False))

    assert reactor.shutdown is False
    assert reactor.control.manual_control is False
    assert reactor.control.rod_fraction < 0.95


def test_chemistry_builds_fouling_and_backpressure():
    reactor = Reactor.default()
    state = reactor.initial_state()
    # Push impurities high to accelerate fouling dynamics.
    state.dissolved_oxygen_ppm = 200.0
    state.sodium_ppm = 100.0
    state.secondary_loop.mass_flow_rate = 20_000.0
    state.secondary_loop.steam_quality = 0.3
    state.secondary_loop.temperature_out = 600.0
    state.secondary_loop.temperature_in = 560.0
    base_hx = state.hx_fouling
    base_foul = state.turbines[0].fouling_penalty if state.turbines else 0.0
    base_pressure = state.turbines[0].condenser_pressure if state.turbines else constants.CONDENSER_BASE_PRESSURE_MPA

    reactor._update_chemistry(state, dt=20.0)

    assert state.hx_fouling > base_hx
    if state.turbines:
        assert state.turbines[0].fouling_penalty > base_foul
        assert state.turbines[0].condenser_pressure >= base_pressure


def test_full_power_reaches_steam_and_turbine_output():
    """Integration: ramp to full power with staged rod control and verify sustained steam/electric output."""
    reactor = Reactor.default()
    reactor.health_monitor.disable_degradation = True
    reactor.allow_external_aux = True
    reactor.relaxed_npsh = True
    reactor.control.set_power_setpoint(3_000.0)
    state = reactor.initial_state()
    reactor.step(
        state,
        dt=1.0,
        command=ReactorCommand(
            generator_units={1: True, 2: True},
            primary_pumps={1: True, 2: True},
            secondary_pumps={1: True, 2: True},
            rod_manual=True,
            rod_position=0.55,
        ),
    )
    checkpoints = {300, 600, 900, 1800, 2700, 3600}
    results = {}
    turbines_started = False
    for i in range(3600):
        cmd = None
        if state.core.power_output_mw >= 2_500.0 and reactor.control.manual_control:
            cmd = ReactorCommand(rod_manual=False)
        if (
            not turbines_started
            and state.secondary_loop.steam_quality > 0.02
            and state.secondary_loop.pressure > 1.0
        ):
            cmd = ReactorCommand(turbine_on=True, turbine_units={1: True, 2: True, 3: True})
            turbines_started = True
        if i == 600 and not turbines_started:
            cmd = ReactorCommand(turbine_on=True, turbine_units={1: True, 2: True, 3: True})
            turbines_started = True
        reactor.step(state, dt=1.0, command=cmd)
        if state.time_elapsed in checkpoints:
            results[state.time_elapsed] = {
                "quality": state.secondary_loop.steam_quality,
                "electric": state.total_electrical_output(),
                "core_temp": state.core.fuel_temperature,
            }

    # At or after 10 minutes of operation, ensure we have meaningful steam and electrical output.
    assert results[600]["quality"] > 0.05
    assert results[600]["electric"] > 50.0
    assert results[3600]["quality"] > 0.1
    assert results[3600]["electric"] > 150.0
    # No runaway core temperatures.
    assert results[3600]["core_temp"] < constants.CORE_MELTDOWN_TEMPERATURE


def test_cooldown_reaches_ambient_without_runaway():
    """Shutdown with pumps running should cool loops toward ambient, no runaway."""
    reactor = Reactor.default()
    reactor.health_monitor.disable_degradation = True
    reactor.allow_external_aux = True
    reactor.relaxed_npsh = True
    state = reactor.initial_state()
    # Start hot.
    reactor.step(
        state,
        dt=1.0,
        command=ReactorCommand(
            generator_units={1: True, 2: True},
            primary_pumps={1: True, 2: True},
            secondary_pumps={1: True, 2: True},
            rod_manual=True,
            rod_position=0.55,
        ),
    )
    turbines_started = False
    for i in range(1800):
        cmd = None
        if not turbines_started and state.secondary_loop.steam_quality > 0.02 and state.secondary_loop.pressure > 1.0:
            cmd = ReactorCommand(turbine_on=True, turbine_units={1: True, 2: True, 3: True})
            turbines_started = True
        if i == 900:
            cmd = ReactorCommand(rod_position=0.95, turbine_on=False, turbine_units={1: False, 2: False, 3: False})
        reactor.step(state, dt=1.0, command=cmd)

    assert not reactor.meltdown
    assert state.core.power_output_mw < 1.0
    assert state.primary_loop.temperature_out < 320.0
    assert state.secondary_loop.temperature_out < 320.0