//! Escu:de Script File (.bin)
use super::list::{EnumScr, EscudeBinList, ListData, NameT, VarT};
use super::ops::base::CustomOps;
use crate::ext::io::*;
use crate::scripts::base::*;
use crate::types::*;
use crate::utils::encoding::{decode_to_string, encode_string};
use crate::utils::struct_pack::StructPack;
use anyhow::Result;
use clap::ValueEnum;
use int_enum::IntEnum;
use std::collections::{BTreeSet, HashMap};
use std::ffi::CString;
use std::io::{Read, Seek, SeekFrom};
use unicode_segmentation::UnicodeSegmentation;

#[derive(Debug, ValueEnum, Clone, Copy)]
/// Game title
pub enum EscudeOp {
    /// パニカルコンフュージョン
    Panicon,
    /// 花嫁と魔王 ～王室のハーレムは下克上～
    Hanaou,
}

#[derive(Debug)]
/// Builder for Escu:de binary script files
pub struct EscudeBinScriptBuilder {}

impl EscudeBinScriptBuilder {
    /// Creates a new instance of `EscudeBinScriptBuilder`
    pub const fn new() -> Self {
        EscudeBinScriptBuilder {}
    }
}

impl ScriptBuilder for EscudeBinScriptBuilder {
    fn default_encoding(&self) -> Encoding {
        Encoding::Cp932
    }

    fn build_script(
        &self,
        data: Vec<u8>,
        _filename: &str,
        encoding: Encoding,
        _archive_encoding: Encoding,
        config: &ExtraConfig,
        _archive: Option<&Box<dyn Script>>,
    ) -> Result<Box<dyn Script + Send + Sync>> {
        Ok(Box::new(EscudeBinScript::new(data, encoding, config)?))
    }

    fn extensions(&self) -> &'static [&'static str] {
        &["bin"]
    }

    fn script_type(&self) -> &'static ScriptType {
        &ScriptType::Escude
    }

    fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
        if buf_len > 8 && buf.starts_with(b"ESCR1_00") {
            return Some(255);
        }
        None
    }
}

#[derive(Debug)]
/// Escu:de binary script file
pub struct EscudeBinScript {
    vms: Vec<u8>,
    unk1: u32,
    strings: Vec<String>,
    names: Option<HashMap<usize, String>>,
}

fn load_enum_script(
    filename: &str,
    encoding: Encoding,
    config: &ExtraConfig,
) -> Result<(Vec<NameT>, Vec<VarT>)> {
    let buf = crate::utils::files::read_file(filename)?;
    let scr = EscudeBinList::new(buf, filename, encoding, config)?;
    let mut names = None;
    let mut vars = None;
    for scr in scr.entries {
        match scr.data {
            ListData::Scr(scr) => match scr {
                EnumScr::Names(name) => {
                    names = Some(name);
                }
                EnumScr::Vars(var) => {
                    vars = Some(var);
                }
                _ => {}
            },
            _ => {}
        }
    }
    Ok((
        names.ok_or_else(|| anyhow::anyhow!("No names data in enum script"))?,
        vars.ok_or_else(|| anyhow::anyhow!("No vars data in enum script"))?,
    ))
}

// fn check_messages_in_vms<'a, T: TryInto<usize> + Copy + 'a, I>(messages: &[String], mes: I)
// where
//     I: Iterator<Item = &'a T>,
// {
//     let mess = mes.filter_map(|m| (*m).try_into().ok()).collect::<HashSet<usize>>();
//     for (i, str) in messages.iter().enumerate() {
//         if str.is_empty() {
//             continue;
//         }
//         if !mess.contains(&i) {
//             eprintln!("WARN: Message at index {i} not referenced in VMs: {}", str);
//             crate::COUNTER.inc_warning();
//         }
//     }
// }

