use std::fmt::{Display, Formatter};
use std::str::FromStr;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum FloatEncoding {
    Ieee,
    Brain,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DTypeKind {
    Bool,
    UInt,
    SInt,
    Float(FloatEncoding),
    Complex(FloatEncoding),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct DType {
    pub kind: DTypeKind,
    pub bits: u16,
}

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("unknown dtype `{name}`")]
pub struct ParseDTypeError {
    pub name: String,
}

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum PromoteError {
    #[error("cannot promote {lhs} and {rhs}")]
    Incompatible { lhs: DType, rhs: DType },
}

impl DType {
    pub const BOOL: Self = Self {
        kind: DTypeKind::Bool,
        bits: 1,
    };
    pub const U8: Self = Self {
        kind: DTypeKind::UInt,
        bits: 8,
    };
    pub const U16: Self = Self {
        kind: DTypeKind::UInt,
        bits: 16,
    };
    pub const U32: Self = Self {
        kind: DTypeKind::UInt,
        bits: 32,
    };
    pub const U64: Self = Self {
        kind: DTypeKind::UInt,
        bits: 64,
    };
    pub const I8: Self = Self {
        kind: DTypeKind::SInt,
        bits: 8,
    };
    pub const I16: Self = Self {
        kind: DTypeKind::SInt,
        bits: 16,
    };
    pub const I32: Self = Self {
        kind: DTypeKind::SInt,
        bits: 32,
    };
    pub const I64: Self = Self {
        kind: DTypeKind::SInt,
        bits: 64,
    };
    pub const F16: Self = Self {
        kind: DTypeKind::Float(FloatEncoding::Ieee),
        bits: 16,
    };
    pub const BF16: Self = Self {
        kind: DTypeKind::Float(FloatEncoding::Brain),
        bits: 16,
    };
    pub const F32: Self = Self {
        kind: DTypeKind::Float(FloatEncoding::Ieee),
        bits: 32,
    };
    pub const F64: Self = Self {
        kind: DTypeKind::Float(FloatEncoding::Ieee),
        bits: 64,
    };
    pub const C64: Self = Self {
        kind: DTypeKind::Complex(FloatEncoding::Ieee),
        bits: 64,
    };
    pub const C128: Self = Self {
        kind: DTypeKind::Complex(FloatEncoding::Ieee),
        bits: 128,
    };

    pub const ALL_DTYPES: [DType; 15] = [
        DType::BOOL,
        DType::U8,
        DType::U16,
        DType::U32,
        DType::U64,
        DType::I8,
        DType::I16,
        DType::I32,
        DType::I64,
        DType::F16,
        DType::BF16,
        DType::F32,
        DType::F64,
        DType::C64,
        DType::C128,
    ];

    pub fn to_str(self) -> &'static str {
        match (self.kind, self.bits) {
            (DTypeKind::Bool, _) => "bool",
            (DTypeKind::UInt, 8) => "u8",
            (DTypeKind::UInt, 16) => "u16",
            (DTypeKind::UInt, 32) => "u32",
            (DTypeKind::UInt, 64) => "u64",
            (DTypeKind::SInt, 8) => "i8",
            (DTypeKind::SInt, 16) => "i16",
            (DTypeKind::SInt, 32) => "i32",
            (DTypeKind::SInt, 64) => "i64",
            (DTypeKind::Float(FloatEncoding::Ieee), 16) => "f16",
            (DTypeKind::Float(FloatEncoding::Brain), 16) => "bf16",
            (DTypeKind::Float(FloatEncoding::Ieee), 32) => "f32",
            (DTypeKind::Float(FloatEncoding::Ieee), 64) => "f64",
            (DTypeKind::Complex(FloatEncoding::Ieee), 64) => "c64",
            (DTypeKind::Complex(FloatEncoding::Ieee), 128) => "c128",
            _ => "<invalid-dtype>",
        }
    }
}

