use std::{rc::Rc, str::FromStr};

use syn::{punctuated::Punctuated, Expr, ExprLit, Lit, PathSegment};

use crate::abi::{
    syscall_enum_variant_identifier,
    ty::{complex::ExternKind, AbiPrimitive, ComplexType, SimpleType},
    Abi, Syscall, TypeEnv,
};

mod common;

// Types
mod bitfield_type;
mod enum_type;
mod extern_block;
mod newtype;
mod struct_type;

mod document;

mod error;
mod syscall;

pub use bitfield_type::{BitfieldDefaultValue, BitfieldRange, BitfieldType, BitfieldTypeField};
pub use common::{Attributes, TypeRepr};
pub use document::{Document, DocumentTypeDefinition};
pub use enum_type::{EnumType, EnumTypeVariant};
pub use error::{ParseError, SyntaxError, UnwrapFancy};
pub use extern_block::{ExternType, ExternTypeBlock};
pub use newtype::NewType;
pub use struct_type::StructType;

pub fn parse_abi_document(input: &str) -> Result<Document, ParseError> {
    let document: Document = syn::parse_str(input)?;
    Ok(document)
}

fn as_single_segment_path(path: &syn::TypePath) -> Option<&PathSegment> {
    if path.qself.is_some() || path.path.leading_colon.is_some() || path.path.segments.len() != 1 {
        return None;
    }

    Some(&path.path.segments[0])
}

fn as_single_type_argument(args: &syn::PathArguments) -> Option<&syn::Type> {
    let syn::PathArguments::AngleBracketed(args) = args else {
        return None;
    };
    if args.args.len() != 1 {
        return None;
    }
    let syn::GenericArgument::Type(ty) = &args.args[0] else {
        return None;
    };

    Some(ty)
}

fn as_const_array_len(len: &Expr) -> Option<usize> {
    let Expr::Lit(ExprLit {
        attrs,
        lit: Lit::Int(lit),
    }) = len
    else {
        return None;
    };

    if !attrs.is_empty() {
        todo!()
    }

    lit.base10_parse().ok()
}

