#![allow(dead_code)]

use fix_lang::{BuiltinId, StringId};
use num_enum::TryFromPrimitive;
use string_interner::Symbol as _;

pub mod disassembler;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct InstructionPtr(pub usize);

#[repr(u8)]
#[derive(Debug, Clone, Copy, TryFromPrimitive)]
#[allow(clippy::enum_variant_names)]
pub enum Op {
    PushSmi,
    PushBigInt,
    PushFloat,
    PushString,
    PushNull,
    PushTrue,
    PushFalse,

    LoadLocal,
    LoadOuter,
    StoreLocal,
    AllocLocals,

    MakeThunk,
    MakeClosure,
    MakePatternClosure,

    Call,
    DispatchCont,

    MakeAttrs,
    MakeEmptyAttrs,
    SelectStatic,
    SelectDynamic,
    HasAttrPathStatic,
    HasAttrPathDynamic,
    HasAttrStatic,
    HasAttrDynamic,
    HasAttrResolve,
    JumpIfSelectSucceeded,
    JumpIfSelectFailed,

    MakeList,
    MakeEmptyList,

    OpAdd,
    OpSub,
    OpMul,
    OpDiv,
    OpEq,
    OpNeq,
    OpLt,
    OpGt,
    OpLeq,
    OpGeq,
    OpConcat,
    OpUpdate,

    OpNeg,
    OpNot,

    JumpIfFalse,
    JumpIfTrue,
    Jump,

    CoerceToString,

    ConcatStrings,
    ResolvePath,

    Assert,

    LookupWith,

    LoadBuiltins,
    LoadBuiltin,

    LoadReplBinding,
    LoadScopedBinding,

    Return,

    Illegal,
}

#[repr(u8)]
#[derive(Debug, Clone, Copy, TryFromPrimitive)]
pub enum OperandType {
    Const,
    BigInt,
    Local,
    BuiltinConst,
    Builtins,
    ReplBinding,
    ScopedImportBinding,
}

pub enum Const {
    Smi(i32),
    Float(f64),
    Bool(bool),
    String(StringId),
    Path(StringId),
    PrimOp {
        id: BuiltinId,
        arity: u8,
        dispatch_ip: u32,
    },
    Null,
}

#[repr(u8)]
#[derive(Debug, Clone, Copy, TryFromPrimitive)]
pub enum AttrKeyType {
    Static,
    Dynamic,
}

pub enum OperandData {
    Const(u32),
    BigInt(i64),
    Local { layer: u8, idx: u32 },
    BuiltinConst(StringId),
    Builtins,
    ReplBinding(StringId),
    ScopedImportBinding { slot_id: u32, name: StringId },
}

#[repr(u8)]
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum Continuation {
    // primops
    PAbort,
    PAdd,
    PAddErrorContext,

    PAll,
    PAllCallPred,
    PAllCheck,

    PAny,
    PAnyCallPred,
    PAnyCheck,

    PAppendContext,
    PAppendContextLoop,
    PAppendContextEntryForced,
    PAppendContextOutputsForced,
    PAppendContextOutputElementLoop,
    PAppendContextOutputElementForced,

    PAttrNames,
    PAttrValues,
    PBaseNameOf,
    PBitAnd,
    PBitOr,
    PBitXor,
    PBreak,
    PCatAttrs,
    PCeil,
    PCompareVersions,
    PConcatLists,
    PConcatMap,
    PConcatStringsSep,
    PConvertHash,

    PDeepSeq,
    PDeepSeqPush,
    PDeepSeqLoop,

    PDerivation,
    PDerivationStrict,
    PDirOf,
    PDiv,
    PElem,
    PElemAt,
    PFetchGit,
    PFetchMercurial,
    PFetchTarball,
    PFetchTree,
    PFetchUrl,

    PFilterForceList,
    PFilterCallPred,
    PFilterCheck,

    PFilterSource,
    PFindFile,
    PFloor,
    PFoldlStrict,
    PFoldlStrictEmpty,
    PFoldlStrictCall1,
    PFoldlStrictCall2,
    PFoldlStrictUpdate,
    PFromJSON,
    PFromTOML,
    PFunctionArgs,
    PGenList,
    PGenericClosure,
    PGetAttr,
    PGetContext,
    PGetEnv,
    PGroupBy,
    PHasAttr,
    PHasContext,
    PHashFile,
    PHashString,
    PHead,

    PImport,
    PImportFinalize,
    PScopedImport,
    PScopedImportFinalize,

    PIntersectAttrs,
    PIsAttrs,
    PIsBool,
    PIsFloat,
    PIsFunction,
    PIsInt,
    PIsList,
    PIsNull,
    PIsPath,
    PIsString,
    PLength,
    PLessThan,
    PListToAttrs,
    PMap,
    PMapAttrs,
    PMatch,
    PMul,
    PParseDrvName,
    PPartition,
    PPath,
    PPathExists,
    PPlaceholder,
    PReadDir,
    PReadFile,
    PReadFileType,
    PRemoveAttrs,
    PReplaceStrings,
    PSeq,
    PSort,
    PSplit,
    PSplitVersion,
    PStorePath,
    PStringLength,
    PSub,
    PSubstring,
    PTail,
    PThrow,
    PToFile,
    PToJSON,
    PToPath,
    PToString,
    PToXML,
    PTrace,
    PTryEval,
    PTypeOf,
    PUnsafeDiscardStringContext,
    PUnsafeGetAttrPos,
    PWarn,
    PZipAttrsWith,
    PUnsafeDiscardOutputDependency,

    ForceResultShallow,
    ForceResultShallowPush,
    ForceResultShallowLoop,
    ForceResultDeepFinish,

    EqStep,
    EqForce,

    CallPattern,
    CallFunctor1,
    CallFunctor2,

    Illegal,
}

