use crate::*;
use serde_json::{json, Value};
use pgrx::{JsonB, pg_test}; 

// Helper macro for asserting success (no changes needed, but ensure it's present)
macro_rules! assert_success_with_json {
  ($result_jsonb:expr, $fmt:literal $(, $($args:tt)*)?) => {
    let condition_result: Option<bool> = $result_jsonb.0.get("success").and_then(Value::as_bool);
    if condition_result != Some(true) {
      let base_msg = format!($fmt $(, $($args)*)?);
      let pretty_json = serde_json::to_string_pretty(&$result_jsonb.0)
        .unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", $result_jsonb.0));
      let panic_msg = format!("Assertion Failed (expected success): {}\nResult JSON:\n{}", base_msg, pretty_json);
      panic!("{}", panic_msg);
    }
  };
  // Simpler version without message
  ($result_jsonb:expr) => {
    let condition_result: Option<bool> = $result_jsonb.0.get("success").and_then(Value::as_bool);
    if condition_result != Some(true) {
      let pretty_json = serde_json::to_string_pretty(&$result_jsonb.0)
        .unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", $result_jsonb.0));
      let panic_msg = format!("Assertion Failed (expected success)\nResult JSON:\n{}", pretty_json);
      panic!("{}", panic_msg);
    }
  };
}

// Updated helper macro for asserting failed JSON results with the new flat error structure
macro_rules! assert_failure_with_json {
  // --- Arms with error count and message substring check ---
  // With custom message:
  ($result:expr, $expected_error_count:expr, $expected_first_message_contains:expr, $fmt:literal $(, $($args:tt)*)?) => {
    let json_result = &$result.0;
    let success = json_result.get("success").and_then(Value::as_bool);
    let error_val_opt = json_result.get("error"); // Changed key
    let base_msg = format!($fmt $(, $($args)*)?);

    if success != Some(false) {
      let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
      panic!("Assertion Failed (expected failure, success was not false): {}\nResult JSON:\n{}", base_msg, pretty_json);
    }
    match error_val_opt {
      Some(error_val) => {
        if error_val.is_array() {
          let errors_array = error_val.as_array().unwrap();
          if errors_array.len() != $expected_error_count {
            let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
            panic!("Assertion Failed (wrong error count): Expected {} errors, got {}. {}\nResult JSON:\n{}", $expected_error_count, errors_array.len(), base_msg, pretty_json);
          }
          if $expected_error_count > 0 {
            let first_error_message = errors_array[0].get("message").and_then(Value::as_str);
            match first_error_message {
              Some(msg) => {
                if !msg.contains($expected_first_message_contains) {
                  let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
                  panic!("Assertion Failed (first error message mismatch): Expected contains '{}', got: '{}'. {}\nResult JSON:\n{}", $expected_first_message_contains, msg, base_msg, pretty_json);
                }
              }
              None => {
                let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
                panic!("Assertion Failed (first error in array has no 'message' string): {}\nResult JSON:\n{}", base_msg, pretty_json);
              }
            }
          }
        } else if error_val.is_object() {
            // Handle single error object case (like 'schema not found')
            if $expected_error_count != 1 {
              let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
              panic!("Assertion Failed (wrong error count): Expected {} errors, but got a single error object. {}\nResult JSON:\n{}", $expected_error_count, base_msg, pretty_json);
            }
            let message = error_val.get("message").and_then(Value::as_str);
            match message {
              Some(msg) => {
                if !msg.contains($expected_first_message_contains) {
                  let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
                  panic!("Assertion Failed (error message mismatch): Expected object message contains '{}', got: '{}'. {}\nResult JSON:\n{}", $expected_first_message_contains, msg, base_msg, pretty_json);
                }
              }
              None => {
                let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
                panic!("Assertion Failed (error object has no 'message' string): {}\nResult JSON:\n{}", base_msg, pretty_json);
              }
            }
        } else {
            let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
            panic!("Assertion Failed ('error' value is not an array or object): {}\nResult JSON:\n{}", base_msg, pretty_json);
        }
      }
      None => {
         let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
         panic!("Assertion Failed (expected 'error' key, but none found): {}\nResult JSON:\n{}", base_msg, pretty_json);
      }
    }
  };
  // Without custom message (calls the one above with ""):
  ($result:expr, $expected_error_count:expr, $expected_first_message_contains:expr) => {
    assert_failure_with_json!($result, $expected_error_count, $expected_first_message_contains, "");
  };

  // --- Arms with error count check only ---
  // With custom message:
  ($result:expr, $expected_error_count:expr, $fmt:literal $(, $($args:tt)*)?) => {
    let json_result = &$result.0;
    let success = json_result.get("success").and_then(Value::as_bool);
    let error_val_opt = json_result.get("error"); // Changed key
    let base_msg = format!($fmt $(, $($args)*)?);

    if success != Some(false) {
      let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
      panic!("Assertion Failed (expected failure, success was not false): {}\nResult JSON:\n{}", base_msg, pretty_json);
    }
     match error_val_opt {
      Some(error_val) => {
        if error_val.is_array() {
            let errors_array = error_val.as_array().unwrap();
            if errors_array.len() != $expected_error_count {
                let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
                panic!("Assertion Failed (wrong error count): Expected {} errors, got {}. {}\nResult JSON:\n{}", $expected_error_count, errors_array.len(), base_msg, pretty_json);
            }
        } else if error_val.is_object() {
            if $expected_error_count != 1 {
                let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
                panic!("Assertion Failed (wrong error count): Expected {} errors, but got a single error object. {}\nResult JSON:\n{}", $expected_error_count, base_msg, pretty_json);
            }
            // Count check passes if expected is 1 and got object
        } else {
            let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
            panic!("Assertion Failed ('error' value is not an array or object): {}\nResult JSON:\n{}", base_msg, pretty_json);
        }
      }
      None => {
         let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
         panic!("Assertion Failed (expected 'error' key, but none found): {}\nResult JSON:\n{}", base_msg, pretty_json);
      }
    }
  };
  // Without custom message (calls the one above with ""):
  ($result:expr, $expected_error_count:expr) => {
    assert_failure_with_json!($result, $expected_error_count, "");
  };

  // --- Arms checking failure only (expects at least one error) ---
  // With custom message:
  ($result:expr, $fmt:literal $(, $($args:tt)*)?) => {
    let json_result = &$result.0;
    let success = json_result.get("success").and_then(Value::as_bool);
    let error_val_opt = json_result.get("error"); // Changed key
    let base_msg = format!($fmt $(, $($args)*)?);

    if success != Some(false) {
      let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
      panic!("Assertion Failed (expected failure, success was not false): {}\nResult JSON:\n{}", base_msg, pretty_json);
    }
    match error_val_opt {
      Some(error_val) => {
        if error_val.is_object() {
          // OK: single error object is a failure
        } else if error_val.is_array() {
          if error_val.as_array().unwrap().is_empty() {
            let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
            panic!("Assertion Failed (expected errors, but 'error' array is empty): {}\nResult JSON:\n{}", base_msg, pretty_json);
          }
          // OK: non-empty error array is a failure
        } else {
           let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
           panic!("Assertion Failed ('error' value is not an array or object): {}\nResult JSON:\n{}", base_msg, pretty_json);
        }
      }
      None => {
         let pretty_json = serde_json::to_string_pretty(&json_result).unwrap_or_else(|_| format!("(Failed to pretty-print JSON: {:?})", json_result));
         panic!("Assertion Failed (expected 'error' key, but none found): {}\nResult JSON:\n{}", base_msg, pretty_json);
      }
    }
  };
  // Without custom message (calls the one above with ""):
  ($result:expr) => {
    assert_failure_with_json!($result, "");
  };
}


