use std::fs;
use std::sync::{Arc, RwLock};

use memmap2::Mmap;
use rustler::{Atom, Env, NifResult, NifStruct, Resource, ResourceArc};

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;

use kiri_kotoba::{build_input_text, CharCategory, InputTextOptions, RewriteRules};

mod atoms {
    rustler::atoms! {
        ok,
        error,
    }
}

/// Tokenizer resource held across NIF calls.
pub struct TokenizerResource {
    inner: RwLock<TokenizerState>,
}

#[rustler::resource_impl]
impl Resource for TokenizerResource {}

/// Shared dictionary data that can be reused across multiple tokenizers.
/// Wraps `SharedDictionary` which holds the pre-built trie via `Arc`,
/// so each tokenizer shares the ~150 MB trie instead of copying it.
pub struct SharedDictResource {
    inner: SharedDictionary,
}

#[rustler::resource_impl]
impl Resource for SharedDictResource {}

/// Full tokenizer resource: dictionary + input text processing + OOV config.
/// Used by `tokenize_text` NIF which does the complete pipeline in Rust.
pub struct FullTokenizerResource {
    lexicon_set: LexiconSet,
    grammar: Grammar,
    system_data: Arc<DictData>,
    oov_config: OovProviderConfig,
    mecab_config: Option<MecabOovConfig>,
    char_category: CharCategory,
    rewrite_rules: RewriteRules,
}

#[rustler::resource_impl]
impl Resource for FullTokenizerResource {}

struct TokenizerState {
    lexicon_set: LexiconSet,
    grammar: Grammar,
    system_data: Arc<DictData>,
    oov_config: OovProviderConfig,
    mecab_config: Option<MecabOovConfig>,
}

/// Options for creating a tokenizer, decoded from Elixir struct.
#[derive(NifStruct)]
#[module = "Kiri.Native.TokenizerOptions"]
pub struct NifTokenizerOptions {
    pub oov_left_id: i32,
    pub oov_right_id: i32,
    pub oov_cost: i32,
    pub oov_pos_id: i32,
    pub mecab_categories: Vec<NifMecabCategory>,
    pub mecab_oovs: Vec<NifMecabOov>,
    pub inhibited_connections: Vec<(i32, i32)>,
}

#[derive(NifStruct)]
#[module = "Kiri.Native.MecabCategory"]
pub struct NifMecabCategory {
    pub category_type: u32,
    pub is_invoke: bool,
    pub is_group: bool,
    pub length: u32,
}

#[derive(NifStruct)]
#[module = "Kiri.Native.MecabOov"]
pub struct NifMecabOov {
    pub category_type: u32,
    pub left_id: i32,
    pub right_id: i32,
    pub cost: i32,
    pub pos_id: i32,
}

