//! llama-sys: Dynamic FFI bindings to llama.cpp via libloading.
//!
//! Loads `llama.dll` (Windows) / `libllama.so` (Linux) / `libllama.dylib` (macOS)
//! at runtime from the path specified by `LLAMA_DLL_PATH` env var, or from a
//! set of well-known locations.
//!
//! This approach requires no `.lib` import library and no source compilation.

#![allow(non_upper_case_globals)]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
#![allow(dead_code)]
#![allow(clippy::all)]

use std::ffi::{c_char, c_void};
use std::path::{Path, PathBuf};

pub use libloading::Library;

// ─── Primitive type aliases ────────────────────────────────────────────────

pub type llama_pos = i32;
pub type llama_token = i32;
pub type llama_seq_id = i32;

// ─── Opaque struct handles ─────────────────────────────────────────────────

#[repr(C)]
pub struct llama_model {
    _private: [u8; 0],
}

#[repr(C)]
pub struct llama_context {
    _private: [u8; 0],
}

#[repr(C)]
pub struct llama_vocab {
    _private: [u8; 0],
}

#[repr(C)]
pub struct llama_sampler {
    _private: [u8; 0],
}

#[repr(C)]
pub struct llama_memory_i {
    _private: [u8; 0],
}

pub type llama_memory_t = *mut llama_memory_i;

// ─── Model params ──────────────────────────────────────────────────────────

#[repr(C)]
pub struct llama_model_params {
    pub devices: *mut c_void,
    pub tensor_buft_overrides: *const c_void,
    pub n_gpu_layers: i32,
    pub split_mode: i32,
    pub main_gpu: i32,
    pub tensor_split: *const f32,
    pub progress_callback: Option<extern "C" fn(f32, *mut c_void) -> bool>,
    pub progress_callback_user_data: *mut c_void,
    pub kv_overrides: *const c_void,
    pub vocab_only: bool,
    pub use_mmap: bool,
    pub use_direct_io: bool,
    pub use_mlock: bool,
    pub check_tensors: bool,
    pub use_extra_bufts: bool,
    pub no_host: bool,
    pub no_alloc: bool,
}

// ─── Context params ────────────────────────────────────────────────────────

#[repr(C)]
pub struct llama_context_params {
    pub n_ctx: u32,
    pub n_batch: u32,
    pub n_ubatch: u32,
    pub n_seq_max: u32,
    pub n_threads: i32,
    pub n_threads_batch: i32,
    pub rope_scaling_type: i32,
    pub pooling_type: i32,
    pub attention_type: i32,
    pub flash_attn_type: i32,
    pub rope_freq_base: f32,
    pub rope_freq_scale: f32,
    pub yarn_ext_factor: f32,
    pub yarn_attn_factor: f32,
    pub yarn_beta_fast: f32,
    pub yarn_beta_slow: f32,
    pub yarn_orig_ctx: u32,
    pub defrag_thold: f32,
    pub cb_eval: Option<extern "C" fn(*mut c_void, bool) -> bool>,
    pub cb_eval_user_data: *mut c_void,
    pub type_k: i32,
    pub type_v: i32,
    pub abort_callback: Option<extern "C" fn(*mut c_void) -> bool>,
    pub abort_callback_data: *mut c_void,
    pub embeddings: bool,
    pub offload_kqv: bool,
    pub no_perf: bool,
    pub op_offload: bool,
    pub swa_full: bool,
    pub kv_unified: bool,
    pub samplers: *mut c_void,
    pub n_samplers: usize,
}

// ─── Batch ─────────────────────────────────────────────────────────────────

#[repr(C)]
pub struct llama_batch {
    pub n_tokens: i32,
    pub token: *mut llama_token,
    pub embd: *mut f32,
    pub pos: *mut llama_pos,
    pub n_seq_id: *mut i32,
    pub seq_id: *mut *mut llama_seq_id,
    pub logits: *mut i8,
}

// ─── Sampler chain params ──────────────────────────────────────────────────

#[repr(C)]
pub struct llama_sampler_chain_params {
    pub no_perf: bool,
}

// ─── Flash attn type ───────────────────────────────────────────────────────

pub const LLAMA_FLASH_ATTN_TYPE_AUTO: i32 = -1;
pub const LLAMA_FLASH_ATTN_TYPE_DISABLED: i32 = 0;
pub const LLAMA_FLASH_ATTN_TYPE_ENABLED: i32 = 1;

