"""Reactor simulation harness and CLI."""

from __future__ import annotations

import copy
import json
import logging
import os
from dataclasses import dataclass, field
from threading import Event
import time
from typing import Callable, Iterable, Optional

from .commands import ReactorCommand
from .logging_utils import configure_logging
from .reactor import Reactor
from .state import PlantState

LOGGER = logging.getLogger(__name__)

CommandProvider = Callable[[float, PlantState], Optional[ReactorCommand]]


@dataclass
class ReactorSimulation:
    reactor: Reactor
    timestep: float = 1.0
    duration: float | None = 3600.0
    command_provider: CommandProvider | None = None
    realtime: bool = False
    stop_event: Event = field(default_factory=Event)
    start_state: PlantState | None = None
    last_state: PlantState | None = field(default=None, init=False)

    def run(self) -> Iterable[PlantState]:
        state = copy.deepcopy(self.start_state) if self.start_state else self.reactor.initial_state()
        elapsed = 0.0
        last_step_wall = time.time()
        while self.duration is None or elapsed < self.duration:
            if self.stop_event.is_set():
                LOGGER.info("Stop signal received, terminating simulation loop")
                break
            snapshot = copy.deepcopy(state)
            yield snapshot
            command = self.command_provider(elapsed, snapshot) if self.command_provider else None
            self.reactor.step(state, self.timestep, command)
            elapsed += self.timestep
            if self.realtime:
                wall_elapsed = time.time() - last_step_wall
                sleep_time = self.timestep - wall_elapsed
                if sleep_time > 0:
                    time.sleep(sleep_time)
                last_step_wall = time.time()
        self.last_state = state
        LOGGER.info("Simulation complete, %.0fs simulated", elapsed)

    def log(self) -> list[dict[str, float]]:
        return [snapshot for snapshot in (s.snapshot() for s in self.run())]

    def stop(self) -> None:
        self.stop_event.set()


def main() -> None:
    log_level = os.getenv("FISSION_LOG_LEVEL", "INFO")
    log_file = os.getenv("FISSION_LOG_FILE")
    configure_logging(log_level, log_file)
    realtime = os.getenv("FISSION_REALTIME", "0") == "1"
    duration_env = os.getenv("FISSION_SIM_DURATION")
    if duration_env:
        duration = None if duration_env.lower() in {"none", "infinite"} else float(duration_env)
    else:
        duration = None if realtime else 600.0
    reactor = Reactor.default()
    sim = ReactorSimulation(reactor, timestep=5.0, duration=duration, realtime=realtime)
    load_path = os.getenv("FISSION_LOAD_STATE")
    save_path = os.getenv("FISSION_SAVE_STATE")
    if load_path:
        sim.start_state = reactor.load_state(load_path)
    try:
        if realtime:
            LOGGER.info("Running in real-time mode (Ctrl+C to stop)...")
            for _ in sim.run():
                pass
        else:
            snapshots = sim.log()
            LOGGER.info("Captured %d snapshots", len(snapshots))
            print(json.dumps(snapshots[-5:], indent=2))
    except KeyboardInterrupt:
        sim.stop()
        LOGGER.warning("Simulation interrupted by user")
    finally:
        if save_path and sim.last_state:
            reactor.save_state(save_path, sim.last_state)


if __name__ == "__main__":
    main()