use std::time::Duration;

use tokio::{
    sync::{mpsc, oneshot, watch},
    task::JoinHandle,
};
use tokio_util::sync::CancellationToken;
use tracing::{Instrument, debug, debug_span, info, trace, warn};

use crate::{
    error::{Error, ErrorHandler, ProcessFailure, ProcessFailureKind, ProcessResult},
    handle_id::{generate_short_id, HandleID},
    process_handler::{IntoProcess, ProcessHandler, SharedProcess},
    shutdown_config::ShutdownConfig,
};

/// Unique identifier for each process invocation (generated each time a process starts)
#[derive(Debug, Clone)]
struct InvocationID(String);

impl InvocationID {
    fn new() -> Self {
        Self(generate_short_id())
    }
}

impl std::fmt::Display for InvocationID {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

pub struct Process {
    handle_id: HandleID,
    command_tx: mpsc::Sender<ProcessCommand>,
    state: watch::Receiver<ProcessState>,
    task: JoinHandle<()>,
}

impl Process {
    pub fn new(process: impl ProcessHandler) -> Self {
        ProcessBuilder::new(process).build()
    }

    pub fn new_with_handle(handle_id: impl Into<HandleID>, process: impl ProcessHandler) -> Self {
        ProcessBuilder::new(process).handle_id(handle_id).build()
    }

    pub fn builder(process: impl ProcessHandler) -> ProcessBuilder {
        ProcessBuilder::new(process)
    }

    fn create(
        handle_id: HandleID,
        process: SharedProcess,
        shutdown_config: ShutdownConfig,
        error_handler: Option<ErrorHandler>,
    ) -> Self {
        let (command_tx, command_rx) = mpsc::channel(1);
        let (event_tx, event_rx) = watch::channel(ProcessState::Pending);

        debug!(handle_id = %handle_id, "creating process");

        let span = debug_span!("process_loop", handle_id = %handle_id);
        let handle = tokio::spawn(
            Self::run_loop(
                handle_id.clone(),
                process,
                command_rx,
                shutdown_config,
                event_tx,
                error_handler,
            )
            .instrument(span),
        );

        Self {
            handle_id,
            state: event_rx,
            command_tx,
            task: handle,
        }
    }
}

pub struct ProcessBuilder {
    process: SharedProcess,
    handle_id: Option<HandleID>,
    shutdown_config: Option<ShutdownConfig>,
    error_handler: Option<ErrorHandler>,
}

impl ProcessBuilder {
    pub fn new(process: impl ProcessHandler) -> Self {
        Self {
            process: process.into_process(),
            handle_id: None,
            shutdown_config: None,
            error_handler: None,
        }
    }

    pub fn handle_id(mut self, handle_id: impl Into<HandleID>) -> Self {
        self.handle_id = Some(handle_id.into());
        self
    }

    pub fn shutdown_config(mut self, config: ShutdownConfig) -> Self {
        self.shutdown_config = Some(config);
        self
    }

    /// Set an error handler that will be called when the process errors or panics
    pub fn on_error<F>(mut self, handler: F) -> Self
    where
        F: Fn(ProcessFailure) + Send + Sync + 'static,
    {
        self.error_handler = Some(std::sync::Arc::new(handler));
        self
    }