impl EscudeBinScript {
    /// Creates a new `EscudeBinScript`
    ///
    /// * `data` - The reader to read the data from
    /// * `encoding` - The encoding of the script
    /// * `config` - Extra configuration options
    pub fn new(data: Vec<u8>, encoding: Encoding, config: &ExtraConfig) -> Result<Self> {
        let mut reader = MemReader::new(data);
        let mut magic = [0u8; 8];
        reader.read_exact(&mut magic)?;
        if &magic != b"ESCR1_00" {
            return Err(anyhow::anyhow!(
                "Invalid Escude binary script magic: {:?}",
                magic
            ));
        }
        let string_count = reader.read_u32()?;
        let mut offsets = Vec::with_capacity(string_count as usize);
        for _ in 0..string_count {
            offsets.push(reader.read_u32()?);
        }
        let vm_count = reader.read_u32()?;
        let mut vms = Vec::with_capacity(vm_count as usize);
        vms.resize(vm_count as usize, 0);
        reader.read_exact(&mut vms)?;
        let unk1 = reader.read_u32()?;
        let mut strings = Vec::with_capacity(string_count as usize);
        if encoding.is_jis() {
            let replaces = StrReplacer::new()?;
            for _ in 0..string_count {
                let s = reader.read_cstring()?;
                let s = replaces.replace(s.as_bytes())?;
                strings.push(decode_to_string(encoding, &s, true)?);
            }
        } else {
            for _ in 0..string_count {
                let s = reader.read_cstring()?;
                strings.push(decode_to_string(encoding, s.as_bytes(), true)?);
            }
        }
        let names = match &config.escude_enum_scr {
            Some(loc) => match load_enum_script(loc, encoding, config) {
                Ok((list, vars)) => match config.escude_op {
                    Some(EscudeOp::Panicon) => {
                        let mut names = HashMap::new();
                        let mut vm = VM::new(&vms);
                        for var in vars {
                            vm.vars.insert(var.value as i32, var.flag as i32);
                        }
                        let _ = vm.run(Some(Box::new(super::ops::panicon::PaniconOps::new())));
                        for (index, name) in vm.names.iter() {
                            if let Some(name) = list.get(*name as usize) {
                                names.insert(*index as usize, name.text.clone());
                            }
                        }
                        // check_messages_in_vms(&strings, vm.mess.iter());
                        Some(names)
                    }
                    Some(EscudeOp::Hanaou) => {
                        let mut names = HashMap::new();
                        let mut vm = VM::new(&vms);
                        for var in vars {
                            vm.vars.insert(var.value as i32, var.flag as i32);
                        }
                        let _ = vm.run(Some(Box::new(super::ops::hanaou::HanaouOps::new())));
                        for (index, name) in vm.names.iter() {
                            if let Some(name) = list.get(*name as usize) {
                                names.insert(*index as usize, name.text.clone());
                            }
                        }
                        // check_messages_in_vms(&strings, vm.mess.iter());
                        Some(names)
                    }
                    None => {
                        let mut names = HashMap::new();
                        let mut vm = VM::new(&vms);
                        for var in vars {
                            vm.vars.insert(var.value as i32, var.flag as i32);
                        }
                        let _ = vm.run(None);
                        for (index, name) in vm.names.iter() {
                            if let Some(name) = list.get(*name as usize) {
                                names.insert(*index as usize, name.text.clone());
                            }
                        }
                        // check_messages_in_vms(&strings, vm.mess.iter());
                        Some(names)
                    }
                },
                Err(e) => {
                    eprintln!(
                        "WARN: Failed to load Escude enum script from {}: {}",
                        loc, e
                    );
                    crate::COUNTER.inc_warning();
                    None
                }
            },
            None => None,
        };
        Ok(EscudeBinScript {
            vms,
            unk1,
            strings,
            names,
        })
    }
}

impl Script for EscudeBinScript {
    fn default_output_script_type(&self) -> OutputScriptType {
        OutputScriptType::Json
    }

    fn default_format_type(&self) -> FormatOptions {
        FormatOptions::None
    }

    fn extract_messages(&self) -> Result<Vec<Message>> {
        Ok(self
            .strings
            .iter()
            .enumerate()
            .map(|(i, s)| Message {
                message: s.replace("<r>", "\n"),
                name: self.names.as_ref().map(|n| n.get(&i).cloned()).flatten(),
            })
            .collect())
    }

    fn import_messages<'a>(
        &'a self,
        messages: Vec<Message>,
        mut writer: Box<dyn WriteSeek + 'a>,
        _filename: &str,
        encoding: Encoding,
        replacement: Option<&'a ReplacementTable>,
    ) -> Result<()> {
        writer.write_all(b"ESCR1_00")?;
        let mut offsets = Vec::with_capacity(messages.len());
        let mut strs = Vec::with_capacity(messages.len());
        let mut len = 0;
        for message in messages {
            offsets.push(len);
            let mut s = message.message.replace("\n", "<r>");
            if let Some(repl) = replacement {
                for (from, to) in &repl.map {
                    s = s.replace(from, to);
                }
            }
            let encoded = encode_string(encoding, &s, false)?;
            len += encoded.len() as u32 + 1;
            strs.push(CString::new(encoded)?);
        }
        writer.write_u32(offsets.len() as u32)?;
        offsets.pack(&mut writer, false, encoding, &None)?;
        writer.write_u32(self.vms.len() as u32)?;
        writer.write_all(&self.vms)?;
        writer.write_u32(self.unk1)?;
        for s in strs {
            writer.write_all(s.as_bytes_with_nul())?;
        }
        Ok(())
    }

    fn is_archive(&self) -> bool {
        false
    }
}

