//! Type checking module.
//!
//! Implements the **checking direction** of bidirectional type checking:
//! given expected parameter types (from function/proc signatures),
//! validate that actual argument types conform.
//!
//! Key entry points:
//! - `check_proc_arg_types` — check built-in proc call arguments
//! - `check_core_fn_arg_types` — check core arithmetic fn arguments
//! - `check_local_fn_call_arg_types` — check local fn variable call arguments
//! - `check_user_fn_arg_types` — check user-defined fn call arguments
//! - `check_function_return_type` — check fn body return type vs declaration
//! - `detect_return_type_hint_from_processed_body` — extract hint-fn return type

use std::cell::RefCell;
use std::collections::HashMap;
use std::sync::Arc;

use crate::calcit::{self, Calcit, CalcitFn, CalcitList, CalcitLocal, CalcitProc, CalcitSyntax, CalcitTypeAnnotation, LocatedWarning};

use super::{
  ScopeTypes, check_enum_tuple_construction, check_tuple_nth_bounds, gen_check_warning, gen_check_warning_code, resolve_type_value,
  tag_annotation,
};

// ---------------------------------------------------------------------------
// Unified arg-type checking loop
// ---------------------------------------------------------------------------

/// Core loop shared by `check_user_fn_arg_types`, `check_local_fn_call_arg_types`,
/// and the proc checking path. Walks `(arg, expected_type)` pairs, resolves
/// actual types from `scope_types`, and emits warnings on mismatch.
///
/// `make_warning` is called with `(arg_index_1based, expected_brief, actual_brief)`
/// and should return the full warning message string.
fn check_arg_types_loop<F>(
  args: &CalcitList,
  expected_types: &[Arc<CalcitTypeAnnotation>],
  scope_types: &ScopeTypes,
  file_ns: &str,
  warning_code: &'static str,
  check_warnings: &RefCell<Vec<LocatedWarning>>,
  make_warning: F,
) where
  F: Fn(usize, &str, &str) -> String,
{
  // Check if we have spreading args — can't check with spread
  for arg in args.iter() {
    if matches!(arg, Calcit::Syntax(CalcitSyntax::ArgSpread, _)) {
      return;
    }
  }

  let mut bindings: HashMap<Arc<str>, Arc<CalcitTypeAnnotation>> = HashMap::new();

  for (idx, (arg, expected_type)) in args.iter().zip(expected_types.iter()).enumerate() {
    if matches!(expected_type.as_ref(), CalcitTypeAnnotation::Dynamic) {
      continue;
    }

    // Handle variadic argument type
    if let CalcitTypeAnnotation::Variadic(inner_type) = expected_type.as_ref() {
      for (rest_idx, rest_arg) in args.iter().skip(idx).enumerate() {
        if let Some(actual_type) = resolve_type_value(rest_arg, scope_types) {
          if !actual_type.as_ref().matches_with_bindings(inner_type.as_ref(), &mut bindings) {
            let expected_str = inner_type.as_ref().to_brief_string();
            let actual_str = actual_type.as_ref().to_brief_string();
            gen_check_warning_code(
              make_warning(idx + rest_idx + 1, &expected_str, &actual_str),
              warning_code,
              file_ns,
              check_warnings,
            );
          }
        }
      }
      return; // Done after variadic
    }

    if let Some(actual_type) = resolve_type_value(arg, scope_types) {
      if !actual_type.as_ref().matches_with_bindings(expected_type.as_ref(), &mut bindings) {
        let expected_str = expected_type.as_ref().to_brief_string();
        let actual_str = actual_type.as_ref().to_brief_string();
        gen_check_warning_code(
          make_warning(idx + 1, &expected_str, &actual_str),
          warning_code,
          file_ns,
          check_warnings,
        );
      }
    }
  }
}

// ---------------------------------------------------------------------------
// Public check functions
// ---------------------------------------------------------------------------