/// PathNode returned to Elixir.
#[derive(NifStruct)]
#[module = "Kiri.Native.PathNode"]
pub struct NifPathNode {
    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 Elixir.
#[derive(NifStruct)]
#[module = "Kiri.Native.WordInfo"]
pub struct NifWordInfo {
    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>,
}

fn on_load(env: Env, _info: rustler::Term) -> bool {
    env.register::<TokenizerResource>().is_ok()
        && env.register::<SharedDictResource>().is_ok()
        && env.register::<FullTokenizerResource>().is_ok()
}

#[rustler::nif]
fn load_dictionary(
    dict_path: String,
    inhibited_connections: Vec<(i32, i32)>,
) -> NifResult<(Atom, ResourceArc<SharedDictResource>)> {
    let file = fs::File::open(&dict_path)
        .map_err(|e| rustler::Error::Term(Box::new(format!("Failed to open dictionary: {e}"))))?;

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

    let inner = if inhibited_connections.is_empty() {
        // Zero-copy: keep the mmap, OS manages the pages
        SharedDictionary::from_mmap(mmap).map_err(|e| {
            rustler::Error::Term(Box::new(format!("Failed to build shared dict: {e}")))
        })?
    } else {
        // Must copy to apply inhibitions
        let mut data: Vec<u8> = mmap[..].to_vec();

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

        for (left_id, right_id) in &inhibited_connections {
            inhibit_connection_in_data(
                &mut data,
                &grammar.connection,
                checked_i16(*left_id, "inhibited_connection.left_id")?,
                checked_i16(*right_id, "inhibited_connection.right_id")?,
            );
        }

        SharedDictionary::new(data).map_err(|e| {
            rustler::Error::Term(Box::new(format!("Failed to build shared dict: {e}")))
        })?
    };

    let resource = ResourceArc::new(SharedDictResource { inner });

    Ok((atoms::ok(), resource))
}

#[rustler::nif]
fn create_tokenizer_from_shared(
    shared: ResourceArc<SharedDictResource>,
    options: NifTokenizerOptions,
) -> NifResult<(Atom, ResourceArc<TokenizerResource>)> {
    let data_arc = shared.inner.data().clone();

    let grammar = shared
        .inner
        .create_grammar()
        .map_err(|e| rustler::Error::Term(Box::new(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 = OovProviderConfig {
        left_id: checked_i16(options.oov_left_id, "oov_left_id")?,
        right_id: checked_i16(options.oov_right_id, "oov_right_id")?,
        cost: checked_i16(options.oov_cost, "oov_cost")?,
        pos_id: checked_i16(options.oov_pos_id, "oov_pos_id")?,
    };

    let mecab_config = build_mecab_config(&options)?;

    let resource = ResourceArc::new(TokenizerResource {
        inner: RwLock::new(TokenizerState {
            lexicon_set,
            grammar,
            system_data: data_arc,
            oov_config,
            mecab_config,
        }),
    });

    Ok((atoms::ok(), resource))
}

#[rustler::nif]
fn create_tokenizer(
    dict_path: String,
    options: NifTokenizerOptions,
) -> NifResult<(Atom, ResourceArc<TokenizerResource>)> {
    let file = fs::File::open(&dict_path)
        .map_err(|e| rustler::Error::Term(Box::new(format!("Failed to open dictionary: {e}"))))?;

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

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

    let header = read_dictionary_header(&data, 0)
        .map_err(|e| rustler::Error::Term(Box::new(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| rustler::Error::Term(Box::new(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
    for (left_id, right_id) in &options.inhibited_connections {
        inhibit_connection_in_data(
            &mut data,
            &grammar.connection,
            checked_i16(*left_id, "inhibited_connection.left_id")?,
            checked_i16(*right_id, "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());

    let oov_config = OovProviderConfig {
        left_id: checked_i16(options.oov_left_id, "oov_left_id")?,
        right_id: checked_i16(options.oov_right_id, "oov_right_id")?,
        cost: checked_i16(options.oov_cost, "oov_cost")?,
        pos_id: checked_i16(options.oov_pos_id, "oov_pos_id")?,
    };

    let mecab_config = build_mecab_config(&options)?;

    let resource = ResourceArc::new(TokenizerResource {
        inner: RwLock::new(TokenizerState {
            lexicon_set,
            grammar,
            system_data: data_arc,
            oov_config,
            mecab_config,
        }),
    });

    Ok((atoms::ok(), resource))
}

#[rustler::nif]
fn add_user_dictionary(
    resource: ResourceArc<TokenizerResource>,
    dict_path: String,
) -> NifResult<Atom> {
    let file = fs::File::open(&dict_path).map_err(|e| {
        rustler::Error::Term(Box::new(format!("Failed to open user dictionary: {e}")))
    })?;

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

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

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

    if !is_user_dictionary(header.version) {
        return Err(rustler::Error::Term(Box::new(
            "Not a user dictionary".to_string(),
        )));
    }

    let has_synonyms = has_synonym_group_ids(header.version);

    let (user_grammar, grammar_bytes) = read_grammar(&data, DICTIONARY_HEADER_SIZE)
        .map_err(|e| rustler::Error::Term(Box::new(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 mut state = resource
        .inner
        .write()
        .map_err(|_| rustler::Error::Term(Box::new("Lock poisoned".to_string())))?;

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

    state
        .lexicon_set
        .add(lexicon, Arc::new(DictData::Owned(data)), pos_offset)
        .map_err(|e| rustler::Error::Term(Box::new(e)))?;

    Ok(atoms::ok())
}

#[rustler::nif]
#[allow(clippy::too_many_arguments)]
fn tokenize(
    resource: ResourceArc<TokenizerResource>,
    bytes: Vec<u8>,
    can_bow: Vec<u8>,
    char_categories: Vec<u32>,
    word_candidate_lengths: Vec<u32>,
    continuous_lengths: Vec<u32>,
    code_point_byte_lengths_flat: Vec<u8>,
    code_point_offsets: Vec<u32>,
) -> NifResult<Vec<NifPathNode>> {
    let state = resource
        .inner
        .read()
        .map_err(|_| rustler::Error::Term(Box::new("Lock poisoned".to_string())))?;

    let dict = DictionaryCtx {
        lexicon: &state.lexicon_set,
        connection: ConnectionCosts::new(&state.grammar.connection, state.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: &state.oov_config,
        mecab: state.mecab_config.as_ref(),
    };
    let result = build_lattice_and_solve(&dict, &input, &oov)
        .map_err(|e| rustler::Error::Term(Box::new(e)))?;

    Ok(result
        .into_iter()
        .map(|n| NifPathNode {
            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())
}

#[rustler::nif]
fn get_word_info(resource: ResourceArc<TokenizerResource>, word_id: i32) -> NifResult<NifWordInfo> {
    let state = resource
        .inner
        .read()
        .map_err(|_| rustler::Error::Term(Box::new("Lock poisoned".to_string())))?;

    let wi = state.lexicon_set.get_word_info(word_id);
    Ok(NifWordInfo {
        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,
    })
}

#[rustler::nif]
fn get_pos(resource: ResourceArc<TokenizerResource>, pos_id: u32) -> NifResult<Vec<String>> {
    let state = resource
        .inner
        .read()
        .map_err(|_| rustler::Error::Term(Box::new("Lock poisoned".to_string())))?;

    let pos =
        state.grammar.pos_list.get(pos_id as usize).ok_or_else(|| {
            rustler::Error::Term(Box::new(format!("POS ID {pos_id} out of range")))
        })?;
    Ok(pos.tags.to_vec())
}

#[rustler::nif]
fn pos_count(resource: ResourceArc<TokenizerResource>) -> NifResult<u32> {
    let state = resource
        .inner
        .read()
        .map_err(|_| rustler::Error::Term(Box::new("Lock poisoned".to_string())))?;
    Ok(state.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) -> NifResult<i16> {
    i16::try_from(value).map_err(|_| {
        rustler::Error::Term(Box::new(format!("{field} value {value} exceeds i16 range")))
    })
}

fn build_mecab_config(opts: &NifTokenizerOptions) -> NifResult<Option<MecabOovConfig>> {
    if opts.mecab_categories.is_empty() || opts.mecab_oovs.is_empty() {
        return Ok(None);
    }

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

    let mut oov_map: Vec<(u32, Vec<MecabOov>)> = Vec::new();
    for raw in &opts.mecab_oovs {
        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,
    }))
}

/// Options for text tokenization via the full pipeline NIF.
#[derive(NifStruct)]
#[module = "Kiri.Native.TextTokenizeOptions"]
pub struct NifTextTokenizeOptions {
    pub disable_normalization: bool,
    pub prolonged_sound_marks: bool,
    pub ignore_yomigana: bool,
}

/// Result from `tokenize_text` NIF: path nodes + modified text + byte-to-original mapping.
#[derive(NifStruct)]
#[module = "Kiri.Native.TextTokenizeResult"]
pub struct NifTextTokenizeResult {
    pub path_nodes: Vec<NifPathNode>,
    pub modified_text: String,
    pub byte_to_original: Vec<u32>,
    pub char_categories: Vec<u32>,
}

#[rustler::nif]
fn create_full_tokenizer(
    shared: ResourceArc<SharedDictResource>,
    options: NifTokenizerOptions,
) -> NifResult<(Atom, ResourceArc<FullTokenizerResource>)> {
    let data_arc = shared.inner.data().clone();

    let grammar = shared
        .inner
        .create_grammar()
        .map_err(|e| rustler::Error::Term(Box::new(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 = OovProviderConfig {
        left_id: checked_i16(options.oov_left_id, "oov_left_id")?,
        right_id: checked_i16(options.oov_right_id, "oov_right_id")?,
        cost: checked_i16(options.oov_cost, "oov_cost")?,
        pos_id: checked_i16(options.oov_pos_id, "oov_pos_id")?,
    };

    let mecab_config = build_mecab_config(&options)?;

    let char_category = CharCategory::default_instance();
    let rewrite_rules = RewriteRules::default_instance();

    let resource = ResourceArc::new(FullTokenizerResource {
        lexicon_set,
        grammar,
        system_data: data_arc,
        oov_config,
        mecab_config,
        char_category,
        rewrite_rules,
    });

    Ok((atoms::ok(), resource))
}

#[rustler::nif(schedule = "DirtyCpu")]
fn tokenize_text(
    resource: ResourceArc<FullTokenizerResource>,
    text: String,
    options: NifTextTokenizeOptions,
) -> NifResult<NifTextTokenizeResult> {
    let state = &*resource;

    let input_opts = InputTextOptions {
        disable_normalization: options.disable_normalization,
        prolonged_sound_marks: options.prolonged_sound_marks,
        ignore_yomigana: options.ignore_yomigana,
    };

    let input = build_input_text(
        &text,
        &state.char_category,
        &state.rewrite_rules,
        &input_opts,
    );

    let lattice_input = input.lattice_input.as_ref();
    let connection = ConnectionCosts::new(&state.grammar.connection, &state.system_data);

    let dict = DictionaryCtx {
        lexicon: &state.lexicon_set,
        connection,
    };
    let oov = OovCtx {
        simple: &state.oov_config,
        mecab: state.mecab_config.as_ref(),
    };

    let result = build_lattice_and_solve(&dict, &lattice_input, &oov)
        .map_err(|e| rustler::Error::Term(Box::new(e)))?;

    let path_nodes = result
        .into_iter()
        .map(|n| NifPathNode {
            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();

    Ok(NifTextTokenizeResult {
        path_nodes,
        modified_text: input.modified_text.clone(),
        byte_to_original: input.byte_to_original,
        char_categories: input.lattice_input.char_categories.clone(),
    })
}

#[rustler::nif]
fn full_tokenizer_get_word_info(
    resource: ResourceArc<FullTokenizerResource>,
    word_id: i32,
) -> NifResult<NifWordInfo> {
    let state = &*resource;
    let wi = state.lexicon_set.get_word_info(word_id);
    Ok(NifWordInfo {
        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,
    })
}

#[rustler::nif]
fn full_tokenizer_get_pos(
    resource: ResourceArc<FullTokenizerResource>,
    pos_id: u32,
) -> NifResult<Vec<String>> {
    let state = &*resource;
    let pos =
        state.grammar.pos_list.get(pos_id as usize).ok_or_else(|| {
            rustler::Error::Term(Box::new(format!("POS ID {pos_id} out of range")))
        })?;
    Ok(pos.tags.to_vec())
}

#[rustler::nif]
fn full_tokenizer_pos_count(resource: ResourceArc<FullTokenizerResource>) -> NifResult<u32> {
    let state = &*resource;
    Ok(state.grammar.pos_list.len() as u32)
}

rustler::init!("Elixir.Kiri.Native", load = on_load);