pub fn all_dtype_names() -> Vec<&'static str> {
    DType::ALL_DTYPES.iter().map(|dt| dt.to_str()).collect()
}

impl Display for DType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.to_str())
    }
}

impl FromStr for DType {
    type Err = ParseDTypeError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "bool" => Ok(DType::BOOL),
            "u8" => Ok(DType::U8),
            "u16" => Ok(DType::U16),
            "u32" => Ok(DType::U32),
            "u64" => Ok(DType::U64),
            "i8" => Ok(DType::I8),
            "i16" => Ok(DType::I16),
            "i32" => Ok(DType::I32),
            "i64" => Ok(DType::I64),
            "f16" => Ok(DType::F16),
            "bf16" => Ok(DType::BF16),
            "f32" => Ok(DType::F32),
            "f64" => Ok(DType::F64),
            "c64" => Ok(DType::C64),
            "c128" => Ok(DType::C128),
            _ => Err(ParseDTypeError {
                name: s.to_string(),
            }),
        }
    }
}

fn signed_from_bits(bits: u16) -> Result<DType, PromoteError> {
    match bits {
        8 => Ok(DType::I8),
        16 => Ok(DType::I16),
        32 => Ok(DType::I32),
        64 => Ok(DType::I64),
        _ => Err(PromoteError::Incompatible {
            lhs: DType::I64,
            rhs: DType::I64,
        }),
    }
}

fn uint_from_bits(bits: u16) -> Result<DType, PromoteError> {
    match bits {
        8 => Ok(DType::U8),
        16 => Ok(DType::U16),
        32 => Ok(DType::U32),
        64 => Ok(DType::U64),
        _ => Err(PromoteError::Incompatible {
            lhs: DType::U64,
            rhs: DType::U64,
        }),
    }
}

fn ieee_float_from_effective_bits(bits: u16) -> DType {
    if bits <= 16 {
        DType::F16
    } else if bits <= 32 {
        DType::F32
    } else {
        DType::F64
    }
}

fn float_effective_bits(dtype: DType) -> Option<u16> {
    match dtype.kind {
        DTypeKind::Float(FloatEncoding::Brain) => Some(16), // bf16
        DTypeKind::Float(_) => Some(dtype.bits),
        _ => None,
    }
}

fn float_merge(lhs: DType, rhs: DType) -> Result<DType, PromoteError> {
    let le = float_effective_bits(lhs).ok_or(PromoteError::Incompatible { lhs, rhs })?;
    let re = float_effective_bits(rhs).ok_or(PromoteError::Incompatible { lhs, rhs })?;
    let mut eff = le.max(re);

    let same_encoding = lhs.kind == rhs.kind;
    if !same_encoding && le == re {
        // Same effective width but different encodings (e.g. bf16 vs f16) => upshift.
        eff = if eff <= 16 { 32 } else { 64 };
    }

    if same_encoding {
        match lhs.kind {
            DTypeKind::Float(FloatEncoding::Brain) => Ok(DType::BF16),
            DTypeKind::Float(_) => Ok(ieee_float_from_effective_bits(eff)),
            _ => Err(PromoteError::Incompatible { lhs, rhs }),
        }
    } else {
        // Mixed encodings always prefer IEEE.
        Ok(ieee_float_from_effective_bits(eff))
    }
}

fn int_to_ieee_float(dtype: DType) -> Result<DType, PromoteError> {
    match dtype.kind {
        DTypeKind::UInt | DTypeKind::SInt => {
            let bits = dtype.bits;
            let eff = (bits.saturating_mul(2)).clamp(16, 64);
            Ok(ieee_float_from_effective_bits(eff))
        }
        _ => Err(PromoteError::Incompatible {
            lhs: dtype,
            rhs: dtype,
        }),
    }
}

fn complex_ieee(bits: u16) -> DType {
    DType {
        kind: DTypeKind::Complex(FloatEncoding::Ieee),
        bits,
    }
}

fn complex_from_bits(bits: u16) -> DType {
    if bits <= 64 {
        complex_ieee(64)
    } else {
        complex_ieee(128)
    }
}