// ─── Dynamic loader ────────────────────────────────────────────────────────

/// Container for all resolved function pointers from llama.dll.
pub struct LlamaLib {
    // Keep the library alive as long as this struct is alive.
    _lib: Library,

    // ── Lifecycle ──────────────────────────────────────────────────────────
    pub backend_init: unsafe extern "C" fn(),
    pub backend_free: unsafe extern "C" fn(),
    pub backend_load_all: unsafe extern "C" fn(),

    // ── Defaults ──────────────────────────────────────────────────────────
    pub model_default_params: unsafe extern "C" fn() -> llama_model_params,
    pub context_default_params: unsafe extern "C" fn() -> llama_context_params,
    pub sampler_chain_default_params: unsafe extern "C" fn() -> llama_sampler_chain_params,

    // ── Model ─────────────────────────────────────────────────────────────
    pub model_load_from_file:
        unsafe extern "C" fn(*const c_char, llama_model_params) -> *mut llama_model,
    pub model_free: unsafe extern "C" fn(*mut llama_model),
    pub model_get_vocab: unsafe extern "C" fn(*const llama_model) -> *const llama_vocab,
    pub model_n_ctx_train: unsafe extern "C" fn(*const llama_model) -> i32,
    pub model_is_recurrent: unsafe extern "C" fn(*const llama_model) -> bool,
    pub model_is_hybrid: unsafe extern "C" fn(*const llama_model) -> bool,
    pub model_desc: unsafe extern "C" fn(*const llama_model, *mut c_char, usize) -> i32,
    pub model_n_params: unsafe extern "C" fn(*const llama_model) -> u64,
    pub model_size: unsafe extern "C" fn(*const llama_model) -> u64,

    // ── Context ───────────────────────────────────────────────────────────
    pub init_from_model:
        unsafe extern "C" fn(*mut llama_model, llama_context_params) -> *mut llama_context,
    pub free: unsafe extern "C" fn(*mut llama_context),
    pub n_ctx: unsafe extern "C" fn(*const llama_context) -> u32,
    pub get_memory: unsafe extern "C" fn(*const llama_context) -> llama_memory_t,

    // ── Batch ─────────────────────────────────────────────────────────────
    pub batch_get_one: unsafe extern "C" fn(*mut llama_token, i32) -> llama_batch,
    pub batch_init: unsafe extern "C" fn(i32, i32, i32) -> llama_batch,
    pub batch_free: unsafe extern "C" fn(llama_batch),

    // ── Decode ────────────────────────────────────────────────────────────
    pub decode: unsafe extern "C" fn(*mut llama_context, llama_batch) -> i32,
    pub set_n_threads: unsafe extern "C" fn(*mut llama_context, i32, i32),
    pub set_warmup: unsafe extern "C" fn(*mut llama_context, bool),
    pub synchronize: unsafe extern "C" fn(*mut llama_context),
    pub memory_clear: unsafe extern "C" fn(llama_memory_t, bool),

    // ── Logits ────────────────────────────────────────────────────────────
    pub get_logits_ith: unsafe extern "C" fn(*mut llama_context, i32) -> *mut f32,

    // ── Vocab / Tokenizer ─────────────────────────────────────────────────
    pub vocab_n_tokens: unsafe extern "C" fn(*const llama_vocab) -> i32,
    pub vocab_is_eog: unsafe extern "C" fn(*const llama_vocab, llama_token) -> bool,
    pub vocab_get_text: unsafe extern "C" fn(*const llama_vocab, llama_token) -> *const c_char,
    pub token_to_piece: Option<
        unsafe extern "C" fn(*const llama_vocab, llama_token, *mut c_char, i32, i32, bool) -> i32,
    >,
    pub tokenize: unsafe extern "C" fn(
        *const llama_vocab,
        *const c_char,
        i32,
        *mut llama_token,
        i32,
        bool,
        bool,
    ) -> i32,
    pub token_bos: unsafe extern "C" fn(*const llama_vocab) -> llama_token,
    pub token_eos: unsafe extern "C" fn(*const llama_vocab) -> llama_token,

