use rhai::{Engine, INT, FLOAT, Dynamic, Map, Array, EvalAltResult, CustomType, TypeBuilder};
use rhai_wrapper::{
    wrap_option_return, wrap_vec_return, wrap_option_vec_return,
    wrap_option_return_result, wrap_vec_return_result, wrap_option_vec_return_result,
    wrap_option_method_result, wrap_vec_method_result, wrap_option_vec_method_result,
    ToRhaiMap, FromRhaiMap
};
use rhai_macros_derive::{ToRhaiMap as ToRhaiMapDerive, FromRhaiMap as FromRhaiMapDerive};
use std::fmt;

// Test structs
#[derive(Debug, Clone, PartialEq, CustomType, ToRhaiMapDerive, FromRhaiMapDerive)]
struct User {
    id: INT,
    name: String,
    age: INT,
}

impl User {
    fn to_rhai_map(&self) -> Map {
        let mut map = Map::new();
        map.insert("id".into(), self.id.into());
        map.insert("name".into(), self.name.clone().into());
        map.insert("age".into(), self.age.into());
        map
    }

    fn from_rhai_map(map: Map) -> Result<Self, String> {
        let id = map.get("id")
            .and_then(|d| d.as_int().ok())
            .ok_or_else(|| "Field 'id' not found or not an INT".to_string())?;
        
        let name = map.get("name")
            .and_then(|d| d.clone().into_string().ok())
            .ok_or_else(|| "Field 'name' not found or not a String".to_string())?;
        
        let age = map.get("age")
            .and_then(|d| d.as_int().ok())
            .ok_or_else(|| "Field 'age' not found or not an INT".to_string())?;
        
        Ok(User { id, name, age })
    }
}

// Mock DB connection type
struct OurDB {
    // Some mock state
    error_mode: bool,
}

impl OurDB {
    fn new(error_mode: bool) -> Self {
        OurDB { error_mode }
    }
}

// Custom error type for testing
#[derive(Debug)]
struct DBError {
    message: String,
}

impl fmt::Display for DBError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "DB Error: {}", self.message)
    }
}

// Test functions for wrap_option_return
fn get_user_by_id(_db: &OurDB, id: INT) -> Option<User> {
    if id > 0 {
        Some(User { id, name: format!("User {}", id), age: 30 })
    } else {
        None
    }
}

// Test functions for wrap_vec_return
fn get_all_users(_db: &OurDB) -> Vec<User> {
    vec![
        User { id: 1, name: "User 1".to_string(), age: 30 },
        User { id: 2, name: "User 2".to_string(), age: 25 },
    ]
}

fn get_users_by_age(_db: &OurDB, min_age: INT) -> Vec<User> {
    vec![
        User { id: 1, name: "User 1".to_string(), age: 30 },
        User { id: 2, name: "User 2".to_string(), age: 25 },
    ].into_iter().filter(|u| u.age >= min_age).collect()
}

fn get_all_user_ids() -> Vec<INT> {
    vec![1, 2, 3]
}

// Test functions for wrap_option_vec_return
fn find_users_by_name(_db: &OurDB, name_part: String) -> Option<Vec<User>> {
    if name_part.is_empty() {
        None
    } else {
        Some(vec![
            User { id: 1, name: format!("{} One", name_part), age: 30 },
            User { id: 2, name: format!("{} Two", name_part), age: 25 },
        ])
    }
}

// Test functions for result-returning wrappers
fn get_user_by_id_result(db: &OurDB, id: INT) -> Result<Option<User>, DBError> {
    if db.error_mode {
        Err(DBError { message: "DB connection error".to_string() })
    } else if id > 0 {
        Ok(Some(User { id, name: format!("User {}", id), age: 30 }))
    } else {
        Ok(None)
    }
}

fn get_all_users_result(db: &OurDB) -> Result<Vec<User>, DBError> {
    if db.error_mode {
        Err(DBError { message: "DB connection error".to_string() })
    } else {
        Ok(vec![
            User { id: 1, name: "User 1".to_string(), age: 30 },
            User { id: 2, name: "User 2".to_string(), age: 25 },
        ])
    }
}