struct StrReplacer {
    pub replacements: HashMap<Vec<u8>, Vec<u8>>,
}

enum JisStr {
    Single(u8),
    Double(u8, u8),
}

impl StrReplacer {
    pub fn new() -> Result<Self> {
        let mut s = StrReplacer {
            replacements: HashMap::new(),
        };
        // 0xa0 to 0xde: Half-width katakana in CP932
        let half_width_katakana = "！？　。「」、…をぁぃぅぇぉゃゅょっーあいうえおかきくけこさしすせそたちつてとなにぬねのはひふへほまみむめもやゆよらりるれろわん゛゜";
        let mut bytes: Vec<u8> = (0xa0..=0xde).collect();
        bytes.insert(0, 0x21);
        bytes.insert(1, 0x3f);
        s.add(&bytes, half_width_katakana)?;
        Ok(s)
    }

    fn add(&mut self, from: &[u8], to: &str) -> Result<()> {
        let encoding = Encoding::Cp932; // Default encoding, can be changed as needed
        let tos = UnicodeSegmentation::graphemes(to, true);
        for (from, to) in from.into_iter().zip(tos) {
            let from_bytes = vec![from.clone()];
            let to_bytes = encode_string(encoding, to, true)?;
            self.replacements.insert(from_bytes, to_bytes);
        }
        Ok(())
    }

    pub fn replace(&self, input: &[u8]) -> Result<Vec<u8>> {
        let mut result = Vec::new();
        let mut reader = MemReaderRef::new(input);
        while let Ok(byte) = reader.read_u8() {
            if byte < 0x80 || (byte >= 0xa0 && byte <= 0xdf) {
                result.push(JisStr::Single(byte));
            } else if (byte >= 0x81 && byte <= 0x9f) || (byte >= 0xe0 && byte <= 0xef) {
                let next_byte = reader.read_u8()?;
                if next_byte < 0x40 || next_byte > 0xfc {
                    return Err(anyhow::anyhow!("Invalid JIS encoding sequence"));
                }
                result.push(JisStr::Double(byte, next_byte));
            } else {
                return Err(anyhow::anyhow!("Invalid byte in JIS encoding: {}", byte));
            }
        }
        let mut output = Vec::new();
        for item in result {
            match item {
                JisStr::Single(byte) => {
                    let vec = vec![byte];
                    if let Some(replacement) = self.replacements.get(&vec) {
                        output.extend_from_slice(replacement);
                    } else {
                        output.push(byte);
                    }
                }
                JisStr::Double(byte1, byte2) => {
                    let key = vec![byte1, byte2];
                    if let Some(replacement) = self.replacements.get(&key) {
                        output.extend_from_slice(replacement);
                    } else {
                        output.push(byte1);
                        output.push(byte2);
                    }
                }
            }
        }
        Ok(output)
    }
}

#[repr(u8)]
#[derive(Debug, IntEnum)]
enum BaseOp {
    End,
    Jump,
    JumpZ,
    Call,
    Ret,
    Push,
    Pop,
    Str,
    SetVar,
    GetVar,
    SetFlag,
    GetFlag,
    Neg,
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    Not,
    And,
    Or,
    Xor,
    Shr,
    Shl,
    Eq,
    Ne,
    Gt,
    Ge,
    Lt,
    Le,
    LNot,
    LAnd,
    LOr,
    FileLine,
}

pub(crate) trait ReadParam<T> {
    fn read_param(&mut self) -> Result<T>;
}

#[derive(Debug)]
pub(crate) struct VM<'a, T: std::fmt::Debug> {
    pub reader: MemReaderRef<'a>,
    pub data: Vec<T>,
    pub stack: Vec<u64>,
    pub strs: Vec<T>,
    pub vars: HashMap<T, T>,
    pub flags: HashMap<T, bool>,
    pub mess: BTreeSet<T>,
    pub names: HashMap<T, T>,
}

impl ReadParam<i32> for MemReaderRef<'_> {
    fn read_param(&mut self) -> Result<i32> {
        Ok(self.read_i32()?)
    }
}