fn jsonb(val: Value) -> JsonB {
  JsonB(val)
}

#[pg_test]
fn test_cache_and_validate_json_schema() {
  clear_json_schemas(); // Call clear directly
  let schema_id = "my_schema";
  let schema = json!({
    "type": "object",
    "properties": {
      "name": { "type": "string" },
      "age": { "type": "integer", "minimum": 0 }
    },
    "required": ["name", "age"]
  });
  let valid_instance = json!({ "name": "Alice", "age": 30 });
  let invalid_instance_type = json!({ "name": "Bob", "age": -5 });
  let invalid_instance_missing = json!({ "name": "Charlie" });

  let cache_result = cache_json_schema(schema_id, jsonb(schema.clone()));
  assert_success_with_json!(cache_result, "Cache operation should succeed.");

  let valid_result = validate_json_schema(schema_id, jsonb(valid_instance));
  assert_success_with_json!(valid_result, "Validation of valid instance should succeed.");

  // Invalid type
  let invalid_result_type = validate_json_schema(schema_id, jsonb(invalid_instance_type));
  assert_failure_with_json!(invalid_result_type, 1, "must be >=0", "Validation with invalid type should fail.");
  let errors_type = invalid_result_type.0["error"].as_array().unwrap(); // Check 'error', expect array
  assert_eq!(errors_type[0]["instance_path"], "/age");
  assert_eq!(errors_type[0]["schema_path"], "urn:my_schema#/properties/age");

  // Missing field
  let invalid_result_missing = validate_json_schema(schema_id, jsonb(invalid_instance_missing));
  assert_failure_with_json!(invalid_result_missing, 1, "missing properties 'age'", "Validation with missing field should fail.");
  let errors_missing = invalid_result_missing.0["error"].as_array().unwrap(); // Check 'error', expect array
  assert_eq!(errors_missing[0]["instance_path"], "");
  assert_eq!(errors_missing[0]["schema_path"], "urn:my_schema#");

  // Schema not found
  let non_existent_id = "non_existent_schema";
  let invalid_schema_result = validate_json_schema(non_existent_id, jsonb(json!({})));
  assert_failure_with_json!(invalid_schema_result, 1, "Schema with id 'non_existent_schema' not found", "Validation with non-existent schema should fail.");
  // Check 'error' is an object for 'schema not found'
  let error_notfound_obj = invalid_schema_result.0["error"].as_object().expect("'error' should be an object for schema not found");
  assert!(error_notfound_obj.contains_key("message")); // Check message exists
  // Removed checks for schema_path/instance_path as they aren't added in lib.rs for this case
}

