use std::fs;
use std::sync::Arc;

use memmap2::Mmap;
use napi::bindgen_prelude::*;
use napi_derive::napi;
use serde::Serialize;

use kiri_engine::dictionary::grammar::read_grammar;
use kiri_engine::dictionary::header::read_dictionary_header;
use kiri_engine::dictionary::lexicon::Lexicon;
use kiri_engine::dictionary::lexicon_set::LexiconSet;
use kiri_engine::inhibit_connection_in_data;
use kiri_engine::oov::mecab::{CategoryInfo, MecabOov, MecabOovConfig};
use kiri_engine::oov::simple::OovProviderConfig;
use kiri_engine::shared::SharedDictionary;
use kiri_engine::tokenizer::{build_lattice_and_solve, DictionaryCtx, LatticeInput, OovCtx};
use kiri_engine::types::{
    has_synonym_group_ids, is_user_dictionary, ConnectionCosts, Grammar, DICTIONARY_HEADER_SIZE,
};
use kiri_engine::DictData;

/// MeCab OOV options passed from JS.
#[napi(object)]
pub struct NativeMecabCategory {
    pub category_type: u32,
    pub is_invoke: bool,
    pub is_group: bool,
    pub length: u32,
}

#[napi(object)]
pub struct NativeMecabOov {
    pub category_type: u32,
    pub left_id: i32,
    pub right_id: i32,
    pub cost: i32,
    pub pos_id: i32,
}

#[napi(object)]
pub struct NativeTokenizerOptions {
    pub oov_left_id: Option<i32>,
    pub oov_right_id: Option<i32>,
    pub oov_cost: Option<i32>,
    pub oov_pos_id: Option<i32>,
    pub mecab_categories: Option<Vec<NativeMecabCategory>>,
    pub mecab_oovs: Option<Vec<NativeMecabOov>>,
    pub inhibited_connections: Option<Vec<Vec<i32>>>,
}

/// PathNode returned to JS.
#[napi(object)]
#[derive(Serialize)]
pub struct NativePathNode {
    pub begin: u32,
    pub end: u32,
    pub word_id: i32,
    pub left_id: i32,
    pub right_id: i32,
    pub cost: i32,
    pub total_cost: i32,
    pub is_oov: bool,
    pub oov_pos_id: Option<i32>,
}

/// WordInfo returned to JS.
#[napi(object)]
pub struct NativeWordInfo {
    pub surface: String,
    pub headword_length: u32,
    pub pos_id: i32,
    pub normalized_form: String,
    pub dictionary_form_word_id: i32,
    pub dictionary_form: String,
    pub reading_form: String,
    pub a_unit_split: Vec<i32>,
    pub b_unit_split: Vec<i32>,
    pub word_structure: Vec<i32>,
    pub synonym_gids: Vec<i32>,
}

/// Pre-parsed shared dictionary that can be reused across multiple tokenizers.
/// The ~150 MB DoubleArrayTrie is shared via Arc, so creating additional
/// tokenizers from this handle is nearly free.
#[napi]
pub struct NativeSharedDictionary {
    inner: SharedDictionary,
}

#[napi]
impl NativeSharedDictionary {
    /// Load a dictionary from disk and pre-parse the trie for sharing.
    #[napi(constructor)]
    pub fn new(dict_path: String, inhibited_connections: Option<Vec<Vec<i32>>>) -> Result<Self> {
        let file = fs::File::open(&dict_path)
            .map_err(|e| Error::from_reason(format!("Failed to open dictionary: {e}")))?;

        let mmap = unsafe { Mmap::map(&file) }
            .map_err(|e| Error::from_reason(format!("Failed to mmap dictionary: {e}")))?;

        let needs_inhibition = inhibited_connections
            .as_ref()
            .is_some_and(|pairs| !pairs.is_empty());

        let inner = if needs_inhibition {
            let mut data: Vec<u8> = mmap[..].to_vec();
            let pairs = inhibited_connections.unwrap();

            let (grammar, _) = read_grammar(&data, DICTIONARY_HEADER_SIZE)
                .map_err(|e| Error::from_reason(format!("Failed to read grammar: {e}")))?;

            for pair in &pairs {
                if pair.len() == 2 {
                    inhibit_connection_in_data(
                        &mut data,
                        &grammar.connection,
                        checked_i16(pair[0], "inhibited_connection.left_id")?,
                        checked_i16(pair[1], "inhibited_connection.right_id")?,
                    );
                }
            }

            SharedDictionary::new(data)
                .map_err(|e| Error::from_reason(format!("Failed to build shared dict: {e}")))?
        } else {
            // Zero-copy: keep the mmap, OS manages the pages
            SharedDictionary::from_mmap(mmap)
                .map_err(|e| Error::from_reason(format!("Failed to build shared dict: {e}")))?
        };

        Ok(Self { inner })
    }
}