impl<'a, T> VM<'a, T>
where
    MemReaderRef<'a>: ReadParam<T>,
    T: TryInto<u64>
        + Default
        + Eq
        + Ord
        + Copy
        + std::fmt::Debug
        + std::fmt::Display
        + std::hash::Hash
        + From<u8>
        + std::ops::Neg<Output = T>
        + std::ops::Add<Output = T>
        + std::ops::Sub<Output = T>
        + std::ops::Mul<Output = T>
        + std::ops::Div<Output = T>
        + std::ops::Rem<Output = T>
        + std::ops::Not<Output = T>
        + std::ops::BitAnd<Output = T>
        + std::ops::BitOr<Output = T>
        + std::ops::BitXor<Output = T>
        + std::ops::Shr<Output = T>
        + std::ops::Shl<Output = T>,
    anyhow::Error: From<<T as TryInto<u64>>::Error>,
{
    pub fn new(data: &'a [u8]) -> Self {
        VM {
            reader: MemReaderRef::new(data),
            data: Vec::new(),
            stack: Vec::new(),
            strs: Vec::new(),
            vars: HashMap::new(),
            flags: HashMap::new(),
            mess: BTreeSet::new(),
            names: HashMap::new(),
        }
    }

    pub fn pop_data(&mut self) -> Result<T> {
        self.data
            .pop()
            .ok_or_else(|| anyhow::anyhow!("No data to pop"))
    }

    fn pop_stack(&mut self) -> Result<u64> {
        self.stack
            .pop()
            .ok_or_else(|| anyhow::anyhow!("No stack to pop"))
    }

    pub fn run(&mut self, mut custom_ops: Option<Box<dyn CustomOps<T>>>) -> Result<()> {
        loop {
            if self.reader.is_eof() {
                break;
            }
            let op = self.reader.read_u8()?;
            if let Ok(op) = BaseOp::try_from(op) {
                // println!("Op code: {op:?}");
                match op {
                    BaseOp::End => break,
                    BaseOp::Jump => {
                        let offset: T = self.reader.read_param()?;
                        let offset: u64 = offset.try_into()?;
                        self.reader.seek(SeekFrom::Start(offset))?;
                    }
                    BaseOp::JumpZ => {
                        let offset: T = self.reader.read_param()?;
                        let offset: u64 = offset.try_into()?;
                        if self.pop_data()? == Default::default() {
                            self.reader.seek(SeekFrom::Start(offset))?;
                        }
                    }
                    BaseOp::Call => {
                        let offset: T = self.reader.read_param()?;
                        let offset: u64 = offset.try_into()?;
                        let pos = self.reader.stream_position()?;
                        self.stack.push(pos);
                        self.reader.seek(SeekFrom::Start(offset))?;
                    }
                    BaseOp::Ret => {
                        if self.stack.is_empty() {
                            let code = self.reader.read_u8()?;
                            if code == 0 && self.reader.is_eof() {
                                break;
                            }
                        }
                        let stack = self.pop_stack()?;
                        self.reader.seek(SeekFrom::Start(stack))?;
                    }
                    BaseOp::Push => {
                        let d = self.reader.read_param()?;
                        self.data.push(d);
                    }
                    BaseOp::Pop => {
                        self.pop_data()?;
                    }
                    BaseOp::Str => {
                        let param = self.reader.read_param()?;
                        self.strs.push(param);
                        self.data.push(param);
                    }
                    BaseOp::SetVar => {
                        let value = self.pop_data()?;
                        let index = self.pop_data()?;
                        self.vars.insert(index, value);
                        self.data.push(value);
                    }
                    BaseOp::GetVar => {
                        let index = self.pop_data()?;
                        let value = self
                            .vars
                            .get(&index)
                            .ok_or_else(|| anyhow::anyhow!("Variable not found: {}", index))?;
                        self.data.push(*value);
                    }
                    BaseOp::SetFlag => {
                        let value = self.pop_data()?;
                        let index = self.pop_data()?;
                        let flag = value != Default::default();
                        self.flags.insert(index, flag);
                        self.data.push(value);
                    }
                    BaseOp::GetFlag => {
                        let index = self.pop_data()?;
                        let flag = self.flags.get(&index).cloned().unwrap_or(false);
                        self.data
                            .push(if flag { T::from(1u8) } else { T::from(0u8) });
                    }
                    BaseOp::Neg => {
                        let value = -self.pop_data()?;
                        self.data.push(value);
                    }
                    BaseOp::Add => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data.push(a + b);
                    }
                    BaseOp::Sub => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data.push(a - b);
                    }
                    BaseOp::Mul => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data.push(a * b);
                    }
                    BaseOp::Div => {
                        let b = self.pop_data()?;
                        if b == Default::default() {
                            return Err(anyhow::anyhow!("Division by zero"));
                        }
                        let a = self.pop_data()?;
                        self.data.push(a / b);
                    }
                    BaseOp::Mod => {
                        let b = self.pop_data()?;
                        if b == Default::default() {
                            return Err(anyhow::anyhow!("Division by zero"));
                        }
                        let a = self.pop_data()?;
                        self.data.push(a % b);
                    }
                    BaseOp::Not => {
                        let value = self.pop_data()?;
                        self.data.push(!value);
                    }
                    BaseOp::And => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data.push(a & b);
                    }
                    BaseOp::Or => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data.push(a | b);
                    }
                    BaseOp::Xor => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data.push(a ^ b);
                    }
                    BaseOp::Shr => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data.push(a >> b);
                    }
                    BaseOp::Shl => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data.push(a << b);
                    }
                    BaseOp::Eq => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data
                            .push(if a == b { T::from(1u8) } else { T::from(0u8) });
                    }
                    BaseOp::Ne => {
                        let b = self.pop_data()?;
                        let a = self.pop_data()?;
                        self.data
                            .push(if a != b { T::from(1u8) } else { T::from(0u8) });
                    }
                    // Original code may contains undefined behavior for these operations
                    BaseOp::Gt => {
                        let a = self.pop_data()?;
                        let b = self.pop_data()?;
                        self.data
                            .push(if a > b { T::from(1u8) } else { T::from(0u8) });
                    }
                    BaseOp::Ge => {
                        let a = self.pop_data()?;
                        let b = self.pop_data()?;
                        self.data
                            .push(if a >= b { T::from(1u8) } else { T::from(0u8) });
                    }
                    BaseOp::Lt => {
                        let a = self.pop_data()?;
                        let b = self.pop_data()?;
                        self.data
                            .push(if a < b { T::from(1u8) } else { T::from(0u8) });
                    }
                    BaseOp::Le => {
                        let a = self.pop_data()?;
                        let b = self.pop_data()?;
                        self.data
                            .push(if a <= b { T::from(1u8) } else { T::from(0u8) });
                    }
                    BaseOp::LNot => {
                        let value = self.pop_data()?;
                        self.data.push(if value == Default::default() {
                            T::from(1u8)
                        } else {
                            T::from(0u8)
                        });
                    }
                    BaseOp::LAnd => {
                        let b = self.pop_data()? != Default::default();
                        let a = self.pop_data()? != Default::default();
                        self.data
                            .push(if a && b { T::from(1u8) } else { T::from(0u8) });
                    }
                    BaseOp::LOr => {
                        let b = self.pop_data()? != Default::default();
                        let a = self.pop_data()? != Default::default();
                        self.data
                            .push(if a || b { T::from(1u8) } else { T::from(0u8) });
                    }
                    BaseOp::FileLine => {
                        let _: T = self.reader.read_param()?;
                    }
                }
                continue;
            }
            if let Some(ops) = &mut custom_ops {
                let nbreak = ops.run(self, op)?;
                if nbreak {
                    break;
                }
            } else {
                return Err(anyhow::anyhow!("Unknown operation: {}", op));
            }
        }
        Ok(())
    }

    pub fn skip_n_params(&mut self, n: u64, nbreak: bool) -> Result<bool> {
        if (self.data.len() as u64) < n {
            println!("{:?}", self.data);
        }
        for _ in 0..n {
            self.pop_data()?;
        }
        Ok(nbreak)
    }

    pub fn skip_params(&mut self, nbreak: bool) -> Result<bool> {
        let count: T = self.reader.read_param()?;
        let count: u64 = count.try_into()?;
        self.skip_n_params(count, nbreak)
    }

    pub fn read_params(&mut self, ncount: Option<u64>) -> Result<Vec<T>> {
        let count = match ncount {
            Some(count) => count,
            None => {
                let count: T = self.reader.read_param()?;
                count.try_into()?
            }
        };
        let data_len = self.data.len();
        if (data_len as u64) < count {
            return Err(anyhow::anyhow!(
                "Not enough data to read {} parameters, only {} parameters available",
                count,
                data_len
            ));
        }
        let mut params = Vec::with_capacity(count as usize);
        params.resize(count as usize, Default::default());
        params.copy_from_slice(&self.data[data_len - count as usize..]);
        self.data.truncate(data_len - count as usize);
        Ok(params)
    }
}
