use crate::engine::{ChatTurn, CompactionOutput, GenerationOutput, Generator};
use crate::metrics::{Metrics, WorkSource};
use crate::settings::{EngineSettings, TurnSettings};
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
use std::time::{Duration, Instant};

const RUNTIME_PANIC_ERROR: &str = "The model runtime hit an internal error and was reset.";

#[derive(Clone)]
pub(crate) struct GenerationService {
    commands: Sender<Command>,
    metrics: Arc<Metrics>,
}

pub(crate) struct ActiveGeneration {
    pub(crate) events: Receiver<GenerationEvent>,
    pub(crate) cancel: Arc<AtomicBool>,
}

impl Drop for ActiveGeneration {
    fn drop(&mut self) {
        self.cancel.store(true, Ordering::Relaxed);
    }
}

pub(crate) enum CheckpointTarget {
    Local {
        checkpoint: PathBuf,
        /// Shared cache for the deterministic rendered system prompt.
        bootstrap: Option<PathBuf>,
    },
    Transient(PathBuf),
}

pub(crate) enum GenerationEvent {
    Loading,
    Activity(&'static str),
    Chunk {
        reasoning: bool,
        content: String,
    },
    Context {
        used: u32,
        limit: u32,
        tokens_per_second: Option<f32>,
    },
    Finished(Result<GenerationOutput, String>),
    Compacted(Result<CompactionOutput, String>),
    Measured(Result<u32, String>),
}

enum Operation {
    Generate,
    Compact {
        reason: String,
        rebuild_system_prompt: String,
    },
    Measure,
}

impl Operation {
    fn tracks_metrics(&self) -> bool {
        !matches!(self, Self::Measure)
    }

    fn error_handler(&self) -> fn(String) -> GenerationEvent {
        match self {
            Self::Generate => |error| GenerationEvent::Finished(Err(error)),
            Self::Compact { .. } => |error| GenerationEvent::Compacted(Err(error)),
            Self::Measure => |error| GenerationEvent::Measured(Err(error)),
        }
    }
}

struct Command {
    engine: EngineSettings,
    turn: TurnSettings,
    messages: Vec<ChatTurn>,
    checkpoint: Option<CheckpointTarget>,
    source: WorkSource,
    operation: Operation,
    idle_timeout: Duration,
    cancel: Arc<AtomicBool>,
    events: Sender<GenerationEvent>,
}

pub(crate) struct CompactionInput {
    pub(crate) engine: EngineSettings,
    pub(crate) turn: TurnSettings,
    pub(crate) messages: Vec<ChatTurn>,
    pub(crate) reason: String,
    pub(crate) rebuild_system_prompt: String,
    pub(crate) checkpoint: PathBuf,
    pub(crate) idle_timeout: Duration,
}

struct RuntimeState {
    loaded: Option<(EngineSettings, Generator)>,
    last_used: Instant,
    idle_timeout: Duration,
}

impl GenerationService {
    pub(crate) fn spawn(metrics: Arc<Metrics>) -> Result<Self, String> {
        let (commands, receiver) = mpsc::channel::<Command>();
        let worker_metrics = Arc::clone(&metrics);
        thread::Builder::new()
            .name("model-runtime".into())
            .spawn(move || run(receiver, worker_metrics))
            .map_err(|error| format!("Could not start the model runtime: {error}"))?;
        Ok(Self { commands, metrics })
    }

    pub(crate) fn generate(
        &self,
        engine: EngineSettings,
        turn: TurnSettings,
        messages: Vec<ChatTurn>,
        checkpoint: CheckpointTarget,
        source: WorkSource,
        idle_timeout: Duration,
    ) -> Result<ActiveGeneration, String> {
        self.submit(CommandRequest {
            engine,
            turn,
            messages,
            checkpoint: Some(checkpoint),
            source,
            operation: Operation::Generate,
            idle_timeout,
        })
    }

    pub(crate) fn compact(&self, input: CompactionInput) -> Result<ActiveGeneration, String> {
        self.submit(CommandRequest {
            engine: input.engine,
            turn: input.turn,
            messages: input.messages,
            checkpoint: Some(CheckpointTarget::Local {
                checkpoint: input.checkpoint,
                bootstrap: None,
            }),
            source: WorkSource::LocalChat,
            operation: Operation::Compact {
                reason: input.reason,
                rebuild_system_prompt: input.rebuild_system_prompt,
            },
            idle_timeout: input.idle_timeout,
        })
    }

