// Copyright © 2026 Collabora, Ltd.
// SPDX-License-Identifier: MIT

use crate::encode_v9::*;
use crate::ir::*;
use kraid_bindings::*;

pub trait Model {
    fn arch(&self) -> u8;

    fn encode_shader(&self, s: &Shader<'_>) -> Vec<u32>;

    fn op_is_supported(&self, op: &Op) -> bool;

    fn op_is_message(&self, op: &Op) -> bool;

    fn op_src_is_staging_reg(&self, op: &Op, src: &Src) -> bool;

    fn op_src_supports_imm32(&self, op: &Op, src: &Src) -> bool;

    fn op_src_supports_swizzle(
        &self,
        op: &Op,
        src: &Src,
        swizzle: Swizzle,
    ) -> bool;

    fn op_src_supports_mod(&self, op: &Op, src: &Src, src_mod: SrcMod) -> bool;

    fn op_dst_is_staging_reg(&self, op: &Op) -> bool;

    fn op_dst_supported_lanes(&self, op: &Op) -> DstLanesSet;

    fn op_dst_supports_lanes(&self, op: &Op, lanes: DstLanes) -> bool {
        self.op_dst_supported_lanes(op).contains(lanes)
    }

    fn small_constants(&self) -> &[SmallConstant];

    fn preload_reg(&self, reg: PreloadReg) -> u8;
}

struct ValhallModel {
    arch: u8,
    sc_table: Vec<SmallConstant>,
}

impl ValhallModel {
    fn new(arch: u8) -> ValhallModel {
        use crate::isa::{SmallConstantTable, v9};
        let sc_table = v9::SmallConstantT::collect(arch);
        ValhallModel { arch, sc_table }
    }
}

impl Model for ValhallModel {
    fn arch(&self) -> u8 {
        self.arch
    }

    fn encode_shader(&self, s: &Shader<'_>) -> Vec<u32> {
        encode_v9(s, self.arch)
    }

    fn op_is_supported(&self, op: &Op) -> bool {
        op.as_virtual().is_some() || v9_op_is_supported(op, self.arch)
    }

    fn op_is_message(&self, op: &Op) -> bool {
        if let Some(vop) = op.as_virtual() {
            vop.is_message()
        } else {
            v9_op_is_message(op, self.arch)
        }
    }

    fn op_src_is_staging_reg(&self, op: &Op, src: &Src) -> bool {
        if let Some(vop) = op.as_virtual() {
            vop.src_is_staging_reg(src)
        } else {
            v9_op_src_is_staging_reg(op, src, self.arch)
        }
    }

    fn op_src_supports_imm32(&self, op: &Op, src: &Src) -> bool {
        if let Some(vop) = op.as_virtual() {
            vop.src_supports_imm32(src)
        } else {
            v9_op_src_supports_imm32(op, src, self.arch)
        }
    }

    fn op_src_supports_swizzle(
        &self,
        op: &Op,
        src: &Src,
        swizzle: Swizzle,
    ) -> bool {
        if let Some(vop) = op.as_virtual() {
            vop.src_supports_swizzle(src, swizzle)
        } else {
            v9_op_src_supports_swizzle(op, src, self.arch, swizzle)
        }
    }

    fn op_src_supports_mod(&self, op: &Op, src: &Src, src_mod: SrcMod) -> bool {
        if let Some(vop) = op.as_virtual() {
            vop.src_supports_mod(src, src_mod)
        } else {
            v9_op_src_supports_mod(op, src, self.arch, src_mod)
        }
    }

    fn op_dst_is_staging_reg(&self, op: &Op) -> bool {
        if let Some(vop) = op.as_virtual() {
            vop.dst_is_staging_reg()
        } else {
            v9_op_dst_is_staging_reg(op, self.arch)
        }
    }

    fn op_dst_supported_lanes(&self, op: &Op) -> DstLanesSet {
        if let Some(vop) = op.as_virtual() {
            vop.dst_supported_lanes()
        } else {
            v9_op_dst_supported_lanes(op, self.arch)
        }
    }

    fn small_constants(&self) -> &[SmallConstant] {
        &self.sc_table
    }

    fn preload_reg(&self, reg: PreloadReg) -> u8 {
        use PreloadReg::*;

        match reg {
            LocalId01 => 55,
            LocalId2 => 56,
            WorkgroupId0 => 57,
            WorkgroupId1 => 58,
            WorkgroupId2 => 59,
            GlobalId0 => 60,
            GlobalId1 => 61,
            GlobalId2 => 62,
            InternalId => 59,
            VertexId => 60,
            InstanceId => 61,
            DrawId => 62,
            ViewId => 63,
            PrimitiveId => 57,
            PrimitiveFlags => 58,
            PositionXY => 59,
            CumulativeCoverage => 60,
            RasterizerSampleCentroid => 61,
            FrameArgLow => 62,
            FrameArgHigh => 63,
        }
    }
}

pub fn model_for_gpu_id(
    gpu_id: u64,
) -> Result<Box<dyn Model + Sync + Send>, &'static str> {
    // SAFETY: pan_arch() just translates one integer to another
    let arch = u8::try_from(unsafe { pan_arch(gpu_id) }).unwrap();

    if arch >= 15 {
        Err("Kraid does not yet support this GPU")
    } else if arch >= 9 {
        Ok(Box::new(ValhallModel::new(arch)))
    } else {
        Err("Kraid only supports Valhall (v9) and later GPUs")
    }
}