/// Check Proc argument types against type signature
pub(crate) fn check_proc_arg_types(
  proc: &CalcitProc,
  args: &CalcitList,
  scope_types: &ScopeTypes,
  file_ns: &str,
  def_name: &str,
  check_warnings: &RefCell<Vec<LocatedWarning>>,
) {
  // Get type signature for this proc
  let Some(signature) = proc.get_type_signature() else {
    return;
  };

  // Check if we have spreading args
  for arg in args.iter() {
    if matches!(arg, Calcit::Syntax(CalcitSyntax::ArgSpread, _)) {
      return;
    }
  }

  // Check argument count
  let arity = signature.arity();
  let min_count = arity.min;
  let max_count = arity.max.unwrap_or(usize::MAX);
  let has_variadic = arity.max.is_none();
  let actual_count = args.len();

  if !has_variadic {
    if actual_count < min_count || actual_count > max_count {
      let expected_str = if min_count == max_count {
        format!("{min_count}")
      } else {
        format!("{min_count}~{max_count}")
      };
      gen_check_warning(
        format!(
          "[Warn] Proc `{}` expects {} args, got {} in call `({} {})`, at {file_ns}/{def_name}",
          proc.as_ref(),
          expected_str,
          actual_count,
          proc.as_ref(),
          args.iter().map(|a| format!("{a}")).collect::<Vec<_>>().join(" ")
        ),
        file_ns,
        check_warnings,
      );
    }
  } else if actual_count < min_count {
    gen_check_warning(
      format!(
        "[Warn] Proc `{}` expects at least {} args, got {} in call `({} {})`, at {file_ns}/{def_name}",
        proc.as_ref(),
        min_count,
        actual_count,
        proc.as_ref(),
        args.iter().map(|a| format!("{a}")).collect::<Vec<_>>().join(" ")
      ),
      file_ns,
      check_warnings,
    );
  }

  if matches!(
    proc,
    CalcitProc::NativeRecord
      | CalcitProc::NativeRecordPartial
      | CalcitProc::NativeRecordGet
      | CalcitProc::NativeRecordNth
      | CalcitProc::NativeLooseRecord
  ) {
    return;
  }

  if matches!(proc, CalcitProc::NativeEnumTupleNew) {
    check_enum_tuple_construction(args, scope_types, file_ns, def_name, check_warnings);
    return;
  }

  if matches!(proc, CalcitProc::NativeTupleNth) {
    check_tuple_nth_bounds(args, scope_types, file_ns, def_name, check_warnings);
    return;
  }

  // Use the unified checker for Proc arg types. Proc's arg_types may have
  // Optional wrappers and variadic markers handled inside the loop.
  let mut bindings: HashMap<Arc<str>, Arc<CalcitTypeAnnotation>> = HashMap::new();

  for (idx, (arg, expected_type)) in args.iter().zip(signature.arg_types.iter()).enumerate() {
    if matches!(expected_type.as_ref(), CalcitTypeAnnotation::Variadic(_)) {
      break;
    }

    let base_type = match expected_type.as_ref() {
      CalcitTypeAnnotation::Optional(inner) => inner,
      _ => expected_type,
    };

    if matches!(base_type.as_ref(), CalcitTypeAnnotation::Dynamic) {
      continue;
    }

    if let Some(actual_type) = resolve_type_value(arg, scope_types) {
      if !actual_type.as_ref().matches_with_bindings(base_type.as_ref(), &mut bindings) {
        let expected_str = base_type.as_ref().to_brief_string();
        let actual_str = actual_type.as_ref().to_brief_string();
        gen_check_warning_code(
          format!(
            "[Warn] Proc `{}` arg {} expects type `{expected_str}`, but got `{actual_str}` in call at {file_ns}/{def_name}",
            proc.as_ref(),
            idx + 1
          ),
          "W_PROC_ARG_TYPE_MISMATCH",
          file_ns,
          check_warnings,
        );
      }
    }
  }
}

/// Check core arithmetic fn arguments (hardcoded :number for +, -, *, /).
pub(crate) fn check_core_fn_arg_types(
  fn_info: &CalcitFn,
  args: &CalcitList,
  scope_types: &ScopeTypes,
  file_ns: &str,
  def_name: &str,
  check_warnings: &RefCell<Vec<LocatedWarning>>,
) {
  if fn_info.def_ns.as_ref() != calcit::CORE_NS {
    return;
  }

  if !fn_info.arg_types.is_empty()
    && fn_info
      .arg_types
      .iter()
      .any(|t| !matches!(t.as_ref(), CalcitTypeAnnotation::Dynamic))
  {
    return;
  }

  let needs_number_args = matches!(fn_info.name.as_ref(), "+" | "-" | "*" | "/");
  if !needs_number_args {
    return;
  }

  let expected_type = tag_annotation("number");

  for (idx, arg) in args.iter().enumerate() {
    if let Some(actual_type) = resolve_type_value(arg, scope_types) {
      if !actual_type.as_ref().matches_annotation(expected_type.as_ref()) {
        let actual_str = actual_type.as_ref().to_brief_string();
        gen_check_warning_code(
          format!(
            "[Warn] Function `calcit.core/{}` arg {} expects type `:number`, but got `{actual_str}` in call at {file_ns}/{def_name}",
            fn_info.name,
            idx + 1
          ),
          "W_CORE_FN_ARG_TYPE_MISMATCH",
          file_ns,
          check_warnings,
        );
      }
    }
  }
}