fn complex_underlying_float(dtype: DType) -> Result<DType, PromoteError> {
    debug_assert!(
        matches!(dtype.kind, DTypeKind::Complex(FloatEncoding::Ieee)),
        "complex dtype must use IEEE encoding"
    );
    match (dtype.kind, dtype.bits) {
        (DTypeKind::Complex(FloatEncoding::Ieee), 64) => Ok(DType::F32),
        (DTypeKind::Complex(FloatEncoding::Ieee), 128) => Ok(DType::F64),
        _ => Err(PromoteError::Incompatible {
            lhs: dtype,
            rhs: dtype,
        }),
    }
}

fn complex_of_float(dtype: DType) -> Result<DType, PromoteError> {
    let eff = float_effective_bits(dtype).ok_or(PromoteError::Incompatible {
        lhs: dtype,
        rhs: dtype,
    })?;
    Ok(if eff <= 32 { DType::C64 } else { DType::C128 })
}

pub fn promote(lhs: DType, rhs: DType) -> Result<DType, PromoteError> {
    if lhs == rhs {
        return Ok(lhs);
    }

    if lhs.kind == DTypeKind::Bool {
        return Ok(rhs);
    }
    if rhs.kind == DTypeKind::Bool {
        return Ok(lhs);
    }

    match (lhs.kind, rhs.kind) {
        (DTypeKind::UInt, DTypeKind::UInt) => uint_from_bits(lhs.bits.max(rhs.bits)),
        (DTypeKind::SInt, DTypeKind::SInt) => signed_from_bits(lhs.bits.max(rhs.bits)),

        (DTypeKind::Float(_), DTypeKind::Float(_)) => float_merge(lhs, rhs),
        (DTypeKind::Complex(_), DTypeKind::Complex(_)) => {
            Ok(complex_from_bits(lhs.bits.max(rhs.bits)))
        }

        // UInt + SInt
        (DTypeKind::UInt, DTypeKind::SInt) | (DTypeKind::SInt, DTypeKind::UInt) => {
            let (uint_dt, sint_dt) = if matches!(lhs.kind, DTypeKind::UInt) {
                (lhs, rhs)
            } else {
                (rhs, lhs)
            };
            let out_bits = if uint_dt.bits < sint_dt.bits {
                sint_dt.bits
            } else {
                (uint_dt.bits.saturating_mul(2)).min(64)
            };
            signed_from_bits(out_bits)
        }

        // Int + Float
        (DTypeKind::UInt, DTypeKind::Float(FloatEncoding::Ieee))
        | (DTypeKind::SInt, DTypeKind::Float(FloatEncoding::Ieee))
        | (DTypeKind::Float(FloatEncoding::Ieee), DTypeKind::UInt)
        | (DTypeKind::Float(FloatEncoding::Ieee), DTypeKind::SInt) => {
            let (int_dt, float_dt) = if matches!(lhs.kind, DTypeKind::UInt | DTypeKind::SInt) {
                (lhs, rhs)
            } else {
                (rhs, lhs)
            };
            let as_float = int_to_ieee_float(int_dt)?;
            float_merge(as_float, float_dt)
        }

        // Float + Complex
        (DTypeKind::Float(FloatEncoding::Ieee), DTypeKind::Complex(_))
        | (DTypeKind::Complex(_), DTypeKind::Float(FloatEncoding::Ieee)) => {
            let (float_dt, complex_dt) = if matches!(lhs.kind, DTypeKind::Float(_)) {
                (lhs, rhs)
            } else {
                (rhs, lhs)
            };
            let cf = complex_underlying_float(complex_dt)?;
            let merged = float_merge(float_dt, cf)?;
            let out = complex_of_float(merged)?;
            Ok(complex_from_bits(out.bits.max(complex_dt.bits)))
        }

        // Int + Complex
        (DTypeKind::UInt, DTypeKind::Complex(_))
        | (DTypeKind::SInt, DTypeKind::Complex(_))
        | (DTypeKind::Complex(_), DTypeKind::UInt)
        | (DTypeKind::Complex(_), DTypeKind::SInt) => {
            let (int_dt, complex_dt) = if matches!(lhs.kind, DTypeKind::UInt | DTypeKind::SInt) {
                (lhs, rhs)
            } else {
                (rhs, lhs)
            };
            let as_float = int_to_ieee_float(int_dt)?;
            let cf = complex_underlying_float(complex_dt)?;
            let merged = float_merge(as_float, cf)?;
            let out = complex_of_float(merged)?;
            Ok(complex_from_bits(out.bits.max(complex_dt.bits)))
        }
        _ => Err(PromoteError::Incompatible { lhs, rhs }),
    }
}