#[pg_test]
fn test_validate_json_schema_not_cached() {
  clear_json_schemas(); // Call clear directly
  let instance = json!({ "foo": "bar" });
  let result = validate_json_schema("non_existent_schema", jsonb(instance));
  // Use the updated macro, expecting count 1 and specific message (handles object case)
  assert_failure_with_json!(result, 1, "Schema with id 'non_existent_schema' not found", "Validation with non-existent schema should fail.");
}

#[pg_test]
fn test_cache_invalid_json_schema() {
  clear_json_schemas(); // Call clear directly
  let schema_id = "invalid_schema";
  // Schema with an invalid type *value*
  let invalid_schema = json!({
    "$id": "urn:invalid_schema",
    "type": ["invalid_type_value"] 
  });

  let cache_result = cache_json_schema(schema_id, jsonb(invalid_schema));

  // Expect 2 leaf errors because the meta-schema validation fails at the type value
  // and within the type array itself.
  assert_failure_with_json!(
    cache_result,
    2, // Expect exactly two leaf errors
    "value must be one of", // Check message substring (present in both)
    "Caching invalid schema should fail with specific meta-schema validation errors."
  );

  // Ensure the error is an array and check specifics
  let error_array = cache_result.0["error"].as_array().expect("Error field should be an array");
  assert_eq!(error_array.len(), 2);
  // Note: Order might vary depending on boon's internal processing, check both possibilities or sort.
  // Assuming the order shown in the logs for now:
  assert_eq!(error_array[0]["instance_path"], "/type");
  assert!(error_array[0]["message"].as_str().unwrap().contains("value must be one of"));
  assert_eq!(error_array[1]["instance_path"], "/type/0");
  assert!(error_array[1]["message"].as_str().unwrap().contains("value must be one of"));
}

#[pg_test]
fn test_validate_json_schema_detailed_validation_errors() {
  clear_json_schemas(); // Call clear directly
  let schema_id = "detailed_errors";
  let schema = json!({
    "type": "object",
    "properties": {
      "address": {
        "type": "object",
        "properties": {
          "street": { "type": "string" },
          "city": { "type": "string", "maxLength": 10 }
        },
        "required": ["street", "city"]
      }
    },
    "required": ["address"]
  });
  let _ = cache_json_schema(schema_id, jsonb(schema));

  let invalid_instance = json!({
    "address": {
      "street": 123, // Wrong type
      "city": "Supercalifragilisticexpialidocious" // Too long
    }
  });

  let result = validate_json_schema(schema_id, jsonb(invalid_instance));

  // Update: Expect 2 errors again, as boon reports both nested errors.
  assert_failure_with_json!(result, 2);

}