#[napi]
pub struct NativeTokenizer {
    lexicon_set: LexiconSet,
    grammar: Grammar,
    /// Arc'd system dict bytes (shared with LexiconSet). Can be mmap or owned.
    system_data: Arc<DictData>,
    oov_config: OovProviderConfig,
    mecab_config: Option<MecabOovConfig>,
}

#[napi]
impl NativeTokenizer {
    /// Load a system dictionary and create the native tokenizer.
    #[napi(constructor)]
    pub fn new(dict_path: String, options: Option<NativeTokenizerOptions>) -> Result<Self> {
        let file = fs::File::open(&dict_path)
            .map_err(|e| Error::from_reason(format!("Failed to open dictionary: {e}")))?;

        // SAFETY: We keep the Mmap alive for the lifetime of NativeTokenizer.
        let mmap = unsafe { Mmap::map(&file) }
            .map_err(|e| Error::from_reason(format!("Failed to mmap dictionary: {e}")))?;

        let mut data: Vec<u8> = mmap[..].to_vec();

        let header = read_dictionary_header(&data, 0)
            .map_err(|e| Error::from_reason(format!("Failed to read header: {e}")))?;

        let has_synonyms = has_synonym_group_ids(header.version);

        let (grammar, grammar_bytes) = read_grammar(&data, DICTIONARY_HEADER_SIZE)
            .map_err(|e| Error::from_reason(format!("Failed to read grammar: {e}")))?;

        let lexicon_offset = DICTIONARY_HEADER_SIZE + grammar_bytes;
        let (lexicon, _) = Lexicon::from_bytes(&data, lexicon_offset, has_synonyms);

        // Apply connection cost inhibitions before wrapping data in Arc
        if let Some(ref opts) = options {
            if let Some(ref pairs) = opts.inhibited_connections {
                for pair in pairs {
                    if pair.len() == 2 {
                        inhibit_connection_in_data(
                            &mut data,
                            &grammar.connection,
                            checked_i16(pair[0], "inhibited_connection.left_id")?,
                            checked_i16(pair[1], "inhibited_connection.right_id")?,
                        );
                    }
                }
            }
        }

        let data_arc = Arc::new(DictData::Owned(data));
        let lexicon_set = LexiconSet::new(lexicon, data_arc.clone(), grammar.pos_list.len());

        // Build OOV config
        let oov_config = if let Some(ref opts) = options {
            OovProviderConfig {
                left_id: checked_i16(opts.oov_left_id.unwrap_or(0), "oov_left_id")?,
                right_id: checked_i16(opts.oov_right_id.unwrap_or(0), "oov_right_id")?,
                cost: checked_i16(opts.oov_cost.unwrap_or(10000), "oov_cost")?,
                pos_id: checked_i16(opts.oov_pos_id.unwrap_or(0), "oov_pos_id")?,
            }
        } else {
            OovProviderConfig::default()
        };

        // Build MeCab OOV config
        let mecab_config = if let Some(ref opts) = options {
            build_mecab_config(opts)?
        } else {
            None
        };

        Ok(Self {
            lexicon_set,
            grammar,
            system_data: data_arc,
            oov_config,
            mecab_config,
        })
    }