fn process_type(t: &syn::Type, known_types: &TypeEnv) -> Result<Rc<ComplexType>, ParseError> {
    match t {
        syn::Type::Path(path) => {
            let segment = as_single_segment_path(path).expect("TODO: more complex paths");

            match segment.ident.to_string().as_str() {
                seg if let Ok(ty) = AbiPrimitive::from_str(seg) => {
                    Ok(Rc::new(ComplexType::primitive(ty)))
                }
                "NonZeroUsize" => Ok(Rc::new(ComplexType::Simple(SimpleType::NonZeroUsize))),
                "Result" if let Some(ty) = as_single_type_argument(&segment.arguments) => {
                    let ty = process_type(ty, known_types)?;
                    Ok(Rc::new(ComplexType::Result(ty)))
                }
                "Option" if let Some(ty) = as_single_type_argument(&segment.arguments) => {
                    let ty = process_type(ty, known_types)?;
                    Ok(Rc::new(ComplexType::Option(ty)))
                }
                "MaybeUninit" if let Some(ty) = as_single_type_argument(&segment.arguments) => {
                    let ty = process_type(ty, known_types)?;
                    Ok(Rc::new(ComplexType::MaybeUninit(ty)))
                }
                name if let Some((ty, _)) = known_types.lookup_local(name) => match ty.as_ref() {
                    ComplexType::Struct { name, kind, .. } => Ok(Rc::new(ComplexType::Struct {
                        name: name.clone(),
                        kind: *kind,
                        arguments: segment.arguments.clone(),
                    })),
                    _ if !segment.arguments.is_empty() => Err(ParseError::syntax1(
                        SyntaxError::UnexpectedAttributes(segment.ident.clone()),
                    )),
                    _ => Ok(ty.clone()),
                },
                name if let Some((ty, kind)) = known_types.lookup_extern(name) => {
                    Ok(Rc::new(ComplexType::Extern {
                        kind: *kind,
                        alias: ty.clone(),
                        arguments: segment.arguments.clone(),
                    }))
                }
                _ => Err(ParseError::syntax1(SyntaxError::UndefinedIdentifier(
                    segment.ident.clone(),
                ))),
            }
        }
        syn::Type::Never(_) => Ok(Rc::new(ComplexType::Simple(SimpleType::Never))),
        syn::Type::Reference(r) => {
            // TODO lifetimes
            if r.lifetime.is_some() {
                todo!()
            }

            let mutable = r.mutability.is_some();

            match r.elem.as_ref() {
                syn::Type::Slice(ty) => {
                    let element = process_type(&ty.elem, known_types)?;
                    Ok(Rc::new(ComplexType::Simple(SimpleType::Slice {
                        mutable,
                        element,
                    })))
                }
                syn::Type::Path(path)
                    if !mutable
                        && as_single_segment_path(path)
                            .map(|seg| seg.ident.to_string() == "str" && seg.arguments.is_empty())
                            .unwrap_or(false) =>
                {
                    Ok(Rc::new(ComplexType::Simple(SimpleType::Str)))
                }
                syn::Type::Array(array) if let Some(length) = as_const_array_len(&array.len) => {
                    let element = process_type(&array.elem, known_types)?;
                    Ok(Rc::new(ComplexType::Simple(SimpleType::Array {
                        mutable,
                        element,
                        length,
                    })))
                }
                _ => {
                    let pointee = process_type(&r.elem, known_types)?;
                    Ok(Rc::new(ComplexType::Simple(SimpleType::Reference {
                        mutable,
                        pointee,
                    })))
                }
            }
        }
        syn::Type::Tuple(t) => {
            if t.elems.is_empty() {
                Ok(Rc::new(ComplexType::primitive(AbiPrimitive::Unit)))
            } else {
                let elems = t
                    .elems
                    .iter()
                    .map(|e| process_type(e, known_types))
                    .collect::<Result<_, _>>()?;
                Ok(Rc::new(ComplexType::Tuple(elems)))
            }
        }
        _ => todo!(),
    }
}

fn require_primitive_repr(repr: &TypeRepr) -> Result<AbiPrimitive, ParseError> {
    let TypeRepr(repr) = repr;
    Ok(AbiPrimitive::from_str(repr.to_string().as_str())
        .map_err(|_| ParseError::syntax1(SyntaxError::PrimitiveReprRequired(repr.clone())))?)
}

fn make_single_ident_type(ident: syn::Ident) -> syn::Type {
    syn::Type::Path(syn::TypePath {
        qself: None,
        path: syn::Path {
            leading_colon: None,
            segments: Punctuated::from_iter([syn::PathSegment {
                arguments: syn::PathArguments::None,
                ident,
            }]),
        },
    })
}

pub enum ProcessedType {
    Extern(syn::Type, ExternKind),
    Local(Rc<ComplexType>, DocumentTypeDefinition),
}

fn process_extern_type(
    def: &ExternType,
    _known_types: &TypeEnv,
) -> Result<(syn::Ident, syn::Type, ExternKind), ParseError> {
    // TODO allow mixing arguments from
    // extern { type X = ::A<B, C> };
    // and use sites like
    // syscall my_call(arg: X<D, E>);

    // Process extern attributes
    let mut kind = ExternKind::None;
    for attr in def.attributes.iter() {
        if attr.path().is_ident("thin") {
            kind = ExternKind::Thin;
        } else if attr.path().is_ident("fat") {
            kind = ExternKind::Fat;
        }
    }

    let alias = def
        .alias
        .clone()
        .unwrap_or_else(|| make_single_ident_type(def.name.clone()));
    Ok((def.name.clone(), alias, kind))
}

fn process_bitfield_type(
    def: &BitfieldType,
    _known_types: &TypeEnv,
) -> Result<(syn::Ident, Rc<ComplexType>), ParseError> {
    let repr = require_primitive_repr(&def.repr)?;
    Ok((
        def.name.clone(),
        Rc::new(ComplexType::Simple(SimpleType::Transparent {
            name: def.name.clone(),
            inner: repr,
        })),
    ))
}

