// ============================================================================
// PFCL Type Checker — App Composition Type Rules
// ============================================================================
//
// Provides type signatures for the app composition intrinsics:
//
//   app.compose       : ∀ Ra Sa Rb Sb. (App<Ra, Sa>, App<Rb, Sb>)
//                         -> App<{a: Ra, b: Rb}, {a: Sa, b: Sb}>
//
//   app.compose_shared : ∀ R Sa Sb. (App<R, Sa>, App<R, Sb>)
//                         -> App<R, {a: Sa, b: Sb}>
//
//   app.with_config    : ∀ R S. (App<R, S>, AppConfig) -> App<R, S>
//
// These are used by `CatalogTypes` implementations to provide type
// information for the inference engine.

use crate::types::*;
use crate::{Box, String, Vec, vec, BTreeMap};

/// Return the polymorphic type scheme for `app.compose`.
///
/// ```text
/// ∀ Ra Sa Rb Sb.
///   (App<Ra, Sa>, App<Rb, Sb>) -> App<{a: Ra, b: Rb}, {a: Sa, b: Sb}>
/// ```
pub fn app_compose_scheme(supply: &mut FreshSupply) -> TypeScheme {
    let ra = supply.fresh();
    let sa = supply.fresh();
    let rb = supply.fresh();
    let sb = supply.fresh();

    let input_a = Type::app(Type::Var(ra), Type::Var(sa));
    let input_b = Type::app(Type::Var(rb), Type::Var(sb));

    // Output resource type: {a: Ra, b: Rb}
    let mut res_fields = BTreeMap::new();
    res_fields.insert("a".into(), Type::Var(ra));
    res_fields.insert("b".into(), Type::Var(rb));
    let res_out = Type::Record(res_fields, None);

    // Output state type: {a: Sa, b: Sb}
    let mut state_fields = BTreeMap::new();
    state_fields.insert("a".into(), Type::Var(sa));
    state_fields.insert("b".into(), Type::Var(sb));
    let state_out = Type::Record(state_fields, None);

    let output = Type::app(res_out, state_out);

    TypeScheme {
        vars: vec![ra, sa, rb, sb],
        ty: Type::Fun(vec![input_a, input_b], Box::new(output)),
    }
}

/// Return the polymorphic type scheme for `app.compose_shared`.
///
/// ```text
/// ∀ R Sa Sb.
///   (App<R, Sa>, App<R, Sb>) -> App<R, {a: Sa, b: Sb}>
/// ```
pub fn app_compose_shared_scheme(supply: &mut FreshSupply) -> TypeScheme {
    let r = supply.fresh();
    let sa = supply.fresh();
    let sb = supply.fresh();

    let input_a = Type::app(Type::Var(r), Type::Var(sa));
    let input_b = Type::app(Type::Var(r), Type::Var(sb));

    // Output state type: {a: Sa, b: Sb}
    let mut state_fields = BTreeMap::new();
    state_fields.insert("a".into(), Type::Var(sa));
    state_fields.insert("b".into(), Type::Var(sb));
    let state_out = Type::Record(state_fields, None);

    // Shared resources: same R
    let output = Type::app(Type::Var(r), state_out);

    TypeScheme {
        vars: vec![r, sa, sb],
        ty: Type::Fun(vec![input_a, input_b], Box::new(output)),
    }
}

/// Return the polymorphic type scheme for `app.with_config`.
///
/// ```text
/// ∀ R S. (App<R, S>, AppConfig) -> App<R, S>
/// ```
pub fn app_with_config_scheme(supply: &mut FreshSupply) -> TypeScheme {
    let r = supply.fresh();
    let s = supply.fresh();

    let input_app = Type::app(Type::Var(r), Type::Var(s));
    let input_config = Type::app_config();
    let output = Type::app(Type::Var(r), Type::Var(s));

    TypeScheme {
        vars: vec![r, s],
        ty: Type::Fun(vec![input_app, input_config], Box::new(output)),
    }
}

/// Check whether a name is one of the app composition intrinsics.
pub fn is_app_intrinsic(name: &str) -> bool {
    matches!(name, "app.compose" | "app.compose_shared" | "app.with_config")
}