    /// Create a tokenizer from a shared dictionary handle.
    /// The ~150 MB trie is shared via Arc (16-byte refcount bump, not a copy).
    #[napi(factory)]
    pub fn from_shared(
        shared: &NativeSharedDictionary,
        options: Option<NativeTokenizerOptions>,
    ) -> Result<Self> {
        let data_arc = shared.inner.data().clone();

        let grammar = shared
            .inner
            .create_grammar()
            .map_err(|e| Error::from_reason(format!("Failed to read grammar: {e}")))?;

        let (lexicon, _) = shared.inner.create_lexicon();

        let lexicon_set = LexiconSet::new(lexicon, data_arc.clone(), grammar.pos_list.len());

        let oov_config = if let Some(ref opts) = options {
            OovProviderConfig {
                left_id: checked_i16(opts.oov_left_id.unwrap_or(0), "oov_left_id")?,
                right_id: checked_i16(opts.oov_right_id.unwrap_or(0), "oov_right_id")?,
                cost: checked_i16(opts.oov_cost.unwrap_or(10000), "oov_cost")?,
                pos_id: checked_i16(opts.oov_pos_id.unwrap_or(0), "oov_pos_id")?,
            }
        } else {
            OovProviderConfig::default()
        };

        let mecab_config = if let Some(ref opts) = options {
            build_mecab_config(opts)?
        } else {
            None
        };

        Ok(Self {
            lexicon_set,
            grammar,
            system_data: data_arc,
            oov_config,
            mecab_config,
        })
    }

    /// Add a user dictionary.
    #[napi]
    pub fn add_user_dictionary(&mut self, dict_path: String) -> Result<()> {
        let file = fs::File::open(&dict_path)
            .map_err(|e| Error::from_reason(format!("Failed to open user dictionary: {e}")))?;

        let mmap = unsafe { Mmap::map(&file) }
            .map_err(|e| Error::from_reason(format!("Failed to mmap user dictionary: {e}")))?;

        let data: Vec<u8> = mmap[..].to_vec();

        let header = read_dictionary_header(&data, 0)
            .map_err(|e| Error::from_reason(format!("Failed to read user dict header: {e}")))?;

        if !is_user_dictionary(header.version) {
            return Err(Error::from_reason("Not a user dictionary"));
        }

        let has_synonyms = has_synonym_group_ids(header.version);

        let (user_grammar, grammar_bytes) = read_grammar(&data, DICTIONARY_HEADER_SIZE)
            .map_err(|e| Error::from_reason(format!("Failed to read user grammar: {e}")))?;

        let lexicon_offset = DICTIONARY_HEADER_SIZE + grammar_bytes;
        let (lexicon, _) = Lexicon::from_bytes(&data, lexicon_offset, has_synonyms);

        let pos_offset = self.grammar.pos_list.len();
        for pos in &user_grammar.pos_list {
            self.grammar.pos_list.push(pos.clone());
        }

        self.lexicon_set
            .add(lexicon, Arc::new(DictData::Owned(data)), pos_offset)
            .map_err(Error::from_reason)?;

        Ok(())
    }

