use crate::{
    attr::{function_id::FunctionID, FunctionIDAttribute},
    signature::ParsedSignature,
};
use proc_macro2::{Span, TokenStream as TokenStream2};
use proc_macro_error::{abort, abort_call_site};
use quote::{format_ident, quote};
use std::collections::HashSet;
use syn::{
    spanned::Spanned,
    visit::{self, Visit},
    Attribute, Error, FnArg, ImplItemFn, ReturnType, Signature, TraitItemFn, Visibility,
};

// ============================================================================
// Special Method Names - Reserved for Framework Functionality
// ============================================================================

/// Fallback function handler for unmatched selectors.
///
/// When a contract receives a call with an unknown function selector,
/// the fallback function (if defined) will be invoked instead of reverting.
///
/// Requirements:
/// - Must have signature: `fn fallback(&self)` with no parameters or return value
/// - Cannot be called directly via selector
/// - Optional: contract will revert on unknown selectors if not defined
pub(crate) const FALLBACK_METHOD: &str = "fallback";

/// Contract initialization function.
///
/// Defines the contract's initialization logic that runs once during deployment.
/// The framework automatically generates a `deploy()` entry point from this method.
///
/// Requirements:
/// - Must be named exactly `constructor`
/// - Can accept any parameters (will be ABI-encoded during deployment)
/// - Cannot have a return value
/// - Cannot be called after deployment (no function selector generated)
/// - Only one constructor per contract is allowed
pub(crate) const CONSTRUCTOR_METHOD: &str = "constructor";

/// Main dispatch function for routing contract calls.
///
/// This function is automatically generated by the `#[router]` macro and serves
/// as the primary entry point for all contract method calls. It reads the function
/// selector from input data and dispatches to the appropriate method.
///
/// The generated `main()` function:
/// - Reads 4-byte function selector from input
/// - Matches selector against all available methods
/// - Decodes parameters and calls the matched method
/// - Encodes and writes the return value
/// - Falls back to `fallback()` or reverts on unknown selector
pub(crate) const MAIN_METHOD: &str = "main";

/// Reserved name for the deployment entry point.
///
/// User-defined `deploy` methods are forbidden because this function
/// is generated automatically by the framework from the `constructor` method.
///
/// The generated `deploy()` function:
/// - Reads initialization parameters from input data
/// - Decodes them according to the constructor's signature
/// - Calls the user-defined `constructor` with decoded parameters
/// - Does not use function selectors (called directly during contract deployment)
pub(crate) const DEPLOY_METHOD: &str = "deploy";

/// Trait defining common behavior for trait and impl methods
pub trait MethodLike: Sized {
    fn sig(&self) -> &Signature;
    fn attrs(&self) -> &Vec<Attribute>;

    /// Extracts function ID attribute if present
    fn function_id_attr(&self) -> syn::Result<Option<(FunctionIDAttribute, Span)>> {
        self.attrs()
            .iter()
            .find(|a| a.path().is_ident("function_id"))
            .map(|attr| {
                let content_span = attr.meta.require_list()?.tokens.span();
                attr.parse_args::<FunctionIDAttribute>()
                    .map(|parsed| (parsed, content_span))
            })
            .transpose()
    }
}

// Implement MethodLike for TraitItemFn
impl MethodLike for TraitItemFn {
    fn sig(&self) -> &Signature {
        &self.sig
    }

    fn attrs(&self) -> &Vec<Attribute> {
        &self.attrs
    }
}

// Implement MethodLike for ImplItemFn
impl MethodLike for ImplItemFn {
    fn sig(&self) -> &Signature {
        &self.sig
    }

    fn attrs(&self) -> &Vec<Attribute> {
        &self.attrs
    }
}

/// Represents a parsed method in a contract
#[derive(Debug, Clone)]
pub struct ParsedMethod<T: MethodLike> {
    /// Function ID calculated from the method signature
    function_id: FunctionID,
    /// Parsed signature of the method
    sig: ParsedSignature,
    /// Inner function implementation
    inner: T,
}

