Skip to content
use std::process::Command;
use std::time::{SystemTime, UNIX_EPOCH};

use log::{info, warn};
use thiserror::Error;

use crate::code::Lang;

pub struct Compiler;

#[derive(Debug, Error)]
pub enum CompileError {
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),
    #[error("Compilation failed: {0}")]
    CompilationFailed(String),
    #[error("Compiler not available for: {0}")]
    CompilerNotAvailable(String),
}

impl Compiler {
    pub fn new() -> Self {
        Self
    }

    pub async fn compile(&self, source: &str, lang: Lang) -> Result<Vec<u8>, CompileError> {
        match lang {
            Lang::Kotlin => self.compile_kotlin(source).await,
            Lang::Rust => self.compile_rust(source).await,
        }
    }

    async fn compile_kotlin(&self, source: &str) -> Result<Vec<u8>, CompileError> {
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_millis();

        let temp_dir = std::env::temp_dir().join(format!("corium_kotlin_{}", timestamp));
        std::fs::create_dir_all(&temp_dir)?;

        let source_file = temp_dir.join("main.kt");
        std::fs::write(&source_file, source)?;

        let kotlinc_available = Command::new("kotlinc")
            .arg("-version")
            .output()
            .is_ok();

        if !kotlinc_available {
            warn!("kotlinc not available, using placeholder WASM");
            let wasm = create_placeholder_wasm(source);
            let _ = std::fs::remove_dir_all(&temp_dir);
            return Ok(wasm);
        }

        let build_dir = temp_dir.join("build");
        std::fs::create_dir_all(&build_dir)?;

        info!("Compiling Kotlin source to WASM...");
        let output = Command::new("kotlinc")
            .arg(&source_file)
            .arg("-target")
            .arg("wasm-js")
            .arg("-output")
            .arg(build_dir.join("output"))
            .output()?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            warn!("kotlinc compilation failed: {}", stderr);
            let wasm = create_placeholder_wasm(source);
            let _ = std::fs::remove_dir_all(&temp_dir);
            return Ok(wasm);
        }

        let wasm_file = build_dir.join("output.wasm");
        if wasm_file.exists() {
            let wasm_bytes = std::fs::read(&wasm_file)?;
            let _ = std::fs::remove_dir_all(&temp_dir);
            info!("Successfully compiled Kotlin to WASM ({} bytes)", wasm_bytes.len());
            Ok(wasm_bytes)
        } else {
            warn!("WASM output file not found, using placeholder");
            let wasm = create_placeholder_wasm(source);
            let _ = std::fs::remove_dir_all(&temp_dir);
            Ok(wasm)
        }
    }

    async fn compile_rust(&self, source: &str) -> Result<Vec<u8>, CompileError> {
        Ok(create_placeholder_wasm(source))
    }
}