    // ── Sampler ───────────────────────────────────────────────────────────
    pub sampler_chain_init: unsafe extern "C" fn(llama_sampler_chain_params) -> *mut llama_sampler,
    pub sampler_chain_add: unsafe extern "C" fn(*mut llama_sampler, *mut llama_sampler),
    pub sampler_free: unsafe extern "C" fn(*mut llama_sampler),
    pub sampler_sample:
        unsafe extern "C" fn(*mut llama_sampler, *mut llama_context, i32) -> llama_token,
    pub sampler_init_greedy: unsafe extern "C" fn() -> *mut llama_sampler,
    pub sampler_init_temp: unsafe extern "C" fn(f32) -> *mut llama_sampler,
    pub sampler_init_top_p: unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
    pub sampler_init_top_k: unsafe extern "C" fn(i32) -> *mut llama_sampler,
    pub sampler_init_min_p: unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
    pub sampler_init_dist: unsafe extern "C" fn(u32) -> *mut llama_sampler,
    pub sampler_init_penalties: unsafe extern "C" fn(i32, f32, f32, f32) -> *mut llama_sampler,

    // ── Misc ──────────────────────────────────────────────────────────────
    pub supports_gpu_offload: unsafe extern "C" fn() -> bool,
    pub time_us: unsafe extern "C" fn() -> i64,
}