    pub(crate) fn measure_context(
        &self,
        engine: EngineSettings,
        turn: TurnSettings,
        messages: Vec<ChatTurn>,
        idle_timeout: Duration,
    ) -> Result<ActiveGeneration, String> {
        self.submit(CommandRequest {
            engine,
            turn,
            messages,
            checkpoint: None,
            source: WorkSource::LocalChat,
            operation: Operation::Measure,
            idle_timeout,
        })
    }

    fn submit(&self, request: CommandRequest) -> Result<ActiveGeneration, String> {
        let cancel = Arc::new(AtomicBool::new(false));
        let (events, receiver) = mpsc::channel();
        let tracked = request.operation.tracks_metrics();
        if tracked {
            self.metrics.request_queued(request.source);
        }
        let command = Command {
            engine: request.engine,
            turn: request.turn,
            messages: request.messages,
            checkpoint: request.checkpoint,
            source: request.source,
            operation: request.operation,
            idle_timeout: request.idle_timeout,
            cancel: Arc::clone(&cancel),
            events,
        };
        if self.commands.send(command).is_err() {
            if tracked {
                self.metrics.request_rejected();
            }
            return Err("The model runtime stopped unexpectedly.".to_owned());
        }
        Ok(ActiveGeneration {
            events: receiver,
            cancel,
        })
    }
}

struct CommandRequest {
    engine: EngineSettings,
    turn: TurnSettings,
    messages: Vec<ChatTurn>,
    checkpoint: Option<CheckpointTarget>,
    source: WorkSource,
    operation: Operation,
    idle_timeout: Duration,
}

fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
    let mut state = RuntimeState {
        loaded: None,
        last_used: Instant::now(),
        idle_timeout: Duration::from_secs(15 * 60),
    };
    loop {
        match commands.recv_timeout(Duration::from_secs(1)) {
            Ok(command) => {
                let request_started = Instant::now();
                let source = command.source;
                let events = command.events.clone();
                let error_event = command.operation.error_handler();
                let tracked = command.operation.tracks_metrics();
                if tracked {
                    metrics.request_started(source);
                }
                if let Err(error) = catch_runtime_panic(|| {
                    run_command(command, &mut state, &metrics, request_started, source);
                }) {
                    if tracked {
                        metrics.request_failed(request_started.elapsed());
                    }
                    if state.loaded.take().is_some() {
                        metrics.unloaded();
                    }
                    let _ = events.send(error_event(error));
                }
            }
            Err(mpsc::RecvTimeoutError::Timeout) => {
                if state.loaded.is_some() && state.last_used.elapsed() >= state.idle_timeout {
                    state.loaded = None;
                    metrics.unloaded();
                }
            }
            Err(mpsc::RecvTimeoutError::Disconnected) => break,
        }
    }
}

