use std::collections::BTreeMap;

use proc_macro2::TokenStream;
use quote::quote;

use crate::config::Config;
use crate::ir::{
    ContentType, OperationIr, OperationResponse, ParamLocation, PrimitiveType, ResponseStatusCode,
    TypeRef,
};

use heck::ToSnakeCase;

use super::{
    SpecInfo, generate_doc_comment, generated_file_banner, generated_file_preamble,
    to_pascal_ident, to_snake_ident, type_ref_to_tokens,
};

/// Module prefix for schema types referenced from operations.
fn types_prefix() -> TokenStream {
    quote! { super::types }
}

/// Generate token streams for a list of operations.
pub fn generate_operations_tokens(
    operations: &[OperationIr],
    default_tag: &str,
    config: &Config,
) -> Vec<TokenStream> {
    let grouped = group_by_tag(operations, default_tag);

    let mut all_items = vec![
        generated_file_preamble(),
        quote! { use orator_axum::serde; },
    ];

    for (tag, ops) in &grouped {
        for op in ops {
            all_items.push(generate_response_enum(op));
            all_items.extend(generate_params_structs(op, config));
            all_items.extend(generate_multipart_body_struct(op));
        }
        all_items.push(generate_api_trait(tag, ops, config));
    }

    all_items
}

/// Generate Rust source code for a list of operations.
pub fn generate_operations(
    operations: &[OperationIr],
    default_tag: &str,
    config: &Config,
    spec_info: &SpecInfo,
) -> String {
    let items = generate_operations_tokens(operations, default_tag, config);
    let file_tokens = quote! { #(#items)* };
    let syntax_tree: syn::File =
        syn::parse2(file_tokens).expect("generated tokens should be valid syntax");
    format!(
        "{}{}",
        generated_file_banner(spec_info),
        prettyplease::unparse(&syntax_tree)
    )
}

pub fn group_by_tag<'a>(
    operations: &'a [OperationIr],
    default_tag: &str,
) -> BTreeMap<String, Vec<&'a OperationIr>> {
    let mut groups: BTreeMap<String, Vec<&'a OperationIr>> = BTreeMap::new();
    for op in operations {
        let tag = op.tag.as_deref().unwrap_or(default_tag);
        groups.entry(tag.to_string()).or_default().push(op);
    }
    groups
}

pub fn status_code_variant_name(response: &OperationResponse) -> String {
    match &response.status_code {
        ResponseStatusCode::Default => "Default".to_string(),
        ResponseStatusCode::Code(code) => match code {
            200 => "Ok".to_string(),
            201 => "Created".to_string(),
            202 => "Accepted".to_string(),
            204 => "NoContent".to_string(),
            301 => "MovedPermanently".to_string(),
            304 => "NotModified".to_string(),
            400 => "BadRequest".to_string(),
            401 => "Unauthorized".to_string(),
            403 => "Forbidden".to_string(),
            404 => "NotFound".to_string(),
            405 => "MethodNotAllowed".to_string(),
            409 => "Conflict".to_string(),
            410 => "Gone".to_string(),
            422 => "UnprocessableEntity".to_string(),
            429 => "TooManyRequests".to_string(),
            500 => "InternalServerError".to_string(),
            502 => "BadGateway".to_string(),
            503 => "ServiceUnavailable".to_string(),
            other => format!("Status{other}"),
        },
    }
}

/// The canonical ordering for parameter locations in generated signatures.
pub const PARAM_LOCATIONS: &[ParamLocation] = &[
    ParamLocation::Path,
    ParamLocation::Query,
    ParamLocation::Header,
    ParamLocation::Cookie,
];

pub fn location_suffix(location: &ParamLocation) -> &'static str {
    match location {
        ParamLocation::Path => "Path",
        ParamLocation::Query => "Query",
        ParamLocation::Header => "Header",
        ParamLocation::Cookie => "Cookie",
    }
}

pub fn location_arg_name(location: &ParamLocation) -> &'static str {
    match location {
        ParamLocation::Path => "path",
        ParamLocation::Query => "query",
        ParamLocation::Header => "header",
        ParamLocation::Cookie => "cookie",
    }
}