#[pg_test]
fn test_validate_json_schema_oneof_validation_errors() {
  clear_json_schemas(); // Call clear directly
  let schema_id = "oneof_schema";
  let schema = json!({
    "oneOf": [
      { // Option 1: Object with string prop
        "type": "object",
        "properties": {
          "string_prop": { "type": "string", "maxLength": 5 }
        },
        "required": ["string_prop"]
      },
      { // Option 2: Object with number prop
        "type": "object",
        "properties": {
          "number_prop": { "type": "number", "minimum": 10 }
        },
        "required": ["number_prop"]
      }
    ]
  });

  let _ = cache_json_schema(schema_id, jsonb(schema));

  // --- Test case 1: Fails string maxLength (in branch 0) AND missing number_prop (in branch 1) ---
  let invalid_string_instance = json!({ "string_prop": "toolongstring" });
  let result_invalid_string = validate_json_schema(schema_id, jsonb(invalid_string_instance));
  // Expect 2 leaf errors. Check count only with the macro.
  assert_failure_with_json!(result_invalid_string, 2);
  // Explicitly check that both expected errors are present, ignoring order
  let errors_string = result_invalid_string.0["error"].as_array().expect("Expected error array for invalid string");
  assert!(errors_string.iter().any(|e| e["instance_path"] == "/string_prop" && e["message"].as_str().unwrap().contains("length must be <=5")), "Missing maxLength error");
  assert!(errors_string.iter().any(|e| e["instance_path"] == "" && e["message"].as_str().unwrap().contains("missing properties 'number_prop'")), "Missing number_prop required error");

  // --- Test case 2: Fails number minimum (in branch 1) AND missing string_prop (in branch 0) ---
  let invalid_number_instance = json!({ "number_prop": 5 });
  let result_invalid_number = validate_json_schema(schema_id, jsonb(invalid_number_instance));
  // Expect 2 leaf errors. Check count only with the macro.
  assert_failure_with_json!(result_invalid_number, 2);
  // Explicitly check that both expected errors are present, ignoring order
  let errors_number = result_invalid_number.0["error"].as_array().expect("Expected error array for invalid number");
  assert!(errors_number.iter().any(|e| e["instance_path"] == "/number_prop" && e["message"].as_str().unwrap().contains("must be >=10")), "Missing minimum error");
  assert!(errors_number.iter().any(|e| e["instance_path"] == "" && e["message"].as_str().unwrap().contains("missing properties 'string_prop'")), "Missing string_prop required error");

  // --- Test case 3: Fails type check (not object) for both branches ---
  // Input: boolean, expected object for both branches
  let invalid_bool_instance = json!(true); // Not an object
  let result_invalid_bool = validate_json_schema(schema_id, jsonb(invalid_bool_instance));
  // Expect only 1 leaf error after filtering, as both original errors have instance_path ""
  assert_failure_with_json!(result_invalid_bool, 1);
  // Explicitly check that the single remaining error is the type error for the root instance path
  let errors_bool = result_invalid_bool.0["error"].as_array().expect("Expected error array for invalid bool");
  assert_eq!(errors_bool.iter().filter(|e| e["instance_path"] == "" && e["message"].as_str().unwrap().contains("want object")).count(), 1, "Expected one 'want object' error at root after filtering");

  // --- Test case 4: Fails missing required for both branches ---
  // Input: empty object, expected string_prop (branch 0) OR number_prop (branch 1)
  let invalid_empty_obj = json!({});
  let result_empty_obj = validate_json_schema(schema_id, jsonb(invalid_empty_obj));
  // Expect only 1 leaf error after filtering, as both original errors have instance_path ""
  assert_failure_with_json!(result_empty_obj, 1);
  // Explicitly check that the single remaining error is one of the expected missing properties errors
  let errors_empty = result_empty_obj.0["error"].as_array().expect("Expected error array for empty object");
  assert_eq!(errors_empty.len(), 1, "Expected exactly one error after filtering empty object");
  let the_error = &errors_empty[0];
  assert_eq!(the_error["instance_path"], "", "Expected instance_path to be empty string");
  let message = the_error["message"].as_str().unwrap();
  assert!(message.contains("missing properties 'string_prop'") || message.contains("missing properties 'number_prop'"),
          "Error message should indicate missing string_prop or number_prop, got: {}", message);
}

#[pg_test]
fn test_clear_json_schemas() {
  clear_json_schemas(); // Call clear directly
  let schema_id = "schema_to_clear";
  let schema = json!({ "type": "string" });
  cache_json_schema(schema_id, jsonb(schema.clone()));

  let show_result1 = show_json_schemas(); 
  assert!(show_result1.contains(&schema_id.to_string())); 

  clear_json_schemas(); 

  let show_result2 = show_json_schemas(); 
  assert!(show_result2.is_empty());   

  let instance = json!("test");
  let validate_result = validate_json_schema(schema_id, jsonb(instance));
  // Use the updated macro, expecting count 1 and specific message (handles object case)
  assert_failure_with_json!(validate_result, 1, "Schema with id 'schema_to_clear' not found", "Validation should fail after clearing schemas.");
}

#[pg_test]
fn test_show_json_schemas() {
  clear_json_schemas(); // Call clear directly
  let schema_id1 = "schema1";
  let schema_id2 = "schema2";
  let schema = json!({ "type": "boolean" });

  cache_json_schema(schema_id1, jsonb(schema.clone()));
  cache_json_schema(schema_id2, jsonb(schema.clone()));

  let mut result = show_json_schemas(); // Make result mutable
  result.sort(); // Sort for deterministic testing
  assert_eq!(result, vec!["schema1".to_string(), "schema2".to_string()]); // Check exact content
  assert!(result.contains(&schema_id1.to_string())); // Keep specific checks too if desired
  assert!(result.contains(&schema_id2.to_string()));
}