#[cfg(test)]
mod tests {
    use super::{DType, DTypeKind, FloatEncoding, PromoteError, promote};

    #[test]
    fn test_from_str_and_display() {
        let bf16: DType = "bf16".parse().expect("bf16 should parse");
        assert_eq!(bf16, DType::BF16);
        assert_eq!(bf16.to_string(), "bf16");

        let dtype: DType = "f32".parse().expect("f32 should parse");
        assert_eq!(dtype, DType::F32);
        assert_eq!(dtype.to_string(), "f32");
    }

    #[test]
    fn test_promote_same_type() {
        assert_eq!(promote(DType::I32, DType::I32), Ok(DType::I32));
    }

    #[test]
    fn test_promote_bool_to_other() {
        assert_eq!(promote(DType::BOOL, DType::U16), Ok(DType::U16));
        assert_eq!(promote(DType::BOOL, DType::F64), Ok(DType::F64));
    }

    #[test]
    fn test_promote_kind_ranking() {
        assert_eq!(promote(DType::U8, DType::I8), Ok(DType::I16));
        assert_eq!(promote(DType::U16, DType::I64), Ok(DType::I64));
        assert_eq!(promote(DType::I32, DType::F16), Ok(DType::F64));
    }

    #[test]
    fn test_promote_float_and_complex_and_bf16() {
        assert_eq!(promote(DType::BF16, DType::BF16), Ok(DType::BF16));
        assert_eq!(promote(DType::BF16, DType::F16), Ok(DType::F32));
        assert_eq!(promote(DType::BF16, DType::F32), Ok(DType::F32));
        assert_eq!(promote(DType::F64, DType::C64), Ok(DType::C128));
        assert_eq!(promote(DType::C64, DType::C128), Ok(DType::C128));
    }

    #[test]
    fn test_promote_non_ieee_float_mixed_families_is_incompatible() {
        assert!(matches!(
            promote(DType::BF16, DType::I32),
            Err(PromoteError::Incompatible { .. })
        ));
        assert!(matches!(
            promote(DType::I32, DType::BF16),
            Err(PromoteError::Incompatible { .. })
        ));
        assert!(matches!(
            promote(DType::BF16, DType::C64),
            Err(PromoteError::Incompatible { .. })
        ));
        assert!(matches!(
            promote(DType::C64, DType::BF16),
            Err(PromoteError::Incompatible { .. })
        ));
    }

    #[test]
    fn test_promote_uint_sint_error() {
        // UInt + SInt is always promotable; u64 mixes may be unsafe but still returns i64.
        assert_eq!(promote(DType::U64, DType::I64), Ok(DType::I64));
        assert_eq!(promote(DType::U32, DType::I16), Ok(DType::I64));
    }

    #[test]
    fn test_dtype_layout() {
        assert_eq!(DType::BF16.kind, DTypeKind::Float(FloatEncoding::Brain));
        assert_eq!(DType::C64.kind, DTypeKind::Complex(FloatEncoding::Ieee));
        assert_eq!(DType::C128.kind, DTypeKind::Complex(FloatEncoding::Ieee));
        assert_eq!(DType::C128.bits, 128);
    }
}