fn find_users_by_name_result(db: &OurDB, name_part: String) -> Result<Option<Vec<User>>, DBError> {
    if db.error_mode {
        Err(DBError { message: "DB connection error".to_string() })
    } else if name_part.is_empty() {
        Ok(None)
    } else {
        Ok(Some(vec![
            User { id: 1, name: format!("{} One", name_part), age: 30 },
            User { id: 2, name: format!("{} Two", name_part), age: 25 },
        ]))
    }
}

// Test methods for method wrappers
struct UserCollection {
    users: Vec<User>,
    error_mode: bool,
}

impl UserCollection {
    fn new(error_mode: bool) -> Self {
        UserCollection {
            users: vec![
                User { id: 1, name: "User 1".to_string(), age: 30 },
                User { id: 2, name: "User 2".to_string(), age: 25 },
            ],
            error_mode,
        }
    }

    fn get_by_id(&self, id: INT) -> Result<Option<User>, DBError> {
        if self.error_mode {
            Err(DBError { message: "Collection error".to_string() })
        } else {
            Ok(self.users.iter().find(|u| u.id == id).cloned())
        }
    }

    fn get_all(&self) -> Result<Vec<User>, DBError> {
        if self.error_mode {
            Err(DBError { message: "Collection error".to_string() })
        } else {
            Ok(self.users.clone())
        }
    }

    fn find_by_name(&self, name_part: String) -> Result<Option<Vec<User>>, DBError> {
        if self.error_mode {
            Err(DBError { message: "Collection error".to_string() })
        } else if name_part.is_empty() {
            Ok(None)
        } else {
            Ok(Some(self.users.iter()
                .filter(|u| u.name.contains(&name_part))
                .cloned()
                .collect()))
        }
    }
}

#[test]
fn test_wrap_option_return() {
    let mut engine = Engine::new();
    let db = OurDB::new(false);
    
    // Register the wrapped function
    engine.register_fn(
        "get_user_by_id_rhai",
        wrap_option_return!(get_user_by_id, &OurDB, INT => User)
    );
    
    // Register the User type
    engine.build_type::<User>();
    
    // Test with existing user
    let script1 = r#"
        let user = get_user_by_id_rhai(db, 1);
        user.id
    "#;
    
    let result1 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script1, 
        "user.id", 
        (db.clone(), )
    ).unwrap();
    
    assert_eq!(result1, 1);
    
    // Test with non-existing user
    let script2 = r#"
        let user = get_user_by_id_rhai(db, 0);
        if user == () { 42 } else { 0 }
    "#;
    
    let result2 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script2, 
        "if user == () { 42 } else { 0 }", 
        (db, )
    ).unwrap();
    
    assert_eq!(result2, 42);
}

#[test]
fn test_wrap_vec_return() {
    let mut engine = Engine::new();
    let db = OurDB::new(false);
    
    // Register the wrapped functions
    engine.register_fn(
        "get_all_users_rhai",
        wrap_vec_return!(get_all_users, &OurDB => User)
    );
    
    engine.register_fn(
        "get_users_by_age_rhai",
        wrap_vec_return!(get_users_by_age, &OurDB, INT => User)
    );
    
    engine.register_fn(
        "get_all_user_ids_rhai",
        wrap_vec_return!(get_all_user_ids, () => INT)
    );
    
    // Register the User type
    engine.build_type::<User>();
    
    // Test get_all_users
    let script1 = r#"
        let users = get_all_users_rhai(db);
        users.len()
    "#;
    
    let result1 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script1, 
        "users.len()", 
        (db.clone(), )
    ).unwrap();
    
    assert_eq!(result1, 2);
    
    // Test get_users_by_age
    let script2 = r#"
        let users = get_users_by_age_rhai(db, 30);
        users.len()
    "#;
    
    let result2 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script2, 
        "users.len()", 
        (db, )
    ).unwrap();
    
    assert_eq!(result2, 1);
    
    // Test get_all_user_ids
    let script3 = r#"
        let ids = get_all_user_ids_rhai();
        ids.len()
    "#;
    
    let result3 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script3, 
        "ids.len()", 
        ()
    ).unwrap();
    
    assert_eq!(result3, 3);
}

