use std::{collections::HashMap, path::Path, rc::Rc};

use proc_macro2::{Span, TokenStream};
use quote::{quote, TokenStreamExt};
use syn::{punctuated::Punctuated, spanned::Spanned, token, Token};

pub mod ty;

use crate::{
    abi::ty::{AbiPrimitive, ComplexType, SimpleType},
    error::Error,
    syntax::{
        parse_abi, Attributes, BitfieldDefaultValue, BitfieldRange, BitfieldType,
        BitfieldTypeField, DocumentTypeDefinition, EnumType, EnumTypeVariant, NewType, ParseError,
        StructType, TypeRepr,
    },
    TargetEnv,
};

use self::ty::{complex::ExternKind, Type};

mod syscall;

pub use syscall::Syscall;

pub trait GenerateTypeDefinition {
    fn generate_type_definition(&self, abi_type: &ComplexType, env: &TargetEnv) -> TokenStream;
}

pub struct TypeEnv {
    pub locals: HashMap<String, (Rc<ComplexType>, DocumentTypeDefinition)>,
    pub externs: HashMap<String, (syn::Type, ExternKind)>,
}

impl TypeEnv {
    pub fn new() -> Self {
        Self {
            locals: HashMap::new(),
            externs: HashMap::new(),
        }
    }

    pub fn lookup_local(&self, name: &str) -> Option<&(Rc<ComplexType>, DocumentTypeDefinition)> {
        self.locals.get(name)
    }

    pub fn lookup_extern(&self, name: &str) -> Option<&(syn::Type, ExternKind)> {
        self.externs.get(name)
    }

    pub fn insert_local(
        &mut self,
        name: String,
        abi_type: Rc<ComplexType>,
        def: DocumentTypeDefinition,
    ) {
        self.locals.insert(name, (abi_type, def));
    }

    pub fn insert_extern(&mut self, name: String, def: syn::Type, kind: ExternKind) {
        self.externs.insert(name, (def, kind));
    }
}

pub struct Abi {
    pub types: TypeEnv,
    // Contains syscalls and their corresponding SyscallFunction enum variants
    pub syscalls: Vec<(syn::Ident, Syscall)>,
}

pub struct AbiBuilder {
    abi: Abi,
    target_env: TargetEnv,
}

pub fn syscall_enum_variant_identifier(name: &syn::Ident) -> syn::Ident {
    let name = name.to_string();
    let string = name.split('_').fold(String::new(), |mut acc, word| {
        word.chars().enumerate().for_each(|(i, c)| {
            if i == 0 {
                acc.extend(c.to_uppercase());
            } else {
                acc.push(c);
            }
        });
        acc
    });
    syn::Ident::new(string.as_str(), name.span())
}

impl AbiBuilder {
    pub fn from_file<P: AsRef<Path>>(path: P, target_env: TargetEnv) -> Result<Self, ParseError> {
        let abi = Abi::read_file(path)?;
        Ok(Self { abi, target_env })
    }

    pub fn from_string<S: AsRef<str>>(data: S, target_env: TargetEnv) -> Result<Self, ParseError> {
        let abi = Abi::read_string(data)?;
        Ok(Self { abi, target_env })
    }