    /// The hot path — builds lattice and returns best path nodes.
    #[napi]
    #[allow(clippy::too_many_arguments)]
    pub fn tokenize(
        &self,
        bytes: Uint8Array,
        can_bow: Uint8Array,
        char_categories: Uint32Array,
        word_candidate_lengths: Uint32Array,
        continuous_lengths: Uint32Array,
        code_point_byte_lengths_flat: Uint8Array,
        code_point_offsets: Uint32Array,
    ) -> Result<Vec<NativePathNode>> {
        let dict = DictionaryCtx {
            lexicon: &self.lexicon_set,
            connection: ConnectionCosts::new(&self.grammar.connection, self.system_data.as_ref()),
        };
        let input = LatticeInput {
            bytes: &bytes,
            can_bow: &can_bow,
            char_categories: &char_categories,
            word_candidate_lengths: &word_candidate_lengths,
            continuous_lengths: &continuous_lengths,
            code_point_byte_lengths_flat: &code_point_byte_lengths_flat,
            code_point_offsets: &code_point_offsets,
        };
        let oov = OovCtx {
            simple: &self.oov_config,
            mecab: self.mecab_config.as_ref(),
        };
        let result = build_lattice_and_solve(&dict, &input, &oov).map_err(Error::from_reason)?;

        Ok(result
            .into_iter()
            .map(|n| NativePathNode {
                begin: n.begin as u32,
                end: n.end as u32,
                word_id: n.word_id,
                left_id: n.left_id as i32,
                right_id: n.right_id as i32,
                cost: n.cost as i32,
                total_cost: n.total_cost,
                is_oov: n.is_oov,
                oov_pos_id: n.oov_pos_id.map(|id| id as i32),
            })
            .collect())
    }

    /// Get word info by packed word ID. Called during TS post-processing.
    #[napi]
    pub fn get_word_info(&self, word_id: i32) -> Result<NativeWordInfo> {
        let wi = self.lexicon_set.get_word_info(word_id);
        Ok(NativeWordInfo {
            surface: wi.surface,
            headword_length: wi.headword_length as u32,
            pos_id: wi.pos_id as i32,
            normalized_form: wi.normalized_form,
            dictionary_form_word_id: wi.dictionary_form_word_id,
            dictionary_form: wi.dictionary_form,
            reading_form: wi.reading_form,
            a_unit_split: wi.a_unit_split,
            b_unit_split: wi.b_unit_split,
            word_structure: wi.word_structure,
            synonym_gids: wi.synonym_gids,
        })
    }

    /// Get POS tags by ID.
    #[napi]
    pub fn get_pos(&self, pos_id: u32) -> Result<Vec<String>> {
        let pos = self
            .grammar
            .pos_list
            .get(pos_id as usize)
            .ok_or_else(|| Error::from_reason(format!("POS ID {pos_id} out of range")))?;
        Ok(pos.tags.to_vec())
    }

    /// Total number of POS entries.
    #[napi(getter)]
    pub fn pos_count(&self) -> u32 {
        self.grammar.pos_list.len() as u32
    }
}

/// Convert i32 to i16, returning an error if the value doesn't fit.
fn checked_i16(value: i32, field: &str) -> Result<i16> {
    i16::try_from(value)
        .map_err(|_| Error::from_reason(format!("{field} value {value} exceeds i16 range")))
}

fn build_mecab_config(opts: &NativeTokenizerOptions) -> Result<Option<MecabOovConfig>> {
    let categories_raw = match opts.mecab_categories.as_ref() {
        Some(c) => c,
        None => return Ok(None),
    };
    let oovs_raw = match opts.mecab_oovs.as_ref() {
        Some(o) => o,
        None => return Ok(None),
    };

    let categories: Vec<CategoryInfo> = categories_raw
        .iter()
        .map(|c| CategoryInfo {
            category_type: c.category_type,
            is_invoke: c.is_invoke,
            is_group: c.is_group,
            length: c.length,
        })
        .collect();

    // Group OOV entries by category type
    let mut oov_map: Vec<(u32, Vec<MecabOov>)> = Vec::new();
    for raw in oovs_raw {
        let oov = MecabOov {
            left_id: checked_i16(raw.left_id, "mecab_oov.left_id")?,
            right_id: checked_i16(raw.right_id, "mecab_oov.right_id")?,
            cost: checked_i16(raw.cost, "mecab_oov.cost")?,
            pos_id: checked_i16(raw.pos_id, "mecab_oov.pos_id")?,
        };
        if let Some(entry) = oov_map.iter_mut().find(|(t, _)| *t == raw.category_type) {
            entry.1.push(oov);
        } else {
            oov_map.push((raw.category_type, vec![oov]));
        }
    }

    Ok(Some(MecabOovConfig {
        categories,
        oov_list: oov_map,
    }))
}
