use std::fmt;

use wasmi::{Caller, Engine, Linker, Module, Store, Val};

use crate::code::Code;

#[derive(Debug)]
pub struct RunResult {
    pub status: u32,
    pub output: String,
}

impl fmt::Display for RunResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{{\"status\": {}, \"output\": {}}}",
            self.status,
            serde_json::to_string(&self.output).unwrap_or_default()
        )
    }
}

#[derive(Debug)]
pub enum RunError {
    Instantiation(String),
    ExportNotFound(String),
    Execution(String),
    InvalidModule,
}

impl fmt::Display for RunError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            RunError::Instantiation(msg) => write!(f, "WASM instantiation failed: {}", msg),
            RunError::ExportNotFound(name) => write!(f, "Export not found: {}", name),
            RunError::Execution(msg) => write!(f, "WASM execution failed: {}", msg),
            RunError::InvalidModule => write!(f, "Invalid WASM module"),
        }
    }
}

impl std::error::Error for RunError {}

pub struct Runner {
    engine: Engine,
}

impl Runner {
    pub fn new() -> Self {
        Self {
            engine: Engine::default(),
        }
    }

    pub fn run(&self, code: &Code, params: &str) -> Result<RunResult, RunError> {
        let wasm_bytes = code.wasm();
        if wasm_bytes.is_empty() {
            return Ok(RunResult {
                status: 200,
                output: format!("No WASM binary available. Source:\n{}", code.raw()),
            });
        }

        let mut store = Store::new(&self.engine, RunnerState::new(params));
        let module = Module::new(&self.engine, wasm_bytes)
            .map_err(|e| RunError::Instantiation(e.to_string()))?;

        let mut linker = <Linker<RunnerState>>::new(&self.engine);
        linker
            .func_wrap("env", "print", host_print)
            .map_err(|e| RunError::Instantiation(e.to_string()))?;

        let instance_pre = linker
            .instantiate(&mut store, &module)
            .map_err(|e| RunError::Instantiation(e.to_string()))?;

        let instance = instance_pre
            .start(&mut store)
            .map_err(|e| RunError::Instantiation(e.to_string()))?;

        let run_func = instance
            .get_func(&store, "run")
            .ok_or_else(|| RunError::ExportNotFound("run".to_string()))?;

        let mut results = [Val::I32(0)];
        run_func
            .call(&mut store, &[], &mut results)
            .map_err(|e| RunError::Execution(e.to_string()))?;

        let state = store.into_data();
        let status = match results.first() {
            Some(Val::I32(s)) => *s as u32,
            _ => 200,
        };

        Ok(RunResult {
            status,
            output: state.output,
        })
    }
}

#[derive(Default)]
struct RunnerState {
    output: String,
    #[allow(dead_code)]
    params: String,
}

impl RunnerState {
    fn new(params: &str) -> Self {
        Self {
            output: String::new(),
            params: params.to_string(),
        }
    }
}

fn host_print(mut caller: Caller<RunnerState>, ptr: i32, len: i32) -> i32 {
    if ptr < 0 || len < 0 {
        return 0;
    }
    let offset = ptr as usize;
    let length = len as usize;

    let output_str = {
        let mem_opt = caller.get_export("memory").and_then(|e| e.into_memory());
        match mem_opt {
            Some(mem) => {
                let data = mem.data(&caller);
                if offset.checked_add(length).map_or(false, |end| end <= data.len()) {
                    String::from_utf8(data[offset..offset + length].to_vec()).ok()
                } else {
                    None
                }
            }
            None => None,
        }
    };

    if let Some(s) = output_str {
        caller.data_mut().output.push_str(&s);
    }

    0
}