fn run_command(
    command: Command,
    state: &mut RuntimeState,
    metrics: &Arc<Metrics>,
    request_started: Instant,
    source: WorkSource,
) {
    let error_event = command.operation.error_handler();
    let tracked = command.operation.tracks_metrics();
    state.idle_timeout = command.idle_timeout;
    if command.cancel.load(Ordering::Relaxed) {
        if tracked {
            metrics.request_failed(request_started.elapsed());
        }
        let _ = command
            .events
            .send(error_event("generation cancelled".into()));
        return;
    }
    if state
        .loaded
        .as_ref()
        .is_none_or(|(settings, _)| settings != &command.engine)
    {
        if state.loaded.take().is_some() {
            metrics.unloaded();
        }
        let _ = command.events.send(GenerationEvent::Loading);
        metrics.loading();
        let load_started = Instant::now();
        state.loaded = match Generator::open(&command.engine, Arc::clone(metrics)) {
            Ok(generator) => {
                let summary = generator.summary();
                metrics.loaded(
                    summary.model,
                    load_started.elapsed(),
                    summary.mapped_bytes,
                    summary.tensor_count,
                    summary.vocabulary_size,
                );
                Some((command.engine.clone(), generator))
            }
            Err(error) => {
                if tracked {
                    metrics.request_failed(request_started.elapsed());
                }
                let _ = command.events.send(error_event(error));
                None
            }
        };
    }
    if let Some((_, generator)) = &mut state.loaded {
        let mut prefill_started = None::<(Instant, u32)>;
        let mut emit = |reasoning, content| {
            let _ = command
                .events
                .send(GenerationEvent::Chunk { reasoning, content });
        };
        let mut progress = |used, limit, tokens_per_second| {
            if let Some(speed) = tokens_per_second {
                metrics.generation_progress(used, limit, speed);
            } else {
                let (started, initial) =
                    prefill_started.get_or_insert_with(|| (Instant::now(), used));
                let elapsed = started.elapsed().as_secs_f32();
                let speed = if elapsed > 0.0 {
                    used.saturating_sub(*initial) as f32 / elapsed
                } else {
                    0.0
                };
                metrics.prefill_progress(used, limit, speed);
            }
            let _ = command.events.send(GenerationEvent::Context {
                used,
                limit,
                tokens_per_second,
            });
        };
        if let Operation::Compact {
            reason,
            rebuild_system_prompt,
        } = &command.operation
        {
            let Some(CheckpointTarget::Local { checkpoint, .. }) = &command.checkpoint else {
                unreachable!("compaction checkpoints are local")
            };
            let result = generator.compact(
                &command.messages,
                &command.turn,
                rebuild_system_prompt,
                reason,
                checkpoint,
                &command.cancel,
                &mut progress,
                |activity| {
                    let _ = command.events.send(GenerationEvent::Activity(activity));
                },
            );
            match &result {
                Ok(_) => {
                    metrics.request_finished(source, request_started.elapsed(), 0, 0, 0, None, 0)
                }
                Err(_) => metrics.request_failed(request_started.elapsed()),
            }
            let _ = command.events.send(GenerationEvent::Compacted(result));
            state.last_used = Instant::now();
            return;
        }
        if matches!(command.operation, Operation::Measure) {
            let result = generator.rendered_history_tokens(&command.messages, &command.turn);
            let _ = command.events.send(GenerationEvent::Measured(result));
            state.last_used = Instant::now();
            return;
        }
        let result = match command
            .checkpoint
            .expect("generation requires a checkpoint target")
        {
            CheckpointTarget::Local {
                checkpoint,
                bootstrap,
            } => generator.generate(
                &checkpoint,
                bootstrap.as_deref(),
                &command.messages,
                &command.turn,
                &command.cancel,
                &mut emit,
                &mut progress,
                |activity| {
                    let _ = command.events.send(GenerationEvent::Activity(activity));
                },
            ),
            CheckpointTarget::Transient(directory) => generator.generate_transient(
                &directory,
                &command.messages,
                &command.turn,
                &command.cancel,
                &mut emit,
                &mut progress,
            ),
        };
        match &result {
            Ok(output) => metrics.request_finished(
                source,
                request_started.elapsed(),
                output.prompt_tokens,
                output.cached_tokens,
                output.completion_tokens,
                output.previous_checkpoint_bytes,
                output.checkpoint_bytes,
            ),
            Err(_) => metrics.request_failed(request_started.elapsed()),
        }
        let _ = command.events.send(GenerationEvent::Finished(result));
        state.last_used = Instant::now();
    }
}

fn catch_runtime_panic<T>(operation: impl FnOnce() -> T) -> Result<T, String> {
    // Generator owns GPU handles and is not unwind-safe; callers discard it on error.
    std::panic::catch_unwind(std::panic::AssertUnwindSafe(operation))
        .map_err(|_| RUNTIME_PANIC_ERROR.to_owned())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn dropping_active_generation_cancels_it() {
        let cancel = Arc::new(AtomicBool::new(false));
        let (_, events) = mpsc::channel();

        drop(ActiveGeneration {
            events,
            cancel: Arc::clone(&cancel),
        });

        assert!(cancel.load(Ordering::Relaxed));
    }

    #[test]
    fn runtime_panic_recovery_accepts_the_next_operation() {
        assert_eq!(
            catch_runtime_panic(|| panic!("injected runtime panic")),
            Err(RUNTIME_PANIC_ERROR.to_owned())
        );
        assert_eq!(catch_runtime_panic(|| 42), Ok(42));
    }

    #[test]
    fn operation_failures_use_the_matching_event() {
        assert!(matches!(
            Operation::Compact {
                reason: String::new(),
                rebuild_system_prompt: String::new(),
            }
            .error_handler()("stopped".into()),
            GenerationEvent::Compacted(Err(error)) if error == "stopped"
        ));
        assert!(matches!(
            Operation::Measure.error_handler()("stopped".into()),
            GenerationEvent::Measured(Err(error)) if error == "stopped"
        ));
    }

    #[test]
    fn operation_tracking_is_consistent() {
        assert!(Operation::Generate.tracks_metrics());
        assert!(
            Operation::Compact {
                reason: String::new(),
                rebuild_system_prompt: String::new(),
            }
            .tracks_metrics()
        );
        assert!(!Operation::Measure.tracks_metrics());
    }
}