    pub fn emit_syscall_dispatcher(
        &self,
        dispatcher_fn: &str,
        impl_mod: &str,
    ) -> Result<syn::File, Error> {
        let dispatcher_fn = syn::Ident::new(dispatcher_fn, Span::call_site());
        let impl_mod = syn::Ident::new(impl_mod, Span::call_site());

        let raw_fn = syn::Ident::new("__fn", Span::call_site());
        let raw_args = syn::Ident::new("__args", Span::call_site());

        let mut branches = Punctuated::<_, Token![,]>::new();

        for (enum_variant, syscall) in self.abi.syscalls.iter() {
            let named_args = (0..syscall.args.len())
                .map(|i| syn::Ident::new(&format!("__a{}", i), Span::call_site()))
                .collect::<Vec<_>>();
            let call_args = named_args
                .iter()
                .cloned()
                .collect::<Punctuated<_, Token![,]>>();
            let syscall_name = &syscall.name;

            let handler_fn = quote! { #impl_mod::#syscall_name };

            let body = if syscall
                .return_type
                .as_ref()
                .map(|p| {
                    matches!(
                        p.as_ref(),
                        ComplexType::Simple(SimpleType::Primitive(AbiPrimitive::Unit))
                    )
                })
                .unwrap_or(true)
            {
                quote! {
                    #handler_fn(#call_args);
                    Ok(0)
                }
            } else {
                quote! {
                    Ok((#handler_fn(#call_args)).into_syscall_register())
                }
            };

            let match_arm = syscall.emit_handler_branch(
                &named_args,
                &raw_args,
                body,
                enum_variant,
                &self.target_env,
            );

            branches.push(match_arm);
        }

        let dispatcher_fn = quote! {
            #[allow(clippy::unnecessary_cast, clippy::identity_op)]
            pub(crate) fn #dispatcher_fn(#raw_fn: SyscallFunction, #raw_args: &[usize]) -> Result<usize, Error> {
                match #raw_fn {
                    #branches
                }
            }
        };

        let file = syn::parse_file(&dispatcher_fn.to_string())?;
        Ok(file)
    }

    pub fn emit_file(&self, emit_types: bool, emit_calls: bool) -> Result<syn::File, ParseError> {
        let mut stream = TokenStream::new();

        if emit_types {
            stream.extend(self.emit_syscall_enum());
            stream.extend(self.emit_type_definitions());
        }

        if emit_calls {
            // Preamble
            stream.extend(quote! {
                #[allow(unused_imports)]
                use abi_lib::{SyscallRegister, SyscallFatRegister};
            });
            stream.extend(self.emit_syscalls());
        }

        let file = syn::parse_file(&stream.to_string())?;
        Ok(file)
    }

    fn emit_syscall_enum(&self) -> TokenStream {
        let syscall_discriminants = self
            .abi
            .syscalls
            .iter()
            .enumerate()
            .map(|(index, (variant, _))| {
                let index = index + 1;
                quote! { #variant = #index }
            })
            .collect::<Punctuated<_, token::Comma>>();

        let syscall_try_from_branches = self
            .abi
            .syscalls
            .iter()
            .enumerate()
            .map(|(index, (variant, _))| {
                let index = index + 1;
                quote! { #index => Ok(Self::#variant) }
            })
            .collect::<Punctuated<_, token::Comma>>();

        quote! {
            #[derive(Clone, Copy, PartialEq, Eq, Debug, PartialOrd, Ord)]
            #[repr(usize)]
            pub enum SyscallFunction {
                #syscall_discriminants
            }

            impl From<SyscallFunction> for usize {
                fn from(value: SyscallFunction) -> usize {
                    value as usize
                }
            }

            impl TryFrom<usize> for SyscallFunction {
                type Error = Error;

                fn try_from(value: usize) -> Result<Self, Error> {
                    match value {
                        #syscall_try_from_branches,
                        _ => Err(Error::UndefinedSyscall)
                    }
                }
            }
        }
    }

    fn emit_type_definitions(&self) -> TokenStream {
        let mut stream = TokenStream::new();

        for (_, (ty, def)) in &self.abi.types.locals {
            stream.extend(def.generate_type_definition(ty, &self.target_env));
        }

        stream
    }

    fn emit_syscalls(&self) -> TokenStream {
        let mut stream = TokenStream::new();

        for (variant, syscall) in &self.abi.syscalls {
            stream.append_all(syscall.emit_stub(variant, &self.target_env));
        }

        stream
    }
}

impl Abi {
    pub fn read_file<P: AsRef<Path>>(path: P) -> Result<Self, ParseError> {
        let data = std::fs::read_to_string(path)?;
        parse_abi(&data)
    }

    pub fn read_string<S: AsRef<str>>(data: S) -> Result<Self, ParseError> {
        parse_abi(data.as_ref())
    }
}

impl GenerateTypeDefinition for StructType {
    fn generate_type_definition(&self, _abi_type: &ComplexType, _env: &TargetEnv) -> TokenStream {
        let Self { inner } = self;
        quote!(#inner)
    }
}

impl GenerateTypeDefinition for EnumType {
    fn generate_type_definition(&self, abi_type: &ComplexType, env: &TargetEnv) -> TokenStream {
        let Self {
            attributes,
            name,
            repr,
            variants,
        } = self;
        let TypeRepr(repr) = &repr;

        // TODO auto-assign numbers to enum variants
        let try_from_branches = variants
            .iter()
            .map(|variant| {
                let EnumTypeVariant {
                    discriminant,
                    value,
                    ..
                } = variant;
                quote! { #value => Ok(Self::#discriminant) }
            })
            .collect::<Punctuated<_, token::Comma>>();

        let direct = quote! {
            #attributes
            #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
            #[repr(#repr)]
            pub enum #name {
                #variants
            }

            impl #name {
                pub const fn into_raw(self) -> #repr {
                    self as #repr
                }

                #[allow(clippy::missing_safety_doc)]
                pub const unsafe fn from_raw(value: #repr) -> Self {
                    core::mem::transmute(value)
                }
            }

            impl From<#name> for #repr {
                #[inline]
                #[allow(clippy::unnecessary_cast)]
                fn from(value: #name) -> #repr {
                    value as #repr
                }
            }

            impl TryFrom<#repr> for #name {
                type Error = Error;

                fn try_from(value: #repr) -> Result<#name, Error> {
                    match value {
                        #try_from_branches,
                        _ => Err(Error::InvalidArgument)
                    }
                }
            }

            impl abi_lib::SyscallRegister for #name {
                fn into_syscall_register(self) -> usize {
                    self as #repr as usize
                }

                fn from_syscall_register(value: usize) -> Self {
                    unsafe { core::mem::transmute(value as #repr) }
                }
            }
        };

        let wrapped = if abi_type.width(env) < env.fat_pointer_width {
            quote! {
                #direct

                unsafe impl abi_lib::ThinWrapped for #name {}
            }
        } else {
            direct
        };

        wrapped
    }
}

fn generate_transparent_struct(
    attributes: &Attributes,
    name: &syn::Ident,
    env: &TargetEnv,
    abi_type: AbiPrimitive,
) -> TokenStream {
    let is_thin = abi_type.width(env) <= env.thin_pointer_width;
    let real_repr = abi_type.as_rust_type();

    let attributes = attributes.filtered(|attr| !attr.path().is_ident("default"));

    let direct = quote! {
        #attributes
        #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
        pub struct #name(pub(crate) #real_repr);

        impl #name {
            pub const fn into_raw(self) -> #real_repr {
                self.0
            }

            pub const fn bits(self) -> #real_repr {
                self.0
            }