#[test]
fn test_wrap_option_vec_return() {
    let mut engine = Engine::new();
    let db = OurDB::new(false);
    
    // Register the wrapped function
    engine.register_fn(
        "find_users_by_name_rhai",
        wrap_option_vec_return!(find_users_by_name, &OurDB, String => User)
    );
    
    // Register the User type
    engine.build_type::<User>();
    
    // Test with found users
    let script1 = r#"
        let users = find_users_by_name_rhai(db, "User");
        users.len()
    "#;
    
    let result1 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script1, 
        "users.len()", 
        (db.clone(), )
    ).unwrap();
    
    assert_eq!(result1, 2);
    
    // Test with no users found
    let script2 = r#"
        let users = find_users_by_name_rhai(db, "");
        if users == () { 42 } else { 0 }
    "#;
    
    let result2 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script2, 
        "if users == () { 42 } else { 0 }", 
        (db, )
    ).unwrap();
    
    assert_eq!(result2, 42);
}

#[test]
fn test_wrap_option_return_result() {
    let mut engine = Engine::new();
    let db_ok = OurDB::new(false);
    let db_err = OurDB::new(true);
    
    // Register the wrapped function
    engine.register_result_fn(
        "get_user_by_id_result_rhai",
        wrap_option_return_result!(get_user_by_id_result, &OurDB, INT => User, DBError)
    );
    
    // Register the User type
    engine.build_type::<User>();
    
    // Test with existing user
    let script1 = r#"
        let user = get_user_by_id_result_rhai(db, 1);
        user.id
    "#;
    
    let result1 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script1, 
        "user.id", 
        (db_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result1, 1);
    
    // Test with non-existing user
    let script2 = r#"
        let user = get_user_by_id_result_rhai(db, 0);
        if user == () { 42 } else { 0 }
    "#;
    
    let result2 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script2, 
        "if user == () { 42 } else { 0 }", 
        (db_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result2, 42);
    
    // Test with error
    let script3 = r#"
        try {
            let user = get_user_by_id_result_rhai(db, 1);
            0
        } catch(err) {
            if err.contains("DB connection error") { 99 } else { 0 }
        }
    "#;
    
    let result3 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script3, 
        "try_catch_block", 
        (db_err, )
    ).unwrap();
    
    assert_eq!(result3, 99);
}

#[test]
fn test_wrap_vec_return_result() {
    let mut engine = Engine::new();
    let db_ok = OurDB::new(false);
    let db_err = OurDB::new(true);
    
    // Register the wrapped function
    engine.register_result_fn(
        "get_all_users_result_rhai",
        wrap_vec_return_result!(get_all_users_result, &OurDB => User, DBError)
    );
    
    // Register the User type
    engine.build_type::<User>();
    
    // Test with successful result
    let script1 = r#"
        let users = get_all_users_result_rhai(db);
        users.len()
    "#;
    
    let result1 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script1, 
        "users.len()", 
        (db_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result1, 2);
    
    // Test with error
    let script2 = r#"
        try {
            let users = get_all_users_result_rhai(db);
            0
        } catch(err) {
            if err.contains("DB connection error") { 99 } else { 0 }
        }
    "#;
    
    let result2 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script2, 
        "try_catch_block", 
        (db_err, )
    ).unwrap();
    
    assert_eq!(result2, 99);
}

#[test]
fn test_wrap_option_vec_return_result() {
    let mut engine = Engine::new();
    let db_ok = OurDB::new(false);
    let db_err = OurDB::new(true);
    
    // Register the wrapped function
    engine.register_result_fn(
        "find_users_by_name_result_rhai",
        wrap_option_vec_return_result!(find_users_by_name_result, &OurDB, String => User, DBError)
    );
    
    // Register the User type
    engine.build_type::<User>();
    
    // Test with found users
    let script1 = r#"
        let users = find_users_by_name_result_rhai(db, "User");
        users.len()
    "#;
    
    let result1 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script1, 
        "users.len()", 
        (db_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result1, 2);
    
    // Test with no users found
    let script2 = r#"
        let users = find_users_by_name_result_rhai(db, "");
        if users == () { 42 } else { 0 }
    "#;
    
    let result2 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script2, 
        "if users == () { 42 } else { 0 }", 
        (db_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result2, 42);
    
    // Test with error
    let script3 = r#"
        try {
            let users = find_users_by_name_result_rhai(db, "User");
            0
        } catch(err) {
            if err.contains("DB connection error") { 99 } else { 0 }
        }
    "#;
    
    let result3 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script3, 
        "try_catch_block", 
        (db_err, )
    ).unwrap();
    
    assert_eq!(result3, 99);
}