    pub fn build(self) -> Process {
        let handle_id = self.handle_id.unwrap_or_else(HandleID::new_random);
        let shutdown_config = self.shutdown_config.unwrap_or_default();
        Process::create(handle_id, self.process, shutdown_config, self.error_handler)
    }
}

impl Process {
    async fn run_loop(
        handle_id: HandleID,
        process: SharedProcess,
        mut commands: mpsc::Receiver<ProcessCommand>,
        shutdown_config: ShutdownConfig,
        state: watch::Sender<ProcessState>,
        error_handler: Option<ErrorHandler>,
    ) {
        let mut current_task: Option<RunningTask> = None;

        trace!("entering run loop");

        loop {
            match current_task.take() {
                Some(mut task) => {
                    tokio::select! {
                        biased;
                        Some(cmd) = commands.recv() => {
                            trace!(command = ?cmd, "received command while running");
                            current_task = Self::handle_command_while_running(
                                cmd,
                                task,
                                &state,
                                &shutdown_config
                            ).await;
                        }

                        result = &mut task.handle => {
                            let invocation_id = task.invocation_id.clone();
                            trace!(%invocation_id, "task completed, handling result");
                            Self::handle_task_completion(
                                result,
                                process.clone(),
                                &handle_id,
                                &state,
                                &shutdown_config,
                                &error_handler,
                                &mut current_task
                            ).await;
                        }
                    }
                }
                None => match commands.recv().await {
                    Some(ProcessCommand::Start) => {
                        if matches!(
                            *state.borrow(),
                            ProcessState::Pending | ProcessState::Stopped | ProcessState::Errored
                        ) {
                            let task = RunningTask::spawn(process.clone(), &handle_id);
                            debug!(
                                invocation_id = %task.invocation_id,
                                "process started"
                            );
                            current_task = Some(task);
                            let _ = state.send(ProcessState::Running);
                        } else {
                            debug!("ignoring start command, process already running");
                        }
                    }
                    Some(ProcessCommand::Stop { response } | ProcessCommand::Kill { response }) => {
                        trace!("received stop/kill command while not running");
                        let _ = response.send(());
                    }
                    None => {
                        trace!("command channel closed, exiting run loop");
                        break;
                    }
                },
            }
        }

        trace!("exiting run loop");
    }

    pub async fn start(&mut self) -> Result<(), Error> {
        info!(handle_id = %self.handle_id, "starting process");
        self.command_tx
            .send(ProcessCommand::Start)
            .await
            .map_err(|_| Error::RunnerTerminated)
    }

    pub async fn stop(&mut self) -> Result<(), Error> {
        info!(handle_id = %self.handle_id, "stopping process");
        let (resp_tx, resp_rx) = oneshot::channel();

        self.command_tx
            .send(ProcessCommand::Stop { response: resp_tx })
            .await
            .map_err(|_| Error::RunnerTerminated)?;

        resp_rx.await.map_err(|_| Error::RunnerTerminated)
    }

    pub async fn kill(&mut self) -> Result<(), Error> {
        info!(handle_id = %self.handle_id, "killing process");
        let (resp_tx, resp_rx) = oneshot::channel();

        self.command_tx
            .send(ProcessCommand::Kill { response: resp_tx })
            .await
            .map_err(|_| Error::RunnerTerminated)?;

        resp_rx.await.map_err(|_| Error::RunnerTerminated)
    }

    pub async fn restart(&mut self) -> Result<(), Error> {
        info!(handle_id = %self.handle_id, "restarting process");

        // Stop without the info log
        let (resp_tx, resp_rx) = oneshot::channel();
        self.command_tx
            .send(ProcessCommand::Stop { response: resp_tx })
            .await
            .map_err(|_| Error::RunnerTerminated)?;
        resp_rx.await.map_err(|_| Error::RunnerTerminated)?;

        // Start without the info log
        self.command_tx
            .send(ProcessCommand::Start)
            .await
            .map_err(|_| Error::RunnerTerminated)
    }

    pub(crate) async fn wait_for(&self, desired: ProcessState) {
        let mut rx = self.state.clone();
        while *rx.borrow_and_update() != desired {
            if rx.changed().await.is_err() {
                break;
            }
        }
    }

    pub fn handle_id(&self) -> &HandleID {
        &self.handle_id
    }

    async fn handle_command_while_running(
        cmd: ProcessCommand,
        task: RunningTask,
        state: &watch::Sender<ProcessState>,
        config: &ShutdownConfig,
    ) -> Option<RunningTask> {
        let invocation_id = task.invocation_id.clone();

        match cmd {
            ProcessCommand::Start => {
                debug!(%invocation_id, "ignoring start command, process already running");
                Some(task)
            }
            ProcessCommand::Stop { response } => {
                debug!(%invocation_id, timeout_secs = config.graceful_timeout.as_secs(), "initiating graceful shutdown");
                task.shutdown_gracefully(config.graceful_timeout).await;
                let _ = state.send(ProcessState::Stopped);
                let _ = response.send(());
                debug!(%invocation_id, "process stopped");
                None
            }
            ProcessCommand::Kill { response } => {
                debug!(%invocation_id, "aborting process");
                task.handle.abort();
                let _ = task.handle.await;
                let _ = state.send(ProcessState::Stopped);
                let _ = response.send(());
                debug!(%invocation_id, "process killed");
                None
            }
        }
    }