fn process_enum_type(
    def: &EnumType,
    _known_types: &TypeEnv,
) -> Result<(syn::Ident, Rc<ComplexType>), ParseError> {
    let repr = require_primitive_repr(&def.repr)?;
    Ok((
        def.name.clone(),
        Rc::new(ComplexType::Simple(SimpleType::Transparent {
            name: def.name.clone(),
            inner: repr,
        })),
    ))
}

fn process_struct_type(
    def: &StructType,
    _known_types: &TypeEnv,
) -> Result<(syn::Ident, Rc<ComplexType>), ParseError> {
    let name = def.inner.ident.clone();
    Ok((
        name.clone(),
        Rc::new(ComplexType::Struct {
            name,
            kind: ExternKind::None,
            arguments: syn::PathArguments::None,
        }),
    ))
}

fn process_newtype(
    def: &NewType,
    _known_types: &TypeEnv,
) -> Result<(syn::Ident, Rc<ComplexType>), ParseError> {
    let repr = require_primitive_repr(&def.repr)?;
    Ok((
        def.name.clone(),
        Rc::new(ComplexType::Simple(SimpleType::Transparent {
            name: def.name.clone(),
            inner: repr,
        })),
    ))
}

fn process_type_definition(
    def: &DocumentTypeDefinition,
    known_types: &TypeEnv,
) -> Result<(syn::Ident, ProcessedType), ParseError> {
    match def {
        DocumentTypeDefinition::ExternType(inner) => process_extern_type(inner, known_types)
            .map(|(ident, def_ty, kind)| (ident, ProcessedType::Extern(def_ty, kind))),
        DocumentTypeDefinition::Enum(inner) => process_enum_type(inner, known_types)
            .map(|(ident, def_ty)| (ident, ProcessedType::Local(def_ty, def.clone()))),
        DocumentTypeDefinition::Bitfield(inner) => process_bitfield_type(inner, known_types)
            .map(|(ident, def_ty)| (ident, ProcessedType::Local(def_ty, def.clone()))),
        DocumentTypeDefinition::NewType(inner) => process_newtype(inner, known_types)
            .map(|(ident, def_ty)| (ident, ProcessedType::Local(def_ty, def.clone()))),
        DocumentTypeDefinition::Struct(inner) => process_struct_type(inner, known_types)
            .map(|(ident, def_ty)| (ident, ProcessedType::Local(def_ty, def.clone()))),
    }
}

pub fn parse_abi(input: &str) -> Result<Abi, ParseError> {
    let doc = parse_abi_document(input)?;
    let mut abi = Abi {
        types: TypeEnv::new(),
        syscalls: Vec::new(),
    };

    ParseError::fold(
        (),
        doc.types.iter().map(|nt| {
            let (name, value) = process_type_definition(nt, &abi.types)?;
            match value {
                ProcessedType::Local(abi_ty, def_ty) => {
                    abi.types.insert_local(name.to_string(), abi_ty, def_ty)
                }
                ProcessedType::Extern(def_ty, kind) => {
                    abi.types.insert_extern(name.to_string(), def_ty, kind)
                }
            }
            Ok(())
        }),
    )?;

    ParseError::fold(
        (),
        doc.syscalls.iter().map(|syscall| {
            let args = syscall
                .arguments
                .iter()
                .map(|arg| {
                    let ty = process_type(&arg.ty, &abi.types)?;
                    Ok((arg.name.clone(), ty))
                })
                .collect::<Result<Vec<_>, ParseError>>()?;
            let return_type = syscall
                .return_type
                .as_ref()
                .map(|ty| process_type(ty, &abi.types))
                .transpose()?;

            let variant = syscall_enum_variant_identifier(&syscall.name);
            let def = Syscall {
                name: syscall.name.clone(),
                args,
                return_type,
            };

            abi.syscalls.push((variant, def));

            Ok(())
        }),
    )?;

    Ok(abi)
}
