#![cfg_attr(not(feature = "std"), no_std)]

extern crate alloc;

use core::{alloc::Layout, mem, ptr, slice};

use alloc::{boxed::Box, string::String, vec::Vec};
use haku::{
    ast::Ast,
    bytecode::{Chunk, Defs, DefsImage, DefsLimits},
    compiler::{compile_expr, ClosureSpec, CompileError, Compiler, Source},
    diagnostic::Diagnostic,
    lexer::{lex, Lexer},
    parser::{self, IntoAstError, Parser},
    source::{SourceCode, Span},
    token::Lexis,
};
use log::{debug, info};
use tiny_skia::{
    BlendMode, Color, FillRule, LineCap, Paint, PathBuilder, Pixmap, PremultipliedColorU8, Rect,
    Shader, Stroke, Transform,
};

pub mod logging;
#[cfg(not(feature = "std"))]
mod panicking;

#[global_allocator]
static ALLOCATOR: dlmalloc::GlobalDlmalloc = dlmalloc::GlobalDlmalloc;

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_alloc(size: usize, align: usize) -> *mut u8 {
    alloc::alloc::alloc(Layout::from_size_align(size, align).unwrap())
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_free(ptr: *mut u8, size: usize, align: usize) {
    alloc::alloc::dealloc(ptr, Layout::from_size_align(size, align).unwrap())
}

#[derive(Debug, Clone, Copy)]
struct Limits {
    max_source_code_len: usize,
    max_chunks: usize,
    max_defs: usize,
    max_tags: usize,
    max_tokens: usize,
    max_parser_events: usize,
    ast_capacity: usize,
    chunk_capacity: usize,
    stack_capacity: usize,
    call_stack_capacity: usize,
    ref_capacity: usize,
    fuel: usize,
    memory: usize,
    render_max_depth: usize,
}

impl Default for Limits {
    fn default() -> Self {
        Self {
            max_source_code_len: 65536,
            max_chunks: 2,
            max_defs: 256,
            max_tags: 256,
            max_tokens: 1024,
            max_parser_events: 1024,
            ast_capacity: 1024,
            chunk_capacity: 65536,
            stack_capacity: 1024,
            call_stack_capacity: 256,
            ref_capacity: 2048,
            fuel: 65536,
            memory: 1024 * 1024,
            render_max_depth: 256,
        }
    }
}

#[unsafe(no_mangle)]
extern "C" fn haku_limits_new() -> *mut Limits {
    let ptr = Box::leak(Box::new(Limits::default())) as *mut _;
    debug!("created limits: {ptr:?}");
    ptr
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_limits_destroy(limits: *mut Limits) {
    debug!("destroying limits: {limits:?}");
    drop(Box::from_raw(limits))
}

macro_rules! limit_setter {
    ($name:tt) => {
        paste::paste! {
            #[unsafe(no_mangle)]
            unsafe extern "C" fn [<haku_limits_set_ $name>](limits: *mut Limits, value: usize) {
                debug!("set limit {} = {value}", stringify!($name));

                let limits = &mut *limits;
                limits.$name = value;
            }
        }
    };
}

limit_setter!(max_source_code_len);
limit_setter!(max_chunks);
limit_setter!(max_defs);
limit_setter!(max_tags);
limit_setter!(max_tokens);
limit_setter!(max_parser_events);
limit_setter!(ast_capacity);
limit_setter!(chunk_capacity);
limit_setter!(stack_capacity);
limit_setter!(call_stack_capacity);
limit_setter!(ref_capacity);
limit_setter!(fuel);
limit_setter!(memory);
limit_setter!(render_max_depth);

#[derive(Debug, Clone)]
struct Instance {
    limits: Limits,

    defs: Defs,
    defs_image: DefsImage,

    compile_result2: Option<CompileResult>,
    diagnostics2: Vec<Diagnostic>,
}

#[derive(Debug, Clone)]
struct CompileResult {
    defs_string: String,
    tags_string: String,
    chunk: Chunk,
    closure_spec: ClosureSpec,
    stats: CompileStats,
}

#[derive(Debug, Clone)]
struct CompileStats {
    ast_size: usize,
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_instance_new(limits: *const Limits) -> *mut Instance {
    let limits = *limits;
    debug!("creating new instance with limits: {limits:?}");

    let defs = Defs::new(&DefsLimits {
        max_defs: limits.max_defs,
        max_tags: limits.max_tags,
    });
    let defs_image = defs.image();

    let instance = Box::new(Instance {
        limits,
        defs,
        defs_image,
        compile_result2: None,
        diagnostics2: Vec::new(),
    });

    let ptr = Box::leak(instance) as *mut _;
    debug!("created instance: {ptr:?}");
    ptr
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_instance_destroy(instance: *mut Instance) {
    debug!("destroying instance: {instance:?}");
    drop(Box::from_raw(instance));
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_reset(instance: *mut Instance) {
    debug!("resetting instance: {instance:?}");
    let instance = &mut *instance;
    instance.defs.restore_image(&instance.defs_image);
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(C)]
enum StatusCode {
    Ok,
    SourceCodeTooLong,
    TooManyTokens,
    TooManyAstNodes,
    TooManyParserEvents,
    ParserUnbalancedEvents,
    ChunkTooBig,
    DiagnosticsEmitted,
}

#[unsafe(no_mangle)]
extern "C" fn haku_is_ok(code: StatusCode) -> bool {
    code == StatusCode::Ok
}

#[unsafe(no_mangle)]
extern "C" fn haku_is_diagnostics_emitted(code: StatusCode) -> bool {
    code == StatusCode::DiagnosticsEmitted
}

#[unsafe(no_mangle)]
extern "C" fn haku_status_string(code: StatusCode) -> *const i8 {
    match code {
        StatusCode::Ok => c"ok",
        StatusCode::SourceCodeTooLong => c"source code is too long",
        StatusCode::TooManyTokens => c"source code has too many tokens",
        StatusCode::TooManyAstNodes => c"source code has too many AST nodes",
        StatusCode::TooManyParserEvents => c"source code has too many parser events",
        StatusCode::ParserUnbalancedEvents => c"parser produced unbalanced events",
        StatusCode::ChunkTooBig => c"compiled bytecode is too large",
        StatusCode::DiagnosticsEmitted => c"diagnostics were emitted",
    }
    .as_ptr()
}

struct PixmapCanvas {
    pixmap: Pixmap,
    pb: PathBuilder,
    transform: Transform,
}

#[unsafe(no_mangle)]
extern "C" fn haku_pixmap_new(width: u32, height: u32) -> *mut PixmapCanvas {
    let ptr = Box::leak(Box::new(PixmapCanvas {
        pixmap: Pixmap::new(width, height).expect("invalid pixmap size"),
        pb: PathBuilder::new(),
        transform: Transform::identity(),
    })) as *mut _;
    debug!("created pixmap with size {width}x{height}: {ptr:?}");
    ptr
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_destroy(c: *mut PixmapCanvas) {
    debug!("destroying pixmap: {c:?}");
    drop(Box::from_raw(c))
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_data(c: *mut PixmapCanvas) -> *mut u8 {
    let c = &mut *c;
    c.pixmap.pixels_mut().as_mut_ptr() as *mut u8
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_clear(c: *mut PixmapCanvas) {
    let c = &mut *c;
    c.pixmap
        .pixels_mut()
        .fill(PremultipliedColorU8::TRANSPARENT);
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_set_translation(c: *mut PixmapCanvas, x: f32, y: f32) {
    let c = &mut *c;
    c.transform = Transform::from_translate(x, y);
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_begin(c: *mut PixmapCanvas) -> bool {
    let c = &mut *c;
    c.pb.clear();
    true
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_line(
    c: *mut PixmapCanvas,
    x1: f32,
    y1: f32,
    x2: f32,
    y2: f32,
) -> bool {
    let c = &mut *c;
    c.pb.move_to(x1, y1);
    c.pb.line_to(x2, y2);
    true
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_rectangle(
    c: *mut PixmapCanvas,
    x: f32,
    y: f32,
    width: f32,
    height: f32,
) -> bool {
    let c = &mut *c;
    if let Some(rect) = Rect::from_xywh(x, y, width, height) {
        c.pb.push_rect(rect);
    }
    true
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_circle(c: *mut PixmapCanvas, x: f32, y: f32, r: f32) -> bool {
    let c = &mut *c;
    c.pb.push_circle(x, y, r);
    true
}

fn default_paint() -> Paint<'static> {
    Paint {
        shader: Shader::SolidColor(Color::BLACK),
        blend_mode: BlendMode::SourceOver,
        anti_alias: false,
        force_hq_pipeline: false,
    }
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_fill(c: *mut PixmapCanvas, r: u8, g: u8, b: u8, a: u8) -> bool {
    let c = &mut *c;
    let pb = mem::take(&mut c.pb);
    if let Some(path) = pb.finish() {
        let paint = Paint {
            shader: Shader::SolidColor(Color::from_rgba8(r, g, b, a)),
            ..default_paint()
        };
        c.pixmap
            .fill_path(&path, &paint, FillRule::EvenOdd, c.transform, None);
    }
    true
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_pixmap_stroke(
    c: *mut PixmapCanvas,
    r: u8,
    g: u8,
    b: u8,
    a: u8,
    thickness: f32,
) -> bool {
    let c = &mut *c;
    let pb = mem::take(&mut c.pb);
    if let Some(path) = pb.finish() {
        let paint = Paint {
            shader: Shader::SolidColor(Color::from_rgba8(r, g, b, a)),
            ..default_paint()
        };
        c.pixmap.stroke_path(
            &path,
            &paint,
            &Stroke {
                width: thickness,
                line_cap: LineCap::Round,
                ..Default::default()
            },
            c.transform,
            None,
        );
    }
    true
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_compile_brush(
    instance: *mut Instance,
    code_len: u32,
    code: *const u8,
) -> StatusCode {
    info!("compiling brush (2)");

    let instance = &mut *instance;

    let code = core::str::from_utf8(slice::from_raw_parts(code, code_len as usize))
        .expect("invalid UTF-8");
    let Some(code) = SourceCode::limited_len(code, instance.limits.max_source_code_len as u32)
    else {
        return StatusCode::SourceCodeTooLong;
    };

    debug!("compiling: lexing");

    let mut lexer = Lexer::new(Lexis::new(instance.limits.max_tokens), code);
    if lex(&mut lexer).is_err() {
        info!("compiling failed: too many tokens");
        return StatusCode::TooManyTokens;
    };

    debug!(
        "compiling: lexed successfully to {} tokens",
        lexer.lexis.len()
    );
    debug!("compiling: parsing");

    let mut ast = Ast::new(instance.limits.ast_capacity);
    let mut parser = Parser::new(
        &lexer.lexis,
        &haku::parser::ParserLimits {
            max_events: instance.limits.max_parser_events,
        },
    );
    parser::toplevel(&mut parser);
    let (root, mut parser_diagnostics) = match parser.into_ast(&mut ast) {
        Ok((r, d)) => (r, d),
        Err(IntoAstError::NodeAlloc(_)) => {
            info!("compiling failed: too many AST nodes");
            return StatusCode::TooManyAstNodes;
        }
        Err(IntoAstError::TooManyEvents) => {
            info!("compiling failed: too many parser events");
            return StatusCode::TooManyParserEvents;
        }
        Err(IntoAstError::UnbalancedEvents) => {
            info!("compiling failed: parser produced unbalanced events");
            return StatusCode::ParserUnbalancedEvents;
        }
    };

    debug!(
        "compiling: parsed successfully into {} AST nodes",
        ast.len()
    );

    let src = Source { code, ast: &ast };

    let mut chunk = Chunk::new(instance.limits.chunk_capacity).unwrap();
    let mut compiler = Compiler::new(&mut instance.defs, &mut chunk);
    if let Err(error) = compile_expr(&mut compiler, &src, root) {
        match error {
            CompileError::Emit => {
                info!("compiling failed: chunk overflowed while emitting code");
                return StatusCode::ChunkTooBig;
            }
        }
    }
    let closure_spec = compiler.closure_spec();

    let mut diagnostics = lexer.diagnostics;
    diagnostics.append(&mut parser_diagnostics);
    diagnostics.append(&mut compiler.diagnostics);
    instance.diagnostics2 = diagnostics;
    if !instance.diagnostics2.is_empty() {
        debug!("compiling failed: diagnostics were emitted");
        return StatusCode::DiagnosticsEmitted;
    }

    debug!(
        "compiling: chunk has {} bytes of bytecode",
        chunk.bytecode.len()
    );
    debug!("compiling: {closure_spec:?}");

    debug!("bytecode: {:?}", chunk.bytecode);
    {
        let mut cursor = 0_usize;
        for info in &chunk.span_info {
            let slice = &chunk.bytecode[cursor..cursor + info.len as usize];
            debug!(
                "{:?} | 0x{:x} {:?} | {:?}",
                info.span,
                cursor,
                slice,
                info.span.slice(src.code),
            );
            cursor += info.len as usize;
        }
    }

    instance.compile_result2 = Some(CompileResult {
        defs_string: instance.defs.serialize_defs(),
        tags_string: instance.defs.serialize_tags(),
        chunk,
        closure_spec,
        stats: CompileStats {
            ast_size: ast.len(),
        },
    });

    info!("brush code compiled into instance");

    StatusCode::Ok
}

unsafe fn unwrap_compile_result<'a>(instance: *const Instance) -> &'a CompileResult {
    let cr = (*instance).compile_result2.as_ref().unwrap();
    cr
}

unsafe fn get_compile_result<'a>(instance: *const Instance) -> Option<&'a CompileResult> {
    (*instance).compile_result2.as_ref()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_num_diagnostics(instance: *const Instance) -> u32 {
    (*instance).diagnostics2.len() as u32
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_diagnostic_start(instance: *const Instance, index: u32) -> u32 {
    (*instance).diagnostics2[index as usize].span().start
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_diagnostic_end(instance: *const Instance, index: u32) -> u32 {
    (*instance).diagnostics2[index as usize].span().end
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_diagnostic_message(instance: *const Instance, index: u32) -> *const u8 {
    (*instance).diagnostics2[index as usize].message().as_ptr()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_diagnostic_message_len(instance: *const Instance, index: u32) -> u32 {
    (*instance).diagnostics2[index as usize].message().len() as u32
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_defs_len(instance: *const Instance) -> usize {
    unwrap_compile_result(instance).defs_string.len()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_defs(instance: *const Instance) -> *const u8 {
    unwrap_compile_result(instance).defs_string.as_ptr()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_tags_len(instance: *const Instance) -> usize {
    unwrap_compile_result(instance).tags_string.len()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_tags(instance: *const Instance) -> *const u8 {
    unwrap_compile_result(instance).tags_string.as_ptr()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_bytecode_len(instance: *const Instance) -> usize {
    unwrap_compile_result(instance).chunk.bytecode.len()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_bytecode(instance: *const Instance) -> *const u8 {
    unwrap_compile_result(instance).chunk.bytecode.as_ptr()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_bytecode_find_span(instance: *const Instance, pc: u16) -> *const Span {
    let chunk = &unwrap_compile_result(instance).chunk;
    match chunk.find_span(pc) {
        Some(r) => r,
        None => ptr::null(),
    }
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_span_start(span: *const Span) -> u32 {
    (*span).start
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_span_end(span: *const Span) -> u32 {
    (*span).end
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_local_count(instance: *const Instance) -> u8 {
    unwrap_compile_result(instance).closure_spec.local_count
}

#[unsafe(no_mangle)]
unsafe extern "C" fn haku_stat_ast_size(instance: *const Instance) -> usize {
    get_compile_result(instance)
        .map(|cr| cr.stats.ast_size)
        .unwrap_or_default()
}