impl TryFrom<u8> for Continuation {
    type Error = u8;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        if (0..Self::Illegal as u8).contains(&value) {
            Ok(unsafe { std::mem::transmute::<u8, Self>(value) })
        } else {
            Err(value)
        }
    }
}

impl Continuation {
    pub fn entry_for_builtin(id: BuiltinId) -> Self {
        use BuiltinId::*;
        match id {
            Abort => Self::PAbort,
            Add => Self::PAdd,
            AddErrorContext => Self::PAddErrorContext,
            All => Self::PAll,
            Any => Self::PAny,
            AppendContext => Self::PAppendContext,
            AttrNames => Self::PAttrNames,
            AttrValues => Self::PAttrValues,
            BaseNameOf => Self::PBaseNameOf,
            BitAnd => Self::PBitAnd,
            BitOr => Self::PBitOr,
            BitXor => Self::PBitXor,
            Break => Self::PBreak,
            CatAttrs => Self::PCatAttrs,
            Ceil => Self::PCeil,
            CompareVersions => Self::PCompareVersions,
            ConcatLists => Self::PConcatLists,
            ConcatMap => Self::PConcatMap,
            ConcatStringsSep => Self::PConcatStringsSep,
            ConvertHash => Self::PConvertHash,
            DeepSeq => Self::PDeepSeq,
            Derivation => Self::PDerivation,
            DerivationStrict => Self::PDerivationStrict,
            DirOf => Self::PDirOf,
            Div => Self::PDiv,
            Elem => Self::PElem,
            ElemAt => Self::PElemAt,
            FetchGit => Self::PFetchGit,
            FetchMercurial => Self::PFetchMercurial,
            FetchTarball => Self::PFetchTarball,
            FetchTree => Self::PFetchTree,
            FetchUrl => Self::PFetchUrl,
            Filter => Self::PFilterForceList,
            FilterSource => Self::PFilterSource,
            FindFile => Self::PFindFile,
            Floor => Self::PFloor,
            FoldlStrict => Self::PFoldlStrict,
            FromJSON => Self::PFromJSON,
            FromTOML => Self::PFromTOML,
            FunctionArgs => Self::PFunctionArgs,
            GenList => Self::PGenList,
            GenericClosure => Self::PGenericClosure,
            GetAttr => Self::PGetAttr,
            GetContext => Self::PGetContext,
            GetEnv => Self::PGetEnv,
            GroupBy => Self::PGroupBy,
            HasAttr => Self::PHasAttr,
            HasContext => Self::PHasContext,
            HashFile => Self::PHashFile,
            HashString => Self::PHashString,
            Head => Self::PHead,
            Import => Self::PImport,
            IntersectAttrs => Self::PIntersectAttrs,
            IsAttrs => Self::PIsAttrs,
            IsBool => Self::PIsBool,
            IsFloat => Self::PIsFloat,
            IsFunction => Self::PIsFunction,
            IsInt => Self::PIsInt,
            IsList => Self::PIsList,
            IsNull => Self::PIsNull,
            IsPath => Self::PIsPath,
            IsString => Self::PIsString,
            Length => Self::PLength,
            LessThan => Self::PLessThan,
            ListToAttrs => Self::PListToAttrs,
            Map => Self::PMap,
            MapAttrs => Self::PMapAttrs,
            Match => Self::PMatch,
            Mul => Self::PMul,
            ParseDrvName => Self::PParseDrvName,
            Partition => Self::PPartition,
            Path => Self::PPath,
            PathExists => Self::PPathExists,
            Placeholder => Self::PPlaceholder,
            ReadDir => Self::PReadDir,
            ReadFile => Self::PReadFile,
            ReadFileType => Self::PReadFileType,
            RemoveAttrs => Self::PRemoveAttrs,
            ReplaceStrings => Self::PReplaceStrings,
            ScopedImport => Self::PScopedImport,
            Seq => Self::PSeq,
            Sort => Self::PSort,
            Split => Self::PSplit,
            SplitVersion => Self::PSplitVersion,
            StorePath => Self::PStorePath,
            StringLength => Self::PStringLength,
            Sub => Self::PSub,
            Substring => Self::PSubstring,
            Tail => Self::PTail,
            Throw => Self::PThrow,
            ToFile => Self::PToFile,
            ToJSON => Self::PToJSON,
            ToPath => Self::PToPath,
            ToString => Self::PToString,
            ToXML => Self::PToXML,
            Trace => Self::PTrace,
            TryEval => Self::PTryEval,
            TypeOf => Self::PTypeOf,
            UnsafeDiscardStringContext => Self::PUnsafeDiscardStringContext,
            UnsafeDiscardOutputDependency => Self::PUnsafeDiscardOutputDependency,
            UnsafeGetAttrPos => Self::PUnsafeGetAttrPos,
            Warn => Self::PWarn,
            ZipAttrsWith => Self::PZipAttrsWith,
        }
    }

