//! High-level wrapper for Haku.

// TODO: This should be used as the basis for haku-wasm as well as haku tests in the future to
// avoid duplicating code.

use eyre::{bail, Context, OptionExt};
use haku::{
    ast::Ast,
    bytecode::{Chunk, Defs, DefsImage, DefsLimits},
    compiler::{Compiler, Source},
    lexer::{lex, Lexer},
    parser::{self, Parser, ParserLimits},
    source::SourceCode,
    token::Lexis,
};
use serde::{Deserialize, Serialize};
use tracing::{info, instrument, Level};

#[derive(Debug, Clone, Deserialize, Serialize)]
// NOTE: For serialization, this struct does _not_ have serde(rename_all = "camelCase") on it,
// because we do some dynamic typing magic over on the JavaScript side to automatically call all
// the appropriate functions for setting these limits on the client side.
pub struct Limits {
    pub max_source_code_len: u32,
    pub max_chunks: usize,
    pub max_defs: usize,
    pub max_tags: usize,
    pub max_tokens: usize,
    pub max_parser_events: usize,
    pub ast_capacity: usize,
    pub chunk_capacity: usize,
    pub stack_capacity: usize,
    pub call_stack_capacity: usize,
    pub ref_capacity: usize,
    pub fuel: usize,
    pub memory: usize,
    pub render_max_depth: usize,
}

pub struct Haku {
    limits: Limits,

    defs: Defs,
    defs_image: DefsImage,

    vm: Option<haku2::Vm>,
}

impl Haku {
    pub fn new(limits: Limits) -> Self {
        let defs = Defs::new(&DefsLimits {
            max_defs: limits.max_defs,
            max_tags: limits.max_tags,
        });

        let defs_image = defs.image();

        Self {
            limits,
            defs,
            defs_image,
            vm: None,
        }
    }

    fn reset(&mut self) {
        self.defs.restore_image(&self.defs_image);
    }

    #[instrument(skip(self, code), err)]
    pub fn set_brush(&mut self, code: &str) -> eyre::Result<()> {
        info!(?code);

        self.reset();

        let code = SourceCode::limited_len(code, self.limits.max_source_code_len)
            .ok_or_eyre("source code is too long")?;

        let mut lexer = Lexer::new(Lexis::new(self.limits.max_tokens), code);
        lex(&mut lexer)?;

        let mut parser = Parser::new(
            &lexer.lexis,
            &ParserLimits {
                max_events: self.limits.max_parser_events,
            },
        );
        parser::toplevel(&mut parser);
        let mut ast = Ast::new(self.limits.ast_capacity);
        let (root, parser_diagnostics) = parser.into_ast(&mut ast)?;

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

        let mut chunk = Chunk::new(self.limits.chunk_capacity)
            .expect("chunk capacity must be representable as a 16-bit number");
        let mut compiler = Compiler::new(&mut self.defs, &mut chunk);
        haku::compiler::compile_expr(&mut compiler, &src, root)
            .context("failed to compile the chunk")?;
        let closure_spec = compiler.closure_spec();

        if !lexer.diagnostics.is_empty()
            || !parser_diagnostics.is_empty()
            || !compiler.diagnostics.is_empty()
        {
            info!(?lexer.diagnostics, ?parser_diagnostics, ?compiler.diagnostics, "diagnostics were emitted");
            bail!("diagnostics were emitted");
        }

        let scratch = self
            .vm
            .take()
            .map(|vm| vm.into_scratch())
            .unwrap_or_else(|| haku2::Scratch::new(self.limits.memory));
        let defs = haku2::Defs::parse(&self.defs.serialize_defs(), &self.defs.serialize_tags());
        // SAFETY: The code is fresh out of the compiler oven, so it is guaranteed to be valid.
        // Well, more or less. There may lurk bugs.
        let code = unsafe { haku2::Code::new(defs, chunk.bytecode, closure_spec.local_count) };
        let limits = haku2::Limits::new(haku2::LimitsSpec {
            stack_capacity: self.limits.stack_capacity,
            call_stack_capacity: self.limits.call_stack_capacity,
        });
        self.vm = Some(haku2::Vm::new(scratch, code, limits));

        info!("brush set successfully");

        Ok(())
    }

    #[instrument(skip(self), err(level = Level::INFO))]
    pub fn eval_brush(&mut self) -> eyre::Result<()> {
        let vm = self
            .vm
            .as_mut()
            .ok_or_eyre("brush is not compiled and ready to be used")?;
        vm.begin(self.limits.fuel as u32)
            .context("an exception occurred during begin()")?;
        Ok(())
    }

    #[instrument(skip(self, canvas, max_depth), err(level = Level::INFO))]
    pub fn render(&mut self, canvas: &mut dyn haku2::Canvas, max_depth: usize) -> eyre::Result<()> {
        let vm = self
            .vm
            .as_mut()
            .ok_or_eyre("VM is not ready for rendering")?;
        vm.render(canvas, max_depth)
            .context("exception while rendering")?;
        Ok(())
    }

    pub fn has_cont(&mut self) -> bool {
        self.vm.as_mut().expect("VM is not started").has_cont()
    }

    pub fn cont(&mut self) -> haku2::Cont<'_> {
        self.vm.as_mut().expect("VM is not started").cont()
    }
}