            #[allow(clippy::missing_safety_doc)]
            pub const unsafe fn from_raw(value: #real_repr) -> Self {
                Self(value)
            }
        }

        impl From<#name> for #real_repr {
            #[inline]
            fn from(value: #name) -> #real_repr {
                value.0
            }
        }

        impl abi_lib::SyscallRegister for #name {
            fn into_syscall_register(self) -> usize {
                self.0 as _
            }

            fn from_syscall_register(value: usize) -> Self {
                Self(value as _)
            }
        }
    };

    let wrapped = if is_thin {
        quote! {
            #direct

            unsafe impl abi_lib::ThinWrapped for #name {}
        }
    } else {
        direct
    };

    wrapped
}

impl GenerateTypeDefinition for NewType {
    fn generate_type_definition(&self, abi_type: &ComplexType, env: &TargetEnv) -> TokenStream {
        let Self {
            attributes, name, ..
        } = self;
        let abi_type = abi_type.as_primitive().unwrap();

        generate_transparent_struct(attributes, name, env, abi_type)
    }
}

impl GenerateTypeDefinition for BitfieldType {
    fn generate_type_definition(&self, abi_type: &ComplexType, env: &TargetEnv) -> TokenStream {
        let Self {
            attributes,
            name,
            fields,
            ..
        } = self;
        let abi_type = abi_type.as_primitive().unwrap();

        // Find default attribute, if any
        let default_attr: Option<Result<BitfieldDefaultValue, _>> = attributes
            .iter()
            .find_map(|attr| attr.path().is_ident("default").then_some(attr.parse_args()));

        let base = generate_transparent_struct(attributes, name, env, abi_type);

        let mut field_impls = TokenStream::new();
        for BitfieldTypeField {
            attributes,
            name,
            range,
        } in fields
        {
            match range {
                BitfieldRange::Single(offset) => {
                    field_impls.extend(quote! {
                        #attributes
                        pub const #name: Self = Self(1 << #offset);
                    });
                }
                BitfieldRange::Range(_start, _end) => todo!(),
            }
        }

        let const_default_fn = if let Some(default_attr) = default_attr {
            // TODO handle error
            let default_attr = default_attr.unwrap();

            match default_attr {
                BitfieldDefaultValue::Empty(_) => {
                    quote! {
                        pub const fn const_default() -> Self {
                            Self::empty()
                        }
                    }
                }
                BitfieldDefaultValue::Fields(items) => {
                    let value = items
                        .into_iter()
                        .map(|ident| quote!(Self::#ident.bits()))
                        .collect::<Punctuated<_, Token![|]>>();

                    quote! {
                        pub const fn const_default() -> Self {
                            Self(#value)
                        }
                    }
                }
            }
        } else {
            TokenStream::new()
        };

        let base = quote! {
            #base

            impl #name {
                #field_impls

                pub const fn empty() -> Self {
                    Self(0)
                }

                pub const fn contains(&self, value: Self) -> bool {
                    self.0 & value.0 == value.0
                }

                pub const fn contains_any(&self, value: Self) -> bool {
                    self.0 & value.0 != 0
                }

                #const_default_fn
            }

            impl core::ops::BitOrAssign for #name {
                fn bitor_assign(&mut self, other: Self) {
                    self.0 |= other.0;
                }
            }

            impl core::ops::BitAndAssign for #name {
                fn bitand_assign(&mut self, other: Self) {
                    self.0 &= other.0;
                }
            }

            impl core::ops::BitAnd for #name {
                type Output = #name;

                fn bitand(self, other: Self) -> Self {
                    Self(self.0 & other.0)
                }
            }

            impl core::ops::BitOr for #name {
                type Output = #name;

                fn bitor(self, other: Self) -> Self {
                    Self(self.0 | other.0)
                }
            }

            impl core::ops::Not for #name {
                type Output = #name;

                fn not(self) -> Self {
                    Self(!self.0)
                }
            }
        };

        base
    }
}

impl GenerateTypeDefinition for DocumentTypeDefinition {
    fn generate_type_definition(&self, abi_type: &ComplexType, env: &TargetEnv) -> TokenStream {
        match self {
            Self::Enum(inner) => inner.generate_type_definition(abi_type, env),
            Self::Bitfield(inner) => inner.generate_type_definition(abi_type, env),
            Self::NewType(inner) => inner.generate_type_definition(abi_type, env),
            Self::Struct(inner) => inner.generate_type_definition(abi_type, env),
            // Extern types require no definition
            Self::ExternType(_) => TokenStream::new(),
        }
    }
}