#[test]
fn test_wrap_option_method_result() {
    let mut engine = Engine::new();
    let collection_ok = UserCollection::new(false);
    let collection_err = UserCollection::new(true);
    
    // Register the wrapped method
    engine.register_result_fn(
        "get_by_id_rhai",
        wrap_option_method_result!(get_by_id, &UserCollection, INT => User, DBError)
    );
    
    // Register the User type
    engine.build_type::<User>();
    
    // Test with existing user
    let script1 = r#"
        let user = get_by_id_rhai(collection, 1);
        user.id
    "#;
    
    let result1 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script1, 
        "user.id", 
        (collection_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result1, 1);
    
    // Test with non-existing user
    let script2 = r#"
        let user = get_by_id_rhai(collection, 999);
        if user == () { 42 } else { 0 }
    "#;
    
    let result2 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script2, 
        "if user == () { 42 } else { 0 }", 
        (collection_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result2, 42);
    
    // Test with error
    let script3 = r#"
        try {
            let user = get_by_id_rhai(collection, 1);
            0
        } catch(err) {
            if err.contains("Collection error") { 99 } else { 0 }
        }
    "#;
    
    let result3 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script3, 
        "try_catch_block", 
        (collection_err, )
    ).unwrap();
    
    assert_eq!(result3, 99);
}

#[test]
fn test_wrap_vec_method_result() {
    let mut engine = Engine::new();
    let collection_ok = UserCollection::new(false);
    let collection_err = UserCollection::new(true);
    
    // Register the wrapped method
    engine.register_result_fn(
        "get_all_rhai",
        wrap_vec_method_result!(get_all, &UserCollection => User, DBError)
    );
    
    // Register the User type
    engine.build_type::<User>();
    
    // Test with successful result
    let script1 = r#"
        let users = get_all_rhai(collection);
        users.len()
    "#;
    
    let result1 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script1, 
        "users.len()", 
        (collection_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result1, 2);
    
    // Test with error
    let script2 = r#"
        try {
            let users = get_all_rhai(collection);
            0
        } catch(err) {
            if err.contains("Collection error") { 99 } else { 0 }
        }
    "#;
    
    let result2 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script2, 
        "try_catch_block", 
        (collection_err, )
    ).unwrap();
    
    assert_eq!(result2, 99);
}

#[test]
fn test_wrap_option_vec_method_result() {
    let mut engine = Engine::new();
    let collection_ok = UserCollection::new(false);
    let collection_err = UserCollection::new(true);
    
    // Register the wrapped method
    engine.register_result_fn(
        "find_by_name_rhai",
        wrap_option_vec_method_result!(find_by_name, &UserCollection, String => User, DBError)
    );
    
    // Register the User type
    engine.build_type::<User>();
    
    // Test with found users
    let script1 = r#"
        let users = find_by_name_rhai(collection, "User");
        users.len()
    "#;
    
    let result1 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script1, 
        "users.len()", 
        (collection_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result1, 2);
    
    // Test with no users found
    let script2 = r#"
        let users = find_by_name_rhai(collection, "");
        if users == () { 42 } else { 0 }
    "#;
    
    let result2 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script2, 
        "if users == () { 42 } else { 0 }", 
        (collection_ok.clone(), )
    ).unwrap();
    
    assert_eq!(result2, 42);
    
    // Test with error
    let script3 = r#"
        try {
            let users = find_by_name_rhai(collection, "User");
            0
        } catch(err) {
            if err.contains("Collection error") { 99 } else { 0 }
        }
    "#;
    
    let result3 = engine.call_fn::<INT>(
        &mut rhai::Scope::new(), 
        script3, 
        "try_catch_block", 
        (collection_err, )
    ).unwrap();
    
    assert_eq!(result3, 99);
}