use crate::{EngineHandle, EpochHandle, InputTransform, PyEngineInput, PyGeneration};
use pyo3::{PyResult, pyclass, pymethods};
use radiate::{
    Chromosome, Engine, EngineIteratorExt, Generation, GeneticEngine, Limit, radiate_err,
};
use radiate_error::{radiate_py_bail, radiate_py_err};
use serde::Serialize;
use std::time::Duration;

#[pyclass(from_py_object)]
#[derive(Clone)]
pub enum PyEngineRunOption {
    Log(bool),
    Checkpoint(usize, String),
    Ui(Duration),
}

#[pymethods]
impl PyEngineRunOption {
    #[staticmethod]
    pub fn log(value: bool) -> Self {
        PyEngineRunOption::Log(value)
    }

    #[staticmethod]
    pub fn checkpoint(interval: usize, path: String) -> Self {
        PyEngineRunOption::Checkpoint(interval, path)
    }

    #[staticmethod]
    pub fn ui() -> Self {
        PyEngineRunOption::Ui(radiate::DEFAULT_RENDER_INTERVAL)
    }
}

#[pyclass(unsendable)]
pub struct PyEngine {
    engine: Option<EngineHandle>,
    limits: Vec<Limit>,
}

impl PyEngine {
    pub fn new(limits: Vec<Limit>, engine: EngineHandle) -> Self {
        Self {
            engine: Some(engine),
            limits,
        }
    }
}

#[pymethods]
impl PyEngine {
    pub fn run(
        &mut self,
        limits: Vec<PyEngineInput>,
        options: Vec<PyEngineRunOption>,
    ) -> PyResult<PyGeneration> {
        use EngineHandle::*;
        let engine = self
            .engine
            .take()
            .ok_or_else(|| radiate_py_err!("Engine has already been run"))?;

        let limits = self
            .limits
            .clone()
            .into_iter()
            .chain(limits.into_iter().filter_map(|input| input.transform()))
            .collect::<Vec<_>>();

        if limits.is_empty() {
            radiate_py_bail!("At least one limit must be provided");
        }

        Ok(PyGeneration::new(match engine {
            UInt8(eng) => EpochHandle::UInt8(run_engine(eng, limits, options)?),
            UInt16(eng) => EpochHandle::UInt16(run_engine(eng, limits, options)?),
            UInt32(eng) => EpochHandle::UInt32(run_engine(eng, limits, options)?),
            UInt64(eng) => EpochHandle::UInt64(run_engine(eng, limits, options)?),
            Int8(eng) => EpochHandle::Int8(run_engine(eng, limits, options)?),
            Int16(eng) => EpochHandle::Int16(run_engine(eng, limits, options)?),
            Int32(eng) => EpochHandle::Int32(run_engine(eng, limits, options)?),
            Int64(eng) => EpochHandle::Int64(run_engine(eng, limits, options)?),
            Float32(eng) => EpochHandle::Float32(run_engine(eng, limits, options)?),
            Float64(eng) => EpochHandle::Float64(run_engine(eng, limits, options)?),
            Char(eng) => EpochHandle::Char(run_engine(eng, limits, options)?),
            Bit(eng) => EpochHandle::Bit(run_engine(eng, limits, options)?),
            Permutation(eng) => EpochHandle::Permutation(run_engine(eng, limits, options)?),
            Graph(eng) => EpochHandle::Graph(run_engine(eng, limits, options)?),
            Tree(eng) => EpochHandle::Tree(run_engine(eng, limits, options)?),
        }))
    }

    pub fn next(&mut self) -> PyResult<PyGeneration> {
        use EngineHandle::*;
        let engine = self
            .engine
            .as_mut()
            .ok_or_else(|| radiate_py_err!("Engine has already been run"))?;

        Ok(PyGeneration::new(match engine {
            UInt8(eng) => EpochHandle::UInt8(eng.next()?),
            UInt16(eng) => EpochHandle::UInt16(eng.next()?),
            UInt32(eng) => EpochHandle::UInt32(eng.next()?),
            UInt64(eng) => EpochHandle::UInt64(eng.next()?),
            Int8(eng) => EpochHandle::Int8(eng.next()?),
            Int16(eng) => EpochHandle::Int16(eng.next()?),
            Int32(eng) => EpochHandle::Int32(eng.next()?),
            Int64(eng) => EpochHandle::Int64(eng.next()?),
            Float32(eng) => EpochHandle::Float32(eng.next()?),
            Float64(eng) => EpochHandle::Float64(eng.next()?),
            Char(eng) => EpochHandle::Char(eng.next()?),
            Bit(eng) => EpochHandle::Bit(eng.next()?),
            Permutation(eng) => EpochHandle::Permutation(eng.next()?),
            Graph(eng) => EpochHandle::Graph(eng.next()?),
            Tree(eng) => EpochHandle::Tree(eng.next()?),
        }))
    }
}

fn run_engine<C, T>(
    engine: GeneticEngine<C, T>,
    limits: Vec<Limit>,
    options: Vec<PyEngineRunOption>,
) -> PyResult<Generation<C, T>>
where
    C: Chromosome + Clone + Serialize + 'static,
    T: Clone + Send + Sync + Serialize + 'static,
{
    let ui_interval = get_ui_option(&options);
    if let Some(interval) = ui_interval {
        iter_engine(radiate::ui((engine, interval)).iter(), limits, options)
    } else {
        iter_engine(engine.iter(), limits, options)
    }
}

fn iter_engine<C, T>(
    engine: impl Iterator<Item = Generation<C, T>> + 'static,
    limits: Vec<Limit>,
    options: Vec<PyEngineRunOption>,
) -> PyResult<Generation<C, T>>
where
    C: Chromosome + Clone + Serialize + 'static,
    T: Clone + Send + Sync + Serialize + 'static,
{
    let log = get_log_option(&options);
    let checkpoint = get_checkpoint_option(&options);

    engine
        .chain_if(log.unwrap_or(false), |eng| eng.logging())
        .chain_if(checkpoint.is_some(), |eng| {
            let (interval, path) = checkpoint.unwrap();
            eng.checkpoint(interval, path)
        })
        .limit(limits)
        .last()
        .ok_or_else(|| radiate_err!(Python: "Failed to run engine and obtain final generation"))
}

fn get_log_option(options: &[PyEngineRunOption]) -> Option<bool> {
    let ui = get_ui_option(options);
    let log = options
        .iter()
        .find(|opt| matches!(opt, PyEngineRunOption::Log(_)))
        .and_then(|opt| {
            if let PyEngineRunOption::Log(val) = opt {
                Some(*val)
            } else {
                None
            }
        });

    if ui.is_some() { Some(false) } else { log }
}

fn get_checkpoint_option(options: &[PyEngineRunOption]) -> Option<(usize, String)> {
    options
        .iter()
        .find(|opt| matches!(opt, PyEngineRunOption::Checkpoint(_, _)))
        .and_then(|opt| {
            if let PyEngineRunOption::Checkpoint(interval, path) = opt {
                Some((*interval, path.clone()))
            } else {
                None
            }
        })
}

fn get_ui_option(options: &[PyEngineRunOption]) -> Option<Duration> {
    options
        .iter()
        .find(|opt| matches!(opt, PyEngineRunOption::Ui(_)))
        .and_then(|opt| {
            if let PyEngineRunOption::Ui(interval) = opt {
                Some(*interval)
            } else {
                None
            }
        })
}