impl<T: MethodLike> ParsedMethod<T> {
    /// Creates a new ParsedMethod with given inner implementation
    pub fn new(inner: T) -> syn::Result<Self> {
        let sig = ParsedSignature::new(inner.sig().clone());
        let function_id = sig.function_abi()?.function_id()?;

        // Handle custom function ID if defined via attribute
        if let Some((attr, attr_span)) = inner.function_id_attr()? {
            let function_id_attr = attr.function_id_bytes()?;

            if attr.is_validation_enabled() && function_id_attr != function_id {
                abort!(
                    attr_span,
                    "Function ID mismatch: Expected 0x{} for '{}', but got 0x{}",
                    hex::encode(function_id),
                    sig.function_abi()?.signature()?,
                    hex::encode(function_id_attr);
                    note = "You're seeing this error because you have validation enabled (validate(true))";
                    help = "To fix this, you can either:";
                    help = "1. Use the expected function ID: #[function_id(\"{}\")]", sig.function_abi()?.signature()?;
                    help = "2. Remove the validate parameter entirely: #[function_id(\"{}\")]", attr.signature().unwrap_or_else(|| "your_signature".to_string());
                    help = "3. Or explicitly disable validation: #[function_id(\"{}\", validate(false))]", attr.signature().unwrap_or_else(|| "your_signature".to_string())
                );
            } else if !attr.is_validation_enabled() {
                // If validation is disabled, use the function ID from the attribute
                return Ok(Self {
                    function_id: function_id_attr,
                    sig,
                    inner,
                });
            }
        }

        Ok(Self {
            function_id,
            sig,
            inner,
        })
    }

    /// Creates a new ParsedMethod for constructor (without function_id calculation)
    pub fn new_constructor(inner: T) -> syn::Result<Self> {
        let sig = ParsedSignature::new(inner.sig().clone());

        // For constructor, use zero function_id as it doesn't need a selector
        Ok(Self {
            function_id: [0, 0, 0, 0],
            sig,
            inner,
        })
    }

    /// Creates a new ParsedMethod from a reference
    pub fn from_ref(inner: &T) -> syn::Result<Self>
    where
        T: Clone,
    {
        Self::new(inner.clone())
    }

    /// Returns the function ID
    pub fn function_id(&self) -> FunctionID {
        self.function_id
    }

    /// Returns a reference to the inner signature
    pub fn sig(&self) -> &Signature {
        self.inner.sig()
    }

    /// Returns the function's parsed signature
    pub fn parsed_signature(&self) -> &ParsedSignature {
        &self.sig
    }

    /// Returns a reference to the inner implementation
    pub fn inner(&self) -> &T {
        &self.inner
    }

    /// Checks if this is a constructor method
    pub fn is_constructor(&self) -> bool {
        self.sig.rust_name() == CONSTRUCTOR_METHOD
    }
    // В method.rs, добавить в impl<T: MethodLike> ParsedMethod<T>