/// Look up the type scheme for an app composition intrinsic by name.
/// Returns None if the name is not an app intrinsic.
pub fn app_intrinsic_scheme(name: &str, supply: &mut FreshSupply) -> Option<TypeScheme> {
    match name {
        "app.compose" => Some(app_compose_scheme(supply)),
        "app.compose_shared" => Some(app_compose_shared_scheme(supply)),
        "app.with_config" => Some(app_with_config_scheme(supply)),
        _ => None,
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::unify::unify;

    #[test]
    fn test_app_type_constructor() {
        let app = Type::app(Type::int(), Type::string());
        assert!(app.is_app());
        let (r, s) = app.as_app().unwrap();
        assert_eq!(*r, Type::int());
        assert_eq!(*s, Type::string());
    }

    #[test]
    fn test_non_app_type() {
        assert!(!Type::int().is_app());
        assert!(!Type::list(Type::int()).is_app());
        assert!(Type::int().as_app().is_none());
    }

    #[test]
    fn test_app_handler_expansion() {
        let handler_ty = Type::app_handler_type(Type::int(), Type::string());
        match &handler_ty {
            Type::Fun(inputs, output) => {
                assert_eq!(inputs.len(), 3);
                assert_eq!(inputs[0], Type::int());
                assert_eq!(inputs[1], Type::string());
                assert_eq!(inputs[2], Type::event());
                match output.as_ref() {
                    Type::Tuple(elems) => {
                        assert_eq!(elems.len(), 2);
                        assert_eq!(elems[0], Type::string());
                        assert_eq!(elems[1], Type::commands());
                    }
                    _ => panic!("expected Tuple output"),
                }
            }
            _ => panic!("expected Fun type"),
        }
    }

    #[test]
    fn test_compose_scheme_structure() {
        let mut supply = FreshSupply::new();
        let scheme = app_compose_scheme(&mut supply);

        // Should have 4 quantified variables: Ra, Sa, Rb, Sb
        assert_eq!(scheme.vars.len(), 4);

        // The type should be a function of 2 args
        match &scheme.ty {
            Type::Fun(inputs, output) => {
                assert_eq!(inputs.len(), 2);

                // Both inputs should be App<_, _>
                assert!(inputs[0].is_app());
                assert!(inputs[1].is_app());

                // Output should be App<Record, Record>
                assert!(output.is_app());
                let (res, state) = output.as_app().unwrap();
                assert!(matches!(res, Type::Record(_, None)));
                assert!(matches!(state, Type::Record(_, None)));
            }
            _ => panic!("expected function type"),
        }
    }

    #[test]
    fn test_compose_shared_scheme_structure() {
        let mut supply = FreshSupply::new();
        let scheme = app_compose_shared_scheme(&mut supply);

        // Should have 3 quantified variables: R, Sa, Sb
        assert_eq!(scheme.vars.len(), 3);

        match &scheme.ty {
            Type::Fun(inputs, output) => {
                assert_eq!(inputs.len(), 2);

                // Both inputs share the same R variable
                let (ra, _) = inputs[0].as_app().unwrap();
                let (rb, _) = inputs[1].as_app().unwrap();
                assert_eq!(ra, rb, "shared compose must use same R");

                // Output preserves the shared R
                let (r_out, _) = output.as_app().unwrap();
                assert_eq!(ra, r_out, "output R must match input R");
            }
            _ => panic!("expected function type"),
        }
    }

    #[test]
    fn test_with_config_scheme_structure() {
        let mut supply = FreshSupply::new();
        let scheme = app_with_config_scheme(&mut supply);

        assert_eq!(scheme.vars.len(), 2);

        match &scheme.ty {
            Type::Fun(inputs, output) => {
                assert_eq!(inputs.len(), 2);
                assert!(inputs[0].is_app());
                assert_eq!(inputs[1], Type::app_config());

                // Output has same R, S as input
                let (r_in, s_in) = inputs[0].as_app().unwrap();
                let (r_out, s_out) = output.as_app().unwrap();
                assert_eq!(r_in, r_out);
                assert_eq!(s_in, s_out);
            }
            _ => panic!("expected function type"),
        }
    }

    #[test]
    fn test_compose_unifies_correctly() {
        // Simulate: compose(counter_app, logger_app)
        // counter: App<{label: String}, Int>
        // logger:  App<{prefix: String}, List<String>>
        let mut supply = FreshSupply::new();
        let scheme = app_compose_scheme(&mut supply);

        // Instantiate the scheme with fresh vars
        let mut inst_supply = FreshSupply::new();
        // Advance past the 4 scheme vars
        for _ in 0..4 { inst_supply.fresh(); }

        let concrete_counter = Type::app(
            Type::Record({
                let mut f = BTreeMap::new();
                f.insert("label".into(), Type::string());
                f
            }, None),
            Type::int(),
        );

        let concrete_logger = Type::app(
            Type::Record({
                let mut f = BTreeMap::new();
                f.insert("prefix".into(), Type::string());
                f
            }, None),
            Type::list(Type::string()),
        );

        // The scheme says: (App<Ra, Sa>, App<Rb, Sb>) -> App<{a:Ra, b:Rb}, {a:Sa, b:Sb}>
        // Unify App<Ra, Sa> with concrete_counter, App<Rb, Sb> with concrete_logger
        match &scheme.ty {
            Type::Fun(inputs, output) => {
                let s1 = unify(&inputs[0], &concrete_counter)
                    .expect("counter should unify with App<Ra, Sa>");
                let s2 = unify(
                    &inputs[1].apply_subst(&s1),
                    &concrete_logger,
                ).expect("logger should unify with App<Rb, Sb>");

                let combined = s2.compose(&s1);
                let result = output.apply_subst(&combined);

                // Result should be App<{a: {label: String}, b: {prefix: String}}, {a: Int, b: List<String>}>
                assert!(result.is_app());
                let (res, state) = result.as_app().unwrap();

                match res {
                    Type::Record(fields, _) => {
                        assert!(fields.contains_key("a"), "result res should have field 'a'");
                        assert!(fields.contains_key("b"), "result res should have field 'b'");
                    }
                    _ => panic!("expected record resource type, got {}", res),
                }
                match state {
                    Type::Record(fields, _) => {
                        assert!(fields.contains_key("a"), "result state should have field 'a'");
                        assert!(fields.contains_key("b"), "result state should have field 'b'");
                    }
                    _ => panic!("expected record state type, got {}", state),
                }
            }
            _ => panic!("expected Fun type"),
        }
    }

    #[test]
    fn test_compose_type_error_on_non_app() {
        // Trying to compose a non-app should fail unification
        let mut supply = FreshSupply::new();
        let scheme = app_compose_scheme(&mut supply);

        let not_an_app = Type::string();

        match &scheme.ty {
            Type::Fun(inputs, _) => {
                let result = unify(&inputs[0], &not_an_app);
                assert!(result.is_err(), "String should not unify with App<Ra, Sa>");
            }
            _ => panic!("expected Fun type"),
        }
    }

    #[test]
    fn test_app_intrinsic_lookup() {
        assert!(is_app_intrinsic("app.compose"));
        assert!(is_app_intrinsic("app.compose_shared"));
        assert!(is_app_intrinsic("app.with_config"));
        assert!(!is_app_intrinsic("list.map"));
        assert!(!is_app_intrinsic("int.add"));
    }

    #[test]
    fn test_app_intrinsic_scheme_lookup() {
        let mut supply = FreshSupply::new();
        assert!(app_intrinsic_scheme("app.compose", &mut supply).is_some());
        assert!(app_intrinsic_scheme("app.compose_shared", &mut supply).is_some());
        assert!(app_intrinsic_scheme("app.with_config", &mut supply).is_some());
        assert!(app_intrinsic_scheme("int.add", &mut supply).is_none());
    }

    #[test]
    fn test_app_display() {
        let app = Type::app(Type::int(), Type::string());
        let s = format!("{}", app);
        assert_eq!(s, "App<Int, String>");
    }

    #[test]
    fn test_compose_scheme_display() {
        let mut supply = FreshSupply::new();
        let scheme = app_compose_scheme(&mut supply);
        let s = format!("{}", scheme);
        // Should contain App and arrow
        assert!(s.contains("App"), "scheme display should contain App: {}", s);
        assert!(s.contains("→"), "scheme display should contain arrow: {}", s);
    }
}