fn generate_response_enum(op: &OperationIr) -> TokenStream {
    let enum_ident = to_pascal_ident(&format!("{}Response", op.operation_id));
    let doc = generate_doc_comment(&op.description);

    let variants: Vec<TokenStream> = op
        .responses
        .iter()
        .map(|resp| {
            let variant_name = status_code_variant_name(resp);
            let variant_ident = to_pascal_ident(&variant_name);
            let variant_doc = generate_doc_comment(&resp.description);

            let is_default = matches!(&resp.status_code, ResponseStatusCode::Default);

            if let Some(body) = &resp.body {
                let body_type = type_ref_to_tokens(&body.type_ref, Some(&types_prefix()));
                if is_default {
                    quote! { #variant_doc #variant_ident(orator_axum::http::StatusCode, #body_type), }
                } else {
                    quote! { #variant_doc #variant_ident(#body_type), }
                }
            } else if is_default {
                quote! { #variant_doc #variant_ident(orator_axum::http::StatusCode), }
            } else {
                quote! { #variant_doc #variant_ident, }
            }
        })
        .collect();

    quote! {
        #doc
        #[derive(Debug)]
        pub enum #enum_ident {
            #(#variants)*
        }
    }
}

fn generate_params_structs(op: &OperationIr, config: &Config) -> Vec<TokenStream> {
    let mut structs = Vec::new();

    for location in PARAM_LOCATIONS {
        if !config.is_location_enabled(location) {
            continue;
        }
        let params: Vec<_> = op
            .parameters
            .iter()
            .filter(|p| &p.location == location)
            .collect();

        if params.is_empty() {
            continue;
        }

        let struct_ident =
            to_pascal_ident(&format!("{}{}", op.operation_id, location_suffix(location)));

        let is_query = *location == ParamLocation::Query;

        let fields: Vec<TokenStream> = params
            .iter()
            .map(|param| {
                let field_ident = to_snake_ident(&param.name);
                let field_type = if param.required {
                    type_ref_to_tokens(&param.type_ref, Some(&types_prefix()))
                } else {
                    let inner = type_ref_to_tokens(&param.type_ref, Some(&types_prefix()));
                    quote! { Option<#inner> }
                };
                let serde_rename = if is_query && param.name != param.name.to_snake_case() {
                    let original = &param.name;
                    quote! { #[serde(rename = #original)] }
                } else {
                    quote! {}
                };
                let field_doc = generate_doc_comment(&param.description);
                quote! { #field_doc #serde_rename pub #field_ident: #field_type, }
            })
            .collect();

        let derives = if is_query {
            quote! {
                #[derive(Debug, Clone, serde::Deserialize)]
                #[serde(crate = "orator_axum::serde")]
            }
        } else {
            quote! { #[derive(Debug, Clone)] }
        };

        structs.push(quote! {
            #derives
            pub struct #struct_ident {
                #(#fields)*
            }
        });
    }

    structs
}

fn generate_multipart_body_struct(op: &OperationIr) -> Option<TokenStream> {
    let body = op.request_body.as_ref()?;
    if body.content_type != ContentType::MultipartFormData {
        return None;
    }
    let fields = body.fields.as_ref()?;
    if fields.is_empty() {
        return None;
    }

    let struct_ident = to_pascal_ident(&format!("{}Body", op.operation_id));

    let field_tokens: Vec<TokenStream> = fields
        .iter()
        .map(|field| {
            let field_ident = to_snake_ident(&field.name);
            let field_type = if field.required {
                multipart_field_type(&field.type_ref)
            } else {
                let inner = multipart_field_type(&field.type_ref);
                quote! { Option<#inner> }
            };
            let field_doc = generate_doc_comment(&field.description);
            quote! { #field_doc pub #field_ident: #field_type, }
        })
        .collect();

    Some(quote! {
        #[derive(Debug)]
        pub struct #struct_ident {
            #(#field_tokens)*
        }
    })
}

/// Map a TypeRef to the appropriate multipart field type.
///
/// Binary fields use `orator_axum::bytes::Bytes` directly.
fn multipart_field_type(type_ref: &TypeRef) -> TokenStream {
    match type_ref {
        TypeRef::Primitive(PrimitiveType::Bytes) => quote! { orator_axum::bytes::Bytes },
        _ => type_ref_to_tokens(type_ref, Some(&types_prefix())),
    }
}

/// Return the body struct name for a multipart operation, if it has fields.
pub fn multipart_body_struct_name(op: &OperationIr) -> Option<String> {
    let body = op.request_body.as_ref()?;
    if body.content_type != ContentType::MultipartFormData {
        return None;
    }
    let fields = body.fields.as_ref()?;
    if fields.is_empty() {
        return None;
    }
    Some(format!("{}Body", op.operation_id))
}

fn generate_api_trait(tag: &str, operations: &[&OperationIr], config: &Config) -> TokenStream {
    let trait_ident = to_pascal_ident(&format!("{tag}Api"));

    let methods: Vec<TokenStream> = operations
        .iter()
        .map(|op| {
            let method_ident = to_snake_ident(&op.operation_id);
            let response_ident = to_pascal_ident(&format!("{}Response", op.operation_id));
            let doc = generate_doc_comment(&op.description);

            let mut extra_args = Vec::new();

            for location in PARAM_LOCATIONS {
                if !config.is_location_enabled(location) {
                    continue;
                }
                let has_params = op.parameters.iter().any(|p| &p.location == location);
                if has_params {
                    let params_ident = to_pascal_ident(&format!(
                        "{}{}",
                        op.operation_id,
                        location_suffix(location)
                    ));
                    let arg_name = to_snake_ident(location_arg_name(location));
                    extra_args.push(quote! { #arg_name: #params_ident });
                }
            }

            if let Some(body) = &op.request_body {
                let raw_type = if let Some(struct_name) = multipart_body_struct_name(op) {
                    let ident = to_pascal_ident(&struct_name);
                    quote! { #ident }
                } else if body.content_type == ContentType::MultipartFormData {
                    // Multipart without inline fields — fall back to raw extractor.
                    // Not wrapped in Option because axum's Multipart doesn't support it.
                    quote! { orator_axum::axum::extract::Multipart }
                } else {
                    type_ref_to_tokens(&body.type_ref, Some(&types_prefix()))
                };
                let body_type =
                    if body.required || body.content_type == ContentType::MultipartFormData {
                        raw_type
                    } else {
                        quote! { Option<#raw_type> }
                    };
                extra_args.push(quote! { body: #body_type });
            }

            quote! {
                #doc
                fn #method_ident(
                    &self,
                    ctx: Self::RequestContext,
                    #(#extra_args),*
                ) -> impl std::future::Future<Output = Result<#response_ident, Self::Error>> + Send;
            }
        })
        .collect();

    quote! {
        pub trait #trait_ident: Send + Sync + 'static {
            type Error: Send;
            type RequestContext: Send;

            #(#methods)*
        }
    }
}