impl LlamaLib {
    /// Load the llama shared library from the given path.
    pub fn load(path: &Path) -> anyhow::Result<Self> {
        use anyhow::Context;

        let lib = unsafe { Library::new(path) }
            .with_context(|| format!("Failed to load llama library from {:?}", path))?;

        // We need transmutes to specific fn types — use a helper macro:
        macro_rules! load {
            ($field:ident, $ty:ty, $name:literal) => {{
                let raw: libloading::Symbol<unsafe extern "C" fn()> = unsafe {
                    lib.get(concat!($name, "\0").as_bytes())
                        .with_context(|| format!("Missing symbol: {}", $name))?
                };
                unsafe { std::mem::transmute::<_, $ty>(*raw) }
            }};
        }

        Ok(Self {
            backend_init: load!(backend_init, unsafe extern "C" fn(), "llama_backend_init"),
            backend_free: load!(backend_free, unsafe extern "C" fn(), "llama_backend_free"),
            backend_load_all: {
                let sym_name = b"ggml_backend_load_all\0";
                if let Ok(raw) = unsafe { lib.get::<unsafe extern "C" fn()>(sym_name) } {
                    unsafe { std::mem::transmute(*raw) }
                } else {
                    let ggml_path = path.with_file_name("ggml.dll");
                    if let Ok(ggml_lib) = unsafe { libloading::Library::new(&ggml_path) } {
                        // Faz a extração do raw *antes* de dar leak no Box, desfazendo ownership borrow issue do Rust
                        let raw_ptr = if let Ok(raw) =
                            unsafe { ggml_lib.get::<unsafe extern "C" fn()>(sym_name) }
                        {
                            Some(unsafe { std::mem::transmute::<_, unsafe extern "C" fn()>(*raw) })
                        } else {
                            None
                        };

                        // Faz leak no handle do módulo pra manter vivo
                        std::mem::forget(ggml_lib);

                        raw_ptr.unwrap_or_else(|| {
                            panic!("Missing symbol ggml_backend_load_all in both llama.dll and ggml.dll")
                        })
                    } else {
                        anyhow::bail!("Missing symbol ggml_backend_load_all in llama.dll and could not load ggml.dll");
                    }
                }
            },

            model_default_params: load!(
                model_default_params,
                unsafe extern "C" fn() -> llama_model_params,
                "llama_model_default_params"
            ),
            context_default_params: load!(
                context_default_params,
                unsafe extern "C" fn() -> llama_context_params,
                "llama_context_default_params"
            ),
            sampler_chain_default_params: load!(
                sampler_chain_default_params,
                unsafe extern "C" fn() -> llama_sampler_chain_params,
                "llama_sampler_chain_default_params"
            ),

            model_load_from_file: load!(
                model_load_from_file,
                unsafe extern "C" fn(*const c_char, llama_model_params) -> *mut llama_model,
                "llama_model_load_from_file"
            ),
            model_free: load!(
                model_free,
                unsafe extern "C" fn(*mut llama_model),
                "llama_model_free"
            ),
            model_get_vocab: load!(
                model_get_vocab,
                unsafe extern "C" fn(*const llama_model) -> *const llama_vocab,
                "llama_model_get_vocab"
            ),
            model_n_ctx_train: load!(
                model_n_ctx_train,
                unsafe extern "C" fn(*const llama_model) -> i32,
                "llama_model_n_ctx_train"
            ),
            model_is_recurrent: load!(
                model_is_recurrent,
                unsafe extern "C" fn(*const llama_model) -> bool,
                "llama_model_is_recurrent"
            ),
            model_is_hybrid: load!(
                model_is_hybrid,
                unsafe extern "C" fn(*const llama_model) -> bool,
                "llama_model_is_hybrid"
            ),
            model_desc: load!(
                model_desc,
                unsafe extern "C" fn(*const llama_model, *mut c_char, usize) -> i32,
                "llama_model_desc"
            ),
            model_n_params: load!(
                model_n_params,
                unsafe extern "C" fn(*const llama_model) -> u64,
                "llama_model_n_params"
            ),
            model_size: load!(
                model_size,
                unsafe extern "C" fn(*const llama_model) -> u64,
                "llama_model_size"
            ),

            init_from_model: load!(
                init_from_model,
                unsafe extern "C" fn(*mut llama_model, llama_context_params) -> *mut llama_context,
                "llama_init_from_model"
            ),
            free: load!(free, unsafe extern "C" fn(*mut llama_context), "llama_free"),
            n_ctx: load!(
                n_ctx,
                unsafe extern "C" fn(*const llama_context) -> u32,
                "llama_n_ctx"
            ),
            get_memory: load!(
                get_memory,
                unsafe extern "C" fn(*const llama_context) -> llama_memory_t,
                "llama_get_memory"
            ),

            batch_get_one: load!(
                batch_get_one,
                unsafe extern "C" fn(*mut llama_token, i32) -> llama_batch,
                "llama_batch_get_one"
            ),
            batch_init: load!(
                batch_init,
                unsafe extern "C" fn(i32, i32, i32) -> llama_batch,
                "llama_batch_init"
            ),
            batch_free: load!(
                batch_free,
                unsafe extern "C" fn(llama_batch),
                "llama_batch_free"
            ),

            decode: load!(
                decode,
                unsafe extern "C" fn(*mut llama_context, llama_batch) -> i32,
                "llama_decode"
            ),
            set_n_threads: load!(
                set_n_threads,
                unsafe extern "C" fn(*mut llama_context, i32, i32),
                "llama_set_n_threads"
            ),
            set_warmup: load!(
                set_warmup,
                unsafe extern "C" fn(*mut llama_context, bool),
                "llama_set_warmup"
            ),
            synchronize: load!(
                synchronize,
                unsafe extern "C" fn(*mut llama_context),
                "llama_synchronize"
            ),
            memory_clear: load!(
                memory_clear,
                unsafe extern "C" fn(llama_memory_t, bool),
                "llama_memory_clear"
            ),

            get_logits_ith: load!(
                get_logits_ith,
                unsafe extern "C" fn(*mut llama_context, i32) -> *mut f32,
                "llama_get_logits_ith"
            ),

            vocab_n_tokens: load!(
                vocab_n_tokens,
                unsafe extern "C" fn(*const llama_vocab) -> i32,
                "llama_vocab_n_tokens"
            ),
            vocab_is_eog: load!(
                vocab_is_eog,
                unsafe extern "C" fn(*const llama_vocab, llama_token) -> bool,
                "llama_vocab_is_eog"
            ),
            vocab_get_text: load!(
                vocab_get_text,
                unsafe extern "C" fn(*const llama_vocab, llama_token) -> *const c_char,
                "llama_vocab_get_text"
            ),
            token_to_piece: unsafe {
                lib.get::<unsafe extern "C" fn(
                    *const llama_vocab,
                    llama_token,
                    *mut c_char,
                    i32,
                    i32,
                    bool,
                ) -> i32>(b"llama_token_to_piece\0")
                    .ok()
                    .map(|raw| *raw)
            },
            tokenize: load!(
                tokenize,
                unsafe extern "C" fn(
                    *const llama_vocab,
                    *const c_char,
                    i32,
                    *mut llama_token,
                    i32,
                    bool,
                    bool,
                ) -> i32,
                "llama_tokenize"
            ),
            token_bos: load!(
                token_bos,
                unsafe extern "C" fn(*const llama_vocab) -> llama_token,
                "llama_vocab_bos"
            ),
            token_eos: load!(
                token_eos,
                unsafe extern "C" fn(*const llama_vocab) -> llama_token,
                "llama_vocab_eos"
            ),

            sampler_chain_init: load!(
                sampler_chain_init,
                unsafe extern "C" fn(llama_sampler_chain_params) -> *mut llama_sampler,
                "llama_sampler_chain_init"
            ),
            sampler_chain_add: load!(
                sampler_chain_add,
                unsafe extern "C" fn(*mut llama_sampler, *mut llama_sampler),
                "llama_sampler_chain_add"
            ),
            sampler_free: load!(
                sampler_free,
                unsafe extern "C" fn(*mut llama_sampler),
                "llama_sampler_free"
            ),
            sampler_sample: load!(
                sampler_sample,
                unsafe extern "C" fn(*mut llama_sampler, *mut llama_context, i32) -> llama_token,
                "llama_sampler_sample"
            ),
            sampler_init_greedy: load!(
                sampler_init_greedy,
                unsafe extern "C" fn() -> *mut llama_sampler,
                "llama_sampler_init_greedy"
            ),
            sampler_init_temp: load!(
                sampler_init_temp,
                unsafe extern "C" fn(f32) -> *mut llama_sampler,
                "llama_sampler_init_temp"
            ),
            sampler_init_top_p: load!(
                sampler_init_top_p,
                unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
                "llama_sampler_init_top_p"
            ),
            sampler_init_top_k: load!(
                sampler_init_top_k,
                unsafe extern "C" fn(i32) -> *mut llama_sampler,
                "llama_sampler_init_top_k"
            ),
            sampler_init_min_p: load!(
                sampler_init_min_p,
                unsafe extern "C" fn(f32, usize) -> *mut llama_sampler,
                "llama_sampler_init_min_p"
            ),
            sampler_init_dist: load!(
                sampler_init_dist,
                unsafe extern "C" fn(u32) -> *mut llama_sampler,
                "llama_sampler_init_dist"
            ),
            sampler_init_penalties: load!(
                sampler_init_penalties,
                unsafe extern "C" fn(i32, f32, f32, f32) -> *mut llama_sampler,
                "llama_sampler_init_penalties"
            ),

            supports_gpu_offload: load!(
                supports_gpu_offload,
                unsafe extern "C" fn() -> bool,
                "llama_supports_gpu_offload"
            ),
            time_us: load!(time_us, unsafe extern "C" fn() -> i64, "llama_time_us"),

            _lib: lib,
        })
    }