    pub fn ip(self) -> u32 {
        self as u32 * 2
    }
}

pub struct BytecodeReader<'a> {
    bytecode: &'a [u8],
    pc: usize,
    inst_start_pc: usize,
}

impl<'a> BytecodeReader<'a> {
    pub fn new(bytecode: &'a [u8], pc: usize) -> Self {
        Self {
            bytecode,
            pc,
            inst_start_pc: pc,
        }
    }

    #[inline(always)]
    pub fn from_after_op(bytecode: &'a [u8], inst_start_pc: usize) -> Self {
        Self {
            bytecode,
            pc: inst_start_pc + 1,
            inst_start_pc,
        }
    }

    #[inline(always)]
    #[cfg_attr(debug_assertions, track_caller)]
    fn read_array<const N: usize>(&mut self) -> [u8; N] {
        let ret = self.bytecode[self.pc..self.pc + N]
            .try_into()
            .expect("read_array failed");
        self.pc += N;
        ret
    }

    #[inline(always)]
    pub fn read_op(&mut self) -> Op {
        self.inst_start_pc = self.pc;
        let byte = self.bytecode[self.pc];
        if !(0..Op::Illegal as u8).contains(&byte) {
            std::hint::cold_path();
            panic!("unknown opcode: {byte:#04x}")
        }
        self.pc += 1;
        unsafe { std::mem::transmute::<u8, Op>(byte) }
    }

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

    #[inline(always)]
    pub fn read_u16(&mut self) -> u16 {
        u16::from_le_bytes(self.read_array())
    }

    #[inline(always)]
    pub fn read_u32(&mut self) -> u32 {
        u32::from_le_bytes(self.read_array())
    }

    #[inline(always)]
    pub fn read_i32(&mut self) -> i32 {
        i32::from_le_bytes(self.read_array())
    }

    #[inline(always)]
    pub fn read_i64(&mut self) -> i64 {
        i64::from_le_bytes(self.read_array())
    }

    #[inline(always)]
    pub fn read_f64(&mut self) -> f64 {
        f64::from_le_bytes(self.read_array())
    }

    #[inline(always)]
    pub fn read_string_id(&mut self) -> StringId {
        let raw = self.read_u32();
        #[allow(clippy::unwrap_used)]
        StringId(string_interner::symbol::SymbolU32::try_from_usize(raw as usize).unwrap())
    }

    #[inline(always)]
    pub fn read_operand_data(&mut self) -> OperandData {
        let tag = self.read_u8();
        let Ok(ty) = OperandType::try_from_primitive(tag)
            .map_err(|err| panic!("unknown operand tag: {:#04x}", err.number));
        match ty {
            OperandType::Const => OperandData::Const(self.read_u32()),
            OperandType::BigInt => OperandData::BigInt(self.read_i64()),
            OperandType::Local => {
                let layer = self.read_u8();
                let idx = self.read_u32();
                OperandData::Local { layer, idx }
            }
            OperandType::BuiltinConst => OperandData::BuiltinConst(self.read_string_id()),
            OperandType::Builtins => OperandData::Builtins,
            OperandType::ReplBinding => OperandData::ReplBinding(self.read_string_id()),
            OperandType::ScopedImportBinding => {
                let slot_id = self.read_u32();
                let name = self.read_string_id();
                OperandData::ScopedImportBinding { slot_id, name }
            }
        }
    }

    pub fn pc(&self) -> usize {
        self.pc
    }

    pub fn set_pc(&mut self, pc: usize) {
        self.pc = pc;
    }

    pub fn inst_start_pc(&self) -> usize {
        self.inst_start_pc
    }
}