use std::fmt::Write;

use colored::Colorize;
use fix_builtins::BuiltinId;
use num_enum::TryFromPrimitive;

use crate::{InstructionPtr, Op, OperandType};

pub trait DisassemblerContext {
    fn resolve_string(&self, id: u32) -> &str;
    fn get_code(&self) -> &[u8];
}

pub struct Disassembler<'a, Ctx> {
    code: &'a [u8],
    ctx: &'a Ctx,
    pc: usize,
}

impl<'a, Ctx: DisassemblerContext> Disassembler<'a, Ctx> {
    pub fn new(ip: InstructionPtr, ctx: &'a Ctx) -> Self {
        Self {
            code: ctx.get_code(),
            ctx,
            pc: ip.0,
        }
    }

    #[inline(always)]
    fn read_u8(&mut self) -> u8 {
        let b = self.code[self.pc];
        self.pc += 1;
        b
    }

    #[inline(always)]
    fn read_u16(&mut self) -> u16 {
        let bytes = self.code[self.pc..self.pc + 2]
            .try_into()
            .expect("no enough bytes");
        self.pc += 2;
        u16::from_le_bytes(bytes)
    }

    #[inline(always)]
    fn read_u32(&mut self) -> u32 {
        let bytes = self.code[self.pc..self.pc + 4]
            .try_into()
            .expect("no enough bytes");
        self.pc += 4;
        u32::from_le_bytes(bytes)
    }

    #[inline(always)]
    fn read_i32(&mut self) -> i32 {
        let bytes = self.code[self.pc..self.pc + 4]
            .try_into()
            .expect("no enough bytes");
        self.pc += 4;
        i32::from_le_bytes(bytes)
    }

    #[inline(always)]
    fn read_i64(&mut self) -> i64 {
        let bytes = self.code[self.pc..self.pc + 8]
            .try_into()
            .expect("no enough bytes");
        self.pc += 8;
        i64::from_le_bytes(bytes)
    }

    #[inline(always)]
    fn read_f64(&mut self) -> f64 {
        let bytes = self.code[self.pc..self.pc + 8]
            .try_into()
            .expect("no enough bytes");
        self.pc += 8;
        f64::from_le_bytes(bytes)
    }

    #[inline(always)]
    fn read_operand_data(&mut self) {
        use OperandType::*;
        let tag = self.read_u8();
        let ty = OperandType::try_from_primitive(tag).expect("invalid operand type");
        match ty {
            Const => {
                self.read_u32();
            }
            BigInt => {
                self.read_i64();
            }
            Local => {
                self.read_u8();
                self.read_u32();
            }
            BuiltinConst => {
                self.read_u32();
            }
            Builtins => {}
            ReplBinding | ScopedImportBinding => {
                self.read_u32();
            }
        }
    }

    pub fn disassemble(&mut self) -> String {
        self.disassemble_impl(false)
    }

    pub fn disassemble_colored(&mut self) -> String {
        self.disassemble_impl(true)
    }

    fn disassemble_impl(&mut self, color: bool) -> String {
        let mut out = String::new();
        if color {
            let _ = writeln!(out, "{}", "=== Bytecode Disassembly ===".bold().white());
            let _ = writeln!(
                out,
                "{} {}",
                "Length:".white(),
                format!("{} bytes", self.code.len()).cyan()
            );
        } else {
            let _ = writeln!(out, "=== Bytecode Disassembly ===");
            let _ = writeln!(out, "Length: {} bytes", self.code.len());
        }

        while self.pc < self.code.len() {
            let start_pos = self.pc;
            let op_byte = self.read_u8();
            let (mnemonic, args) = self.decode_instruction(op_byte, start_pos);

            let bytes_slice = &self.code[start_pos + 1..self.pc];
            let mut chunks = bytes_slice.chunks(4);

            let first_chunk = chunks.next().unwrap_or(&[]);
            let bytes_str = {
                let mut temp = format!("{:02x}", self.code[start_pos]);
                for b in first_chunk {
                    let _ = write!(&mut temp, " {:02x}", b);
                }
                temp
            };

            if color {
                let sep = if args.is_empty() { "" } else { " " };
                let _ = writeln!(
                    out,
                    "{}  {:<14} | {}{}{}",
                    format!("{:04x}", start_pos).dimmed(),
                    bytes_str.green(),
                    mnemonic.yellow().bold(),
                    sep,
                    args.cyan()
                );
            } else {
                let op_str = if args.is_empty() {
                    mnemonic.to_string()
                } else {
                    format!("{} {}", mnemonic, args)
                };
                let _ = writeln!(out, "{:04x}  {:<14} | {}", start_pos, bytes_str, op_str);
            }

            for chunk in chunks {
                let bytes_str = {
                    let mut temp = String::from("  ");
                    for b in chunk {
                        let _ = write!(&mut temp, " {:02x}", b);
                    }
                    temp
                };

                let extra_width = if start_pos > 0 {
                    start_pos.ilog2() >> 4
                } else {
                    0
                };

                if color {
                    let _ = write!(out, "    ");
                    for _ in 0..extra_width {
                        let _ = write!(out, " ");
                    }
                    let _ = writeln!(out, "  {:<14} |", bytes_str.green());
                } else {
                    let _ = write!(out, "    ");
                    for _ in 0..extra_width {
                        let _ = write!(out, " ");
                    }
                    let _ = writeln!(out, "  {:<14} |", bytes_str);
                }
            }
        }
        out
    }