/// Check argument types when calling a local variable with a known Fn type.
pub(crate) fn check_local_fn_call_arg_types(
  local: &CalcitLocal,
  args: &CalcitList,
  scope_types: &ScopeTypes,
  file_ns: &str,
  def_name: &str,
  check_warnings: &RefCell<Vec<LocatedWarning>>,
) {
  let local_type = if matches!(*local.type_info, CalcitTypeAnnotation::Dynamic) {
    match scope_types.get(&local.sym) {
      Some(t) => t.as_ref(),
      None => return,
    }
  } else {
    &*local.type_info
  };

  let CalcitTypeAnnotation::Fn(fn_annot) = local_type else {
    return;
  };

  if fn_annot.arg_types.is_empty()
    || fn_annot
      .arg_types
      .iter()
      .all(|t| matches!(t.as_ref(), CalcitTypeAnnotation::Dynamic))
  {
    return;
  }

  let local_sym = local.sym.clone();
  let def_name = def_name.to_owned();
  let file_ns_owned = file_ns.to_owned();
  check_arg_types_loop(
    args,
    &fn_annot.arg_types,
    scope_types,
    file_ns,
    "W_LOCAL_FN_ARG_TYPE_MISMATCH",
    check_warnings,
    |arg_idx, expected_str, actual_str| {
      format!(
        "[Warn] calling `{local_sym}` arg {arg_idx} expects type `{expected_str}`, but got `{actual_str}` at {file_ns_owned}/{def_name}"
      )
    },
  );
}

/// Check user-defined function argument types against type annotations.
pub(crate) fn check_user_fn_arg_types(
  fn_info: &CalcitFn,
  args: &CalcitList,
  scope_types: &ScopeTypes,
  file_ns: &str,
  def_name: &str,
  check_warnings: &RefCell<Vec<LocatedWarning>>,
) {
  if fn_info.arg_types.is_empty() || fn_info.arg_types.iter().all(|t| matches!(**t, CalcitTypeAnnotation::Dynamic)) {
    return;
  }

  let fn_def_ns = fn_info.def_ns.clone();
  let fn_name = fn_info.name.clone();
  let def_name = def_name.to_owned();
  let file_ns_owned = file_ns.to_owned();
  check_arg_types_loop(
    args,
    &fn_info.arg_types,
    scope_types,
    file_ns,
    "W_FN_ARG_TYPE_MISMATCH",
    check_warnings,
    |arg_idx, expected_str, actual_str| {
      format!(
        "[Warn] Function `{fn_def_ns}/{fn_name}` arg {arg_idx} expects type `{expected_str}`, but got `{actual_str}` in call at {file_ns_owned}/{def_name}"
      )
    },
  );
}

// ---------------------------------------------------------------------------
// Return-type checking
// ---------------------------------------------------------------------------

/// Extract return type hint from processed function body.
pub(crate) fn detect_return_type_hint_from_processed_body(processed_body: &[Calcit]) -> Arc<CalcitTypeAnnotation> {
  for form in processed_body {
    if let Some(hint) = CalcitTypeAnnotation::extract_return_type_from_hint_form(form) {
      return hint;
    }
  }
  crate::calcit::DYNAMIC_TYPE.clone()
}

/// Check function return type matches declared return_type.
pub(crate) fn check_function_return_type(
  fn_body: &[Calcit],
  declared_return_type: &Arc<CalcitTypeAnnotation>,
  scope_types: &ScopeTypes,
  file_ns: &str,
  def_name: &str,
  check_warnings: &RefCell<Vec<LocatedWarning>>,
) {
  if matches!(**declared_return_type, CalcitTypeAnnotation::Dynamic) {
    return;
  }

  if declared_return_type.contains_type_var() {
    return;
  }

  if fn_body.is_empty() {
    return;
  }

  let last_expr = &fn_body[fn_body.len() - 1];

  let Some(actual_type) = resolve_type_value(last_expr, scope_types) else {
    return;
  };

  if matches!(actual_type.as_ref(), CalcitTypeAnnotation::Dynamic | CalcitTypeAnnotation::DynFn) {
    return;
  }

  if !actual_type.as_ref().matches_annotation(declared_return_type.as_ref()) {
    let expected_str = declared_return_type.as_ref().to_brief_string();
    let actual_str = actual_type.as_ref().to_brief_string();
    gen_check_warning_code(
      format!("[Warn] Function `{file_ns}/{def_name}` declares return type `{expected_str}`, but body returns `{actual_str}`"),
      "W_FN_RETURN_TYPE_MISMATCH",
      file_ns,
      check_warnings,
    );
  }
}