    /// Resolve the DLL path from env var or known locations.
    pub fn find_dll() -> anyhow::Result<PathBuf> {
        // 1. Explicit env override
        if let Ok(p) = std::env::var("LLAMA_DLL_PATH") {
            let path = PathBuf::from(p);
            if path.exists() {
                return Ok(path);
            }
        }

        // 2. Load directly relying on the Windows Library loader searching the .exe dir
        // since our build.rs now copies all DLLs to target/{debug,release}/
        #[cfg(target_os = "windows")]
        {
            return Ok(PathBuf::from("llama.dll"));
        }

        // 3. Unix-like: rely on LD_LIBRARY_PATH / rpath
        #[cfg(not(target_os = "windows"))]
        {
            for name in &["libllama.so", "libllama.dylib"] {
                if let Ok(lib) = Library::new(name) {
                    drop(lib);
                    return Ok(PathBuf::from(name));
                }
            }

            anyhow::bail!(
                "Cannot find llama shared library. Set LLAMA_DLL_PATH to the full path of llama.dll"
            )
        }
    }

    /// Convenience: find + load in one call.
    pub fn load_auto() -> anyhow::Result<Self> {
        let path = Self::find_dll()?;
        tracing::info!(path = %path.display(), "Loading llama library");
        Self::load(&path)
    }
}

// SAFETY: LlamaLib only holds a Library and function pointers.
// All calls must be externally synchronized (enforced by InferenceEngine).
unsafe impl Send for LlamaLib {}
unsafe impl Sync for LlamaLib {}