    /// Generates the deploy method body for a constructor.
    /// This includes parameter decoding and constructor call logic.
    pub fn generate_deploy_body(&self) -> TokenStream2 {
        assert!(
            self.is_constructor(),
            "generate_deploy_body should only be called for constructor methods"
        );

        let fn_name = format_ident!("constructor");
        let params = self.parsed_signature().parameters();
        let param_count = params.len();

        let call_struct = format_ident!("ConstructorCall");

        // Generate parameter handling based on parameter count
        let param_handling = match param_count {
            0 => quote! {},
            1 => quote! {
                let param0 = match #call_struct::decode(&&call_data[..]) {
                    Ok(decoded) => decoded.0.0,
                    Err(err) => {
                        panic!("Failed to decode constructor parameters: {:?}", err);
                    }
                };
            },
            _ => {
                let param_names = (0..param_count)
                    .map(|i| format_ident!("param{}", i))
                    .collect::<Vec<_>>();
                let param_indices = (0..param_count).map(syn::Index::from).collect::<Vec<_>>();

                quote! {
                    let (#(#param_names),*) = match #call_struct::decode(&&call_data[..]) {
                        Ok(decoded) => (#(decoded.0.#param_indices),*),
                        Err(err) => {
                            panic!("Failed to decode constructor parameters: {:?}", err);
                        }
                    };
                }
            }
        };

        // Generate function call based on parameter count
        let fn_call = match param_count {
            0 => quote! { self.#fn_name() },
            1 => quote! { self.#fn_name(param0) },
            _ => {
                let param_names = (0..param_count)
                    .map(|i| format_ident!("param{}", i))
                    .collect::<Vec<_>>();
                quote! { self.#fn_name(#(#param_names),*) }
            }
        };

        // Generate complete deploy body
        quote! {
            let input_length = self.sdk.input_size();
            let mut call_data = alloc::vec![0u8; input_length as usize];
            self.sdk.read(&mut call_data, 0);

            #param_handling
            #fn_call;
        }
    }
}

/// Collector for gathering methods from trait or impl blocks
pub struct MethodCollector<T: MethodLike> {
    /// Collected methods
    pub methods: Vec<ParsedMethod<T>>,
    /// Constructor method
    pub constructor: Option<ParsedMethod<T>>,
    /// Source span for error reporting
    pub span: Span,
    /// Collection of errors encountered during parsing
    pub errors: Vec<Error>,
    /// Set of method selectors to detect collisions
    pub selectors: HashSet<FunctionID>,
    /// Whether this is a trait implementation (for router)
    pub is_trait_impl: bool,
}

impl<T: MethodLike> MethodCollector<T> {
    /// Creates a new method collector for trait methods
    pub fn new(span: Span) -> Self {
        Self {
            methods: Vec::new(),
            constructor: None,
            span,
            errors: Vec::new(),
            selectors: HashSet::new(),
            is_trait_impl: false,
        }
    }

    /// Creates a new method collector for impl methods with trait flag
    pub fn new_for_impl(span: Span, is_trait_impl: bool) -> Self {
        Self {
            methods: Vec::new(),
            constructor: None,
            span,
            errors: Vec::new(),
            selectors: HashSet::new(),
            is_trait_impl,
        }
    }

    /// Validates that there are no function selector collisions
    pub fn validate_selectors(&self) -> Result<(), Error> {
        if self.selectors.len() == self.methods.len() {
            return Ok(());
        }

        // Find the first method with a collision
        for method in &self.methods {
            let selector = method.function_id();
            let count = self
                .methods
                .iter()
                .filter(|m| m.function_id() == selector)
                .count();

            if count > 1 {
                return Err(Error::new(
                    method.parsed_signature().span(),
                    format!(
                        "Function selector collision detected for '{}'. Selector: {:02x}{:02x}{:02x}{:02x}",
                        method.parsed_signature().rust_name(),
                        selector[0], selector[1], selector[2], selector[3]
                    ),
                ));
            }
        }

        Ok(())
    }

    /// Checks if any errors were encountered during method collection
    pub fn has_errors(&self) -> bool {
        !self.errors.is_empty()
    }

    /// Adds a method to the collection after checking for selector collision
    pub fn add_method(&mut self, method: ParsedMethod<T>) {
        let selector = method.function_id();
        if !self.selectors.insert(selector) {
            self.add_error(
                method.parsed_signature().span(),
                format!(
                    "Function selector collision detected for '{}'. Selector: {:02x}{:02x}{:02x}{:02x}",
                    method.parsed_signature().rust_name(),
                    selector[0], selector[1], selector[2], selector[3]
                ),
            );
        } else {
            self.methods.push(method);
        }
    }

    /// Adds an error to the collection
    pub fn add_error(&mut self, span: Span, message: String) {
        self.errors.push(Error::new(span, message));
    }

    /// Validates the signature of the fallback method
    ///
    /// A valid fallback method:
    /// - Takes only a single parameter (`self`)
    /// - Has no return type
    fn validate_fallback_signature(&self, sig: &Signature) -> bool {
        // Check if there's exactly 1 input and it's a receiver
        let has_only_self =
            sig.inputs.len() == 1 && matches!(sig.inputs.first(), Some(FnArg::Receiver(_)));

        let has_no_return = matches!(sig.output, ReturnType::Default);

        has_only_self && has_no_return
    }

    /// Handles fallback method validation
    fn handle_fallback_method(&mut self, method_sig: &Signature, span: Span) {
        if !self.validate_fallback_signature(method_sig) {
            self.add_error(
                span,
                format!("{} method must have signature 'fn {}(&self)' with no parameters and no return value",
                        FALLBACK_METHOD, FALLBACK_METHOD),
            );
        }
    }

    /// Handles constructor method
    fn handle_constructor_method(&mut self, method: &T)
    where
        T: Clone,
    {
        // Check if constructor was already defined
        if self.constructor.is_some() {
            self.add_error(
                method.sig().span(),
                "Only one constructor method is allowed".to_string(),
            );
            return;
        }

        // Create ParsedMethod for constructor without function_id calculation
        match ParsedMethod::new_constructor(method.clone()) {
            Ok(parsed_constructor) => {
                self.constructor = Some(parsed_constructor);
            }
            Err(err) => {
                self.add_error(
                    method.sig().span(),
                    format!("Failed to parse constructor: {}", err),
                );
            }
        }
    }

    /// Handles regular method parsing and validation
    fn handle_regular_method(&mut self, method: &T)
    where
        T: Clone,
    {
        match ParsedMethod::from_ref(method) {
            Ok(parsed_method) => {
                self.add_method(parsed_method);
            }
            Err(err) => {
                self.add_error(
                    method.sig().span(),
                    format!("Failed to parse method: {}", err),
                );
            }
        }
    }
}
// Implementation for MethodCollector<TraitItemFn>
impl Visit<'_> for MethodCollector<TraitItemFn> {
    fn visit_trait_item_fn(&mut self, method: &TraitItemFn) {
        match method.sig.ident.to_string().as_str() {
            // Reserved methods that cannot be defined by user
            MAIN_METHOD => {
                self.add_error(
                    method.sig.span(),
                    "Method 'main' is reserved and generated automatically by the router. It serves as the dispatch entry point for all contract calls".to_string(),
                );
            }
            DEPLOY_METHOD => {
                self.add_error(
                    method.sig.span(),
                    "Method 'deploy' is reserved and generated automatically. Define 'constructor' instead to create initialization logic".to_string(),
                );
            }
            // Special methods with custom handling
            CONSTRUCTOR_METHOD => self.handle_constructor_method(method),
            FALLBACK_METHOD => self.handle_fallback_method(&method.sig, method.sig.span()),
            // Regular user-defined methods
            _ => self.handle_regular_method(method),
        }

        visit::visit_trait_item_fn(self, method);
    }
}

// Implementation for MethodCollector<ImplItemFn>
impl Visit<'_> for MethodCollector<ImplItemFn> {
    fn visit_impl_item_fn(&mut self, method: &ImplItemFn) {
        match method.sig.ident.to_string().as_str() {
            // Reserved methods that cannot be defined by user in router
            MAIN_METHOD => {
                self.add_error(
                    method.sig.span(),
                    "Method 'main' is reserved and generated automatically by the router. It serves as the dispatch entry point for all contract calls".to_string(),
                );
            }
            DEPLOY_METHOD => {
                self.add_error(
                    method.sig.span(),
                    "Method 'deploy' is reserved and generated automatically. Define 'constructor' instead to create initialization logic".to_string(),
                );
            }
            // Special methods with custom handling
            CONSTRUCTOR_METHOD => {
                let is_public = self.is_trait_impl || matches!(method.vis, Visibility::Public(_));
                if is_public {
                    self.handle_constructor_method(method);
                }
            }
            FALLBACK_METHOD => self.handle_fallback_method(&method.sig, method.sig.span()),
            // Regular user-defined methods
            _ => {
                let is_public = self.is_trait_impl || matches!(method.vis, Visibility::Public(_));
                if is_public {
                    self.handle_regular_method(method);
                }
            }
        }

        visit::visit_impl_item_fn(self, method);
    }
}

// Implement From for ImplItemFn
impl From<ImplItemFn> for ParsedMethod<ImplItemFn> {
    fn from(function: ImplItemFn) -> Self {
        match Self::new(function) {
            Ok(method) => method,
            Err(err) => abort_call_site!("Failed to parse method: {}", err),
        }
    }
}

// Implement From for reference to ImplItemFn
impl From<&ImplItemFn> for ParsedMethod<ImplItemFn> {
    fn from(function: &ImplItemFn) -> Self {
        match Self::from_ref(function) {
            Ok(method) => method,
            Err(err) => abort_call_site!("Failed to parse method from reference: {}", err),
        }
    }
}

// Implement From for TraitItemFn
impl From<TraitItemFn> for ParsedMethod<TraitItemFn> {
    fn from(function: TraitItemFn) -> Self {
        match Self::new(function) {
            Ok(method) => method,
            Err(err) => abort_call_site!("Failed to parse method: {}", err),
        }
    }
}

// Implement From for reference to TraitItemFn
impl From<&TraitItemFn> for ParsedMethod<TraitItemFn> {
    fn from(function: &TraitItemFn) -> Self {
        match Self::from_ref(function) {
            Ok(method) => method,
            Err(err) => abort_call_site!("Failed to parse method from reference: {}", err),
        }
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use syn::parse_quote;

    #[test]
    fn test_validate_fallback_signature() {
        let collector = MethodCollector::<TraitItemFn>::new(Span::call_site());

        // Valid fallback signature
        let valid_sig: Signature = parse_quote! {
            fn fallback(&self)
        };
        assert!(collector.validate_fallback_signature(&valid_sig));

        // Invalid: with parameters
        let invalid_with_params: Signature = parse_quote! {
            fn fallback(&self, param: u32)
        };
        assert!(!collector.validate_fallback_signature(&invalid_with_params));

        // Invalid: with return type
        let invalid_with_return: Signature = parse_quote! {
            fn fallback(&self) -> u32
        };
        assert!(!collector.validate_fallback_signature(&invalid_with_return));

        // Invalid: without self
        let invalid_without_self: Signature = parse_quote! {
            fn fallback()
        };
        assert!(!collector.validate_fallback_signature(&invalid_without_self));
    }

    #[test]
    fn test_function_id_collision() {
        let mut collector = MethodCollector::<TraitItemFn>::new(Span::call_site());

        // Create two trait methods with the same function ID selector
        let trait_fn1: TraitItemFn = parse_quote! {
            #[function_id("transfer(address,uint256)", validate(false))]
            fn first_method(&self);
        };

        let trait_fn2: TraitItemFn = parse_quote! {
            #[function_id("transfer(address,uint256)", validate(false))]
            fn second_method(&self);
        };

        // Process the first method
        let parsed1 = ParsedMethod::from(trait_fn1);
        collector.add_method(parsed1);
        assert_eq!(collector.methods.len(), 1);
        assert_eq!(collector.errors.len(), 0);

        // Process the second method with the same function ID
        let parsed2 = ParsedMethod::from(trait_fn2);
        collector.add_method(parsed2);

        // Should detect collision
        assert_eq!(collector.methods.len(), 1); // Second method not added
        assert_eq!(collector.errors.len(), 1); // Error reported
        assert!(collector.has_errors());
    }

    #[test]
    fn test_from_ref_impl() {
        // Create a simple impl item function
        let impl_fn: ImplItemFn = parse_quote! {
            pub fn simple_function(&self) {}
        };

        // Test that from_ref works without cloning unnecessarily
        let result = ParsedMethod::from_ref(&impl_fn);
        assert!(result.is_ok());
    }

    #[test]
    fn test_constructor_method_collection() {
        let mut collector = MethodCollector::<ImplItemFn>::new_for_impl(Span::call_site(), false);

        // Create a constructor method
        let constructor_fn: ImplItemFn = parse_quote! {
            pub fn constructor(&mut self, initial_value: u32) {
                self.value = initial_value;
            }
        };

        // Visit the constructor
        collector.visit_impl_item_fn(&constructor_fn);

        // Constructor should be stored separately, not in methods
        assert_eq!(collector.methods.len(), 0);
        assert!(collector.constructor.is_some());

        let constructor = collector.constructor.as_ref().unwrap();
        assert_eq!(constructor.parsed_signature().rust_name(), "constructor");
        // Constructor should have zero function_id
        assert_eq!(constructor.function_id(), [0, 0, 0, 0]);
    }

    #[test]
    fn test_multiple_constructors_error() {
        let mut collector = MethodCollector::<ImplItemFn>::new_for_impl(Span::call_site(), false);

        // Create two constructor methods
        let constructor1: ImplItemFn = parse_quote! {
            pub fn constructor(&mut self, value: u32) {}
        };

        let constructor2: ImplItemFn = parse_quote! {
            pub fn constructor(&mut self, value: u64) {}
        };

        // Visit first constructor
        collector.visit_impl_item_fn(&constructor1);
        assert!(collector.constructor.is_some());
        assert_eq!(collector.errors.len(), 0);

        // Visit second constructor - should produce error
        collector.visit_impl_item_fn(&constructor2);
        assert_eq!(collector.errors.len(), 1);
        assert!(collector.errors[0]
            .to_string()
            .contains("Only one constructor method is allowed"));
    }

    #[test]
    fn test_deploy_method_forbidden() {
        let mut collector = MethodCollector::<ImplItemFn>::new_for_impl(Span::call_site(), false);

        // Create a deploy method (should be forbidden)
        let deploy_fn: ImplItemFn = parse_quote! {
            pub fn deploy(&mut self) {
                println!("This should not be allowed");
            }
        };

        // Visit the deploy method
        collector.visit_impl_item_fn(&deploy_fn);

        // Should have an error about deploy being reserved
        assert_eq!(collector.errors.len(), 1);
        assert!(collector.errors[0]
            .to_string()
            .contains("'deploy' is reserved"));
        // Should not be added to methods or constructor
        assert_eq!(collector.methods.len(), 0);
        assert!(collector.constructor.is_none());
    }

    #[test]
    fn test_constructor_with_regular_methods() {
        let mut collector = MethodCollector::<ImplItemFn>::new_for_impl(Span::call_site(), false);

        // Create a mix of methods
        let constructor_fn: ImplItemFn = parse_quote! {
            pub fn constructor(&mut self, initial: u32) {}
        };

        let regular_fn: ImplItemFn = parse_quote! {
            pub fn transfer(&mut self, to: Address, amount: u32) {}
        };

        let another_fn: ImplItemFn = parse_quote! {
            pub fn balance_of(&self, owner: Address) -> u32 {}
        };

        // Visit all methods
        collector.visit_impl_item_fn(&constructor_fn);
        collector.visit_impl_item_fn(&regular_fn);
        collector.visit_impl_item_fn(&another_fn);

        // Should have constructor separate from regular methods
        assert!(collector.constructor.is_some());
        assert_eq!(collector.methods.len(), 2);

        // Check constructor
        let constructor = collector.constructor.as_ref().unwrap();
        assert_eq!(constructor.parsed_signature().rust_name(), "constructor");

        // Check regular methods
        let method_names: Vec<String> = collector
            .methods
            .iter()
            .map(|m| m.parsed_signature().rust_name())
            .collect();
        assert!(method_names.contains(&"transfer".to_string()));
        assert!(method_names.contains(&"balance_of".to_string()));
    }

    #[test]
    fn test_parsed_method_is_constructor() {
        // Test constructor method
        let constructor_fn: ImplItemFn = parse_quote! {
            pub fn constructor(&mut self, value: u32) {}
        };

        let parsed_constructor = ParsedMethod::new_constructor(constructor_fn).unwrap();
        assert!(parsed_constructor.is_constructor());
        assert_eq!(parsed_constructor.function_id(), [0, 0, 0, 0]);

        // Test regular method
        let regular_fn: ImplItemFn = parse_quote! {
            pub fn transfer(&mut self, to: Address, amount: u32) {}
        };

        let parsed_regular = ParsedMethod::new(regular_fn).unwrap();
        assert!(!parsed_regular.is_constructor());
        assert_ne!(parsed_regular.function_id(), [0, 0, 0, 0]);
    }
}