    async fn handle_task_completion(
        result: Result<ProcessResult, tokio::task::JoinError>,
        process: SharedProcess,
        handle_id: &HandleID,
        state: &watch::Sender<ProcessState>,
        config: &ShutdownConfig,
        error_handler: &Option<ErrorHandler>,
        current_task: &mut Option<RunningTask>,
    ) {
        match result {
            Ok(Ok(())) => {
                debug!("process completed successfully");
                if config.restart_on_success {
                    if !config.restart_delay.is_zero() {
                        debug!(delay_ms = config.restart_delay.as_millis(), "waiting before restart");
                    }
                    tokio::time::sleep(config.restart_delay).await;
                    let task = RunningTask::spawn(process, handle_id);
                    debug!(invocation_id = %task.invocation_id, "process auto-restarted");
                    *current_task = Some(task);
                } else {
                    debug!("restart_on_success disabled, not restarting");
                    let _ = state.send(ProcessState::Stopped);
                }
            }
            Ok(Err(e)) => {
                warn!(error = %e, "process returned error");
                if let Some(handler) = error_handler {
                    handler(ProcessFailure {
                        handle_id: handle_id.clone(),
                        kind: ProcessFailureKind::Returned(e.to_string()),
                    });
                }
                let _ = state.send(ProcessState::Errored);
            }
            Err(e) if e.is_cancelled() => {
                trace!("task was cancelled/aborted");
            }
            Err(e) => {
                warn!(error = %e, "process panicked");
                if let Some(handler) = error_handler {
                    handler(ProcessFailure {
                        handle_id: handle_id.clone(),
                        kind: ProcessFailureKind::Panicked(e.to_string()),
                    });
                }
                let _ = state.send(ProcessState::Errored);
            }
        }
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        trace!(handle_id = %self.handle_id, "dropping process, aborting task");
        self.task.abort();
    }
}

struct RunningTask {
    handle: JoinHandle<ProcessResult>,
    cancellation: CancellationToken,
    invocation_id: InvocationID,
}

impl RunningTask {
    pub fn spawn(process: SharedProcess, handle_id: &HandleID) -> Self {
        let cancellation = CancellationToken::new();
        let invocation_id = InvocationID::new();

        trace!(%handle_id, invocation_id = %invocation_id, "spawning process task");

        let span = debug_span!(
            "process_invocation",
            %handle_id,
            invocation_id = %invocation_id
        );

        let handle = tokio::spawn({
            let cancellation = cancellation.child_token();
            async move { process.process.call_async(cancellation).await }.instrument(span)
        });

        Self {
            handle,
            cancellation,
            invocation_id,
        }
    }

    async fn shutdown_gracefully(mut self, timeout: Duration) {
        trace!(invocation_id = %self.invocation_id, "cancelling process token");
        self.cancellation.cancel();

        match tokio::time::timeout(timeout, &mut self.handle).await {
            Ok(_) => {
                debug!(invocation_id = %self.invocation_id, "process stopped gracefully");
            }
            Err(_) => {
                warn!(
                    invocation_id = %self.invocation_id,
                    timeout_secs = timeout.as_secs(),
                    "graceful shutdown timed out, aborting"
                );
                self.handle.abort();
            }
        }
    }
}

#[derive(Debug)]
enum ProcessCommand {
    Start,
    Stop { response: oneshot::Sender<()> },
    Kill { response: oneshot::Sender<()> },
}

#[derive(PartialEq)]
pub(crate) enum ProcessState {
    Pending,
    Running,
    Errored,
    Stopped,
}