    fn decode_instruction(&mut self, op_byte: u8, current_pc: usize) -> (&'static str, String) {
        let op = Op::try_from_primitive(op_byte).expect("invalid op code");

        match op {
            Op::PushSmi => {
                let val = self.read_i32();
                ("PushSmi", format!("{}", val))
            }
            Op::PushBigInt => {
                let val = self.read_i64();
                ("PushBigInt", format!("{}", val))
            }
            Op::PushFloat => {
                let val = self.read_f64();
                ("PushFloat", format!("{}", val))
            }
            Op::PushString => {
                let idx = self.read_u32();
                let s = self.ctx.resolve_string(idx);
                let len = s.len();
                let mut s_fmt = format!("{:?}", s);
                if s_fmt.len() > 60 {
                    s_fmt.truncate(57);
                    #[allow(clippy::unwrap_used)]
                    write!(s_fmt, "...\" (total {len} bytes)").unwrap();
                }
                ("PushString", format!("@{} {}", idx, s_fmt))
            }
            Op::PushNull => ("PushNull", String::new()),
            Op::PushTrue => ("PushTrue", String::new()),
            Op::PushFalse => ("PushFalse", String::new()),

            Op::LoadLocal => {
                let idx = self.read_u32();
                ("LoadLocal", format!("[{}]", idx))
            }
            Op::LoadOuter => {
                let depth = self.read_u8();
                let idx = self.read_u32();
                ("LoadOuter", format!("depth={} [{}]", depth, idx))
            }
            Op::StoreLocal => {
                let idx = self.read_u32();
                ("StoreLocal", format!("[{}]", idx))
            }
            Op::AllocLocals => {
                let count = self.read_u32();
                ("AllocLocals", format!("count={}", count))
            }

            Op::MakeThunk => {
                let offset = self.read_u32();
                ("MakeThunk", format!("-> {:04x}", offset))
            }
            Op::MakeClosure => {
                let offset = self.read_u32();
                let slots = self.read_u32();
                ("MakeClosure", format!("-> {:04x} slots={}", offset, slots))
            }
            Op::MakePatternClosure => {
                let offset = self.read_u32();
                let slots = self.read_u32();
                let req_count = self.read_u16();
                let opt_count = self.read_u16();
                let ellipsis = self.read_u8() != 0;

                let mut arg_str = format!(
                    "-> {:04x} slots={} req={} opt={} ...={})",
                    offset, slots, req_count, opt_count, ellipsis
                );

                arg_str.push_str(" Args=[");
                for _ in 0..req_count {
                    let idx = self.read_u32();
                    arg_str.push_str(&format!("Req({}) ", self.ctx.resolve_string(idx)));
                }
                for _ in 0..opt_count {
                    let idx = self.read_u32();
                    arg_str.push_str(&format!("Opt({}) ", self.ctx.resolve_string(idx)));
                }

                let total_args = req_count + opt_count;
                for _ in 0..total_args {
                    let _name_idx = self.read_u32();
                    let _span_id = self.read_u32();
                }
                arg_str.push(']');

                ("MakePatternClosure", arg_str)
            }

            Op::Call => {
                self.read_operand_data();
                ("Call", "arg=?".into())
            }
            Op::DispatchPrimOp => {
                let id = BuiltinId::try_from_primitive(self.read_u8()).expect("invalid builtin id");
                ("DispatchPrimOp", format!("id={id:?}"))
            }

            Op::MakeAttrs => {
                let static_count = self.read_u32();
                let dynamic_count = self.read_u32();
                let mut args = format!("static={} dynamic={}", static_count, dynamic_count);

                for _ in 0..static_count {
                    let key_id = self.read_u32();
                    let _ = write!(args, " [{}={}", self.ctx.resolve_string(key_id), key_id);
                    self.read_operand_data();
                    let _span_id = self.read_u32();
                    args.push(']');
                }

                for _ in 0..dynamic_count {
                    let _ = write!(args, " [dyn");
                    self.read_operand_data();
                    let _span_id = self.read_u32();
                    args.push(']');
                }

                ("MakeAttrs", args)
            }
            Op::MakeEmptyAttrs => ("MakeEmptyAttrs", String::new()),

            Op::SelectStatic => {
                let span_id = self.read_u32();
                let key_id = self.read_u32();
                (
                    "SelectStatic",
                    format!("key={} span={}", self.ctx.resolve_string(key_id), span_id),
                )
            }
            Op::SelectDynamic => {
                let span_id = self.read_u32();
                ("SelectDynamic", format!("span={}", span_id))
            }
            Op::HasAttrPathStatic => {
                let span_id = self.read_u32();
                let key_id = self.read_u32();
                (
                    "HasAttrPathStatic",
                    format!("key={} span={}", self.ctx.resolve_string(key_id), span_id),
                )
            }
            Op::HasAttrPathDynamic => {
                let span_id = self.read_u32();
                ("HasAttrPathDynamic", format!("span={}", span_id))
            }
            Op::HasAttrStatic => {
                let key_id = self.read_u32();
                (
                    "HasAttrStatic",
                    format!("key={}", self.ctx.resolve_string(key_id)),
                )
            }
            Op::HasAttrDynamic => ("HasAttrDynamic", String::new()),
            Op::HasAttrResolve => ("HasAttrResolve", String::new()),
            Op::JumpIfSelectSucceeded => {
                let offset = self.read_i32();
                let target = (current_pc as isize + 1 + 4 + offset as isize) as usize;
                (
                    "JumpIfSelectSucceeded",
                    format!("-> {:04x} offset={}", target, offset),
                )
            }
            Op::JumpIfSelectFailed => {
                let offset = self.read_i32();
                let target = (current_pc as isize + 1 + 4 + offset as isize) as usize;
                (
                    "JumpIfSelectFailed",
                    format!("-> {:04x} offset={}", target, offset),
                )
            }

            Op::MakeList => {
                let count = self.read_u32();
                for _ in 0..count {
                    self.read_operand_data();
                }
                ("MakeList", format!("size={}", count))
            }
            Op::MakeEmptyList => ("MakeEmptyList", String::new()),

            Op::OpAdd => ("OpAdd", String::new()),
            Op::OpSub => ("OpSub", String::new()),
            Op::OpMul => ("OpMul", String::new()),
            Op::OpDiv => ("OpDiv", String::new()),
            Op::OpEq => ("OpEq", String::new()),
            Op::OpNeq => ("OpNeq", String::new()),
            Op::OpLt => ("OpLt", String::new()),
            Op::OpGt => ("OpGt", String::new()),
            Op::OpLeq => ("OpLeq", String::new()),
            Op::OpGeq => ("OpGeq", String::new()),
            Op::OpConcat => ("OpConcat", String::new()),
            Op::OpUpdate => ("OpUpdate", String::new()),
            Op::OpNeg => ("OpNeg", String::new()),
            Op::OpNot => ("OpNot", String::new()),

            Op::JumpIfFalse => {
                let offset = self.read_i32();
                let target = (current_pc as isize + 1 + 4 + offset as isize) as usize;
                (
                    "JumpIfFalse",
                    format!("-> {:04x} offset={}", target, offset),
                )
            }
            Op::JumpIfTrue => {
                let offset = self.read_i32();
                let target = (current_pc as isize + 1 + 4 + offset as isize) as usize;
                ("JumpIfTrue", format!("-> {:04x} offset={}", target, offset))
            }
            Op::Jump => {
                let offset = self.read_i32();
                let target = (current_pc as isize + 1 + 4 + offset as isize) as usize;
                ("Jump", format!("-> {:04x} offset={}", target, offset))
            }

            Op::ConcatStrings => {
                let count = self.read_u16();
                let force = self.read_u8();
                ("ConcatStrings", format!("count={} force={}", count, force))
            }
            Op::CoerceToString => ("CoerceToString", String::new()),
            Op::ResolvePath => ("ResolvePath", String::new()),
            Op::Assert => {
                let raw_idx = self.read_u32();
                let span_id = self.read_u32();
                ("Assert", format!("text_id={} span={}", raw_idx, span_id))
            }
            Op::LookupWith => {
                let idx = self.read_u32();
                let name = self.ctx.resolve_string(idx);
                let n = self.read_u8();
                for _ in 0..n {
                    self.read_operand_data();
                }
                ("LookupWith", format!("sym={:?} n={}", name, n))
            }

            Op::LoadBuiltins => ("LoadBuiltins", String::new()),
            Op::LoadBuiltin => {
                let id = self.read_u8();
                ("LoadBuiltin", format!("id={}", id))
            }
            Op::LoadReplBinding => {
                let idx = self.read_u32();
                let name = self.ctx.resolve_string(idx);
                ("LoadReplBinding", format!("{:?}", name))
            }
            Op::LoadScopedBinding => {
                let idx = self.read_u32();
                let name = self.ctx.resolve_string(idx);
                ("LoadScopedBinding", format!("{:?}", name))
            }
            Op::Return => ("Return", String::new()),
            Op::Illegal => ("Illegal", String::new()),
        }
    }
}