pub fn create_placeholder_wasm(source: &str) -> Vec<u8> {
    let source_bytes = source.as_bytes();
    let source_len = source_bytes.len() as u32;
    let min_pages = (source_len / 65536) + 1;

    let mut wasm = vec![0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00];

    let type_payload = vec![
        0x02,
        0x60, 0x02, 0x7f, 0x7f, 0x01, 0x7f,
        0x60, 0x00, 0x01, 0x7f,
    ];
    wasm.push(0x01);
    wasm.extend_from_slice(&encode_leb128(type_payload.len() as u32));
    wasm.extend_from_slice(&type_payload);

    let import_payload = vec![
        0x01,
        0x03, 0x65, 0x6e, 0x76,
        0x05, 0x70, 0x72, 0x69, 0x6e, 0x74,
        0x00, 0x00,
    ];
    wasm.push(0x02);
    wasm.extend_from_slice(&encode_leb128(import_payload.len() as u32));
    wasm.extend_from_slice(&import_payload);

    let func_payload = vec![0x01, 0x01];
    wasm.push(0x03);
    wasm.extend_from_slice(&encode_leb128(func_payload.len() as u32));
    wasm.extend_from_slice(&func_payload);

    let mut mem_payload = vec![0x01, 0x00];
    mem_payload.extend_from_slice(&encode_leb128(min_pages));
    wasm.push(0x05);
    wasm.extend_from_slice(&encode_leb128(mem_payload.len() as u32));
    wasm.extend_from_slice(&mem_payload);

    let export_payload = vec![
        0x02,
        0x03, 0x72, 0x75, 0x6e, 0x00, 0x01,
        0x06, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x02, 0x00,
    ];
    wasm.push(0x07);
    wasm.extend_from_slice(&encode_leb128(export_payload.len() as u32));
    wasm.extend_from_slice(&export_payload);

    let source_len_sleb = encode_sleb128(source_len as i64);
    let mut func_body = vec![0x00];
    func_body.extend_from_slice(&[0x41, 0x00]);
    func_body.push(0x41);
    func_body.extend_from_slice(&source_len_sleb);
    func_body.extend_from_slice(&[0x10, 0x00]);
    func_body.push(0x1a);
    func_body.extend_from_slice(&[0x41, 0x00]);
    func_body.push(0x0b);

    let func_body_size = encode_leb128(func_body.len() as u32);
    let mut code_payload = vec![0x01];
    code_payload.extend_from_slice(&func_body_size);
    code_payload.extend_from_slice(&func_body);
    wasm.push(0x0a);
    wasm.extend_from_slice(&encode_leb128(code_payload.len() as u32));
    wasm.extend_from_slice(&code_payload);

    let mut data_payload = vec![0x01, 0x00, 0x41, 0x00, 0x0b];
    data_payload.extend_from_slice(&encode_leb128(source_bytes.len() as u32));
    data_payload.extend_from_slice(source_bytes);
    wasm.push(0x0b);
    wasm.extend_from_slice(&encode_leb128(data_payload.len() as u32));
    wasm.extend_from_slice(&data_payload);

    wasm
}

fn encode_leb128(mut value: u32) -> Vec<u8> {
    let mut result = Vec::new();
    loop {
        let mut byte = (value & 0x7f) as u8;
        value >>= 7;
        if value != 0 {
            byte |= 0x80;
        }
        result.push(byte);
        if value == 0 {
            break;
        }
    }
    result
}

fn encode_sleb128(mut value: i64) -> Vec<u8> {
    let mut result = Vec::new();
    loop {
        let mut byte = (value & 0x7f) as u8;
        value >>= 7;
        if (value == 0 && (byte & 0x40) == 0) || (value == -1 && (byte & 0x40) != 0) {
            result.push(byte);
            break;
        }
        byte |= 0x80;
        result.push(byte);
    }
    result
}

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

    #[test]
    fn test_create_placeholder_wasm() {
        let source = "fun main() {}";
        let wasm = create_placeholder_wasm(source);
        assert!(wasm.len() > 8);
        assert_eq!(&wasm[0..8], &[0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00]);
    }

    #[test]
    fn test_create_placeholder_wasm_contains_source() {
        let source = "fun main() { println(\"hello\") }";
        let wasm = create_placeholder_wasm(source);
        let wasm_str = String::from_utf8_lossy(&wasm);
        assert!(wasm_str.contains(source));
    }

    #[tokio::test]
    async fn test_compile_kotlin_placeholder() {
        let compiler = Compiler::new();
        let result = compiler.compile("fun main() {}", Lang::Kotlin).await;
        assert!(result.is_ok());
        let wasm = result.unwrap();
        assert!(wasm.starts_with(&[0x00, 0x61, 0x73, 0x6d]));
    }

    #[tokio::test]
    async fn test_compile_rust_placeholder() {
        let compiler = Compiler::new();
        let result = compiler.compile("fn main() {}", Lang::Rust).await;
        assert!(result.is_ok());
    }
}