use std::collections::{HashMap, HashSet};
use serde_json::Value;

pub fn compose(val: &mut Value, errors: &mut Vec<crate::drop::Error>) -> Result<(), String> {
  let mut traits = HashMap::new();
  let mut schemas = HashMap::new();

  // 1. Gather all traits and schemas from enums, types, and puncs arrays
  let arrays = ["enums", "types", "puncs"];
  for arr_name in &arrays {
    if let Some(arr) = val.get(arr_name).and_then(|v| v.as_array()) {
      for item in arr {
        if let Some(item_traits) = item.get("traits").and_then(|v| v.as_object()) {
          for (name, trait_val) in item_traits {
            traits.insert(name.clone(), trait_val.clone());
          }
        }
        if let Some(item_schemas) = item.get("schemas").and_then(|v| v.as_object()) {
          for (name, schema_val) in item_schemas {
            schemas.insert(name.clone(), schema_val.clone());
          }
        }
      }
    }
  }

  // 2. Resolve inclusions recursively in all schema objects
  for arr_name in &arrays {
    if let Some(arr) = val.get_mut(arr_name).and_then(|v| v.as_array_mut()) {
      for item in arr {
        if let Some(item_schemas) = item.get_mut("schemas").and_then(|v| v.as_object_mut()) {
          for (schema_id, schema_val) in item_schemas {
            let mut visited = HashSet::new();
            resolve_in_place(
              schema_val,
              &traits,
              &schemas,
              errors,
              schema_id,
              schema_id,
              &mut visited,
            );
          }
        }
      }
    }
  }

  // 3. Strip the "traits" block from each item in enums, types, puncs so it doesn't serialize
  for arr_name in &arrays {
    if let Some(arr) = val.get_mut(arr_name).and_then(|v| v.as_array_mut()) {
      for item in arr {
        if let Some(obj) = item.as_object_mut() {
          obj.remove("traits");
        }
      }
    }
  }

  Ok(())
}

fn resolve_in_place(
  current: &mut Value,
  traits: &HashMap<String, Value>,
  schemas: &HashMap<String, Value>,
  errors: &mut Vec<crate::drop::Error>,
  schema_id: &str,
  path: &str,
  visited: &mut HashSet<String>,
) {
  if !current.is_object() {
    return;
  }

  let include_opt = current.as_object_mut().and_then(|obj| obj.remove("include"));
  if let Some(include_val) = include_opt {
    if let Some(include_arr) = include_val.as_array() {
      let mut merged_props = serde_json::Map::new();
      let mut merged_required = HashSet::new();
      let mut merged_display = HashSet::new();
      let mut merged_dependencies = serde_json::Map::new();
      let mut merged_pattern_props = serde_json::Map::new();

      // Read current values first to let host override included properties
      if let Some(req) = current.get("required").and_then(|v| v.as_array()) {
        for r in req {
          if let Some(s) = r.as_str() {
            merged_required.insert(s.to_string());
          }
        }
      }
      if let Some(disp) = current.get("display").and_then(|v| v.as_array()) {
        for d in disp {
          if let Some(s) = d.as_str() {
            merged_display.insert(s.to_string());
          }
        }
      }
      if let Some(deps) = current.get("dependencies").and_then(|v| v.as_object()) {
        for (k, v) in deps {
          merged_dependencies.insert(k.clone(), v.clone());
        }
      }
      if let Some(pat_props) = current.get("patternProperties").and_then(|v| v.as_object()) {
        for (k, v) in pat_props {
          merged_pattern_props.insert(k.clone(), v.clone());
        }
      }
      if let Some(props) = current.get("properties").and_then(|v| v.as_object()) {
        for (k, v) in props {
          merged_props.insert(k.clone(), v.clone());
        }
      }

      for inc in include_arr {
        if let Some(inc_name) = inc.as_str() {
          if visited.contains(inc_name) {
            errors.push(crate::drop::Error {
              code: "CIRCULAR_INCLUDE_DETECTED".to_string(),
              message: format!("Circular inclusion detected for '{}'", inc_name),
              details: crate::drop::ErrorDetails {
                schema: Some(schema_id.to_string()),
                path: Some(path.to_string()),
                ..Default::default()
              },
            });
            continue;
          }

          let target_opt = traits.get(inc_name).or_else(|| schemas.get(inc_name));
          if let Some(target_val) = target_opt {
            let mut resolved_target = target_val.clone();
            visited.insert(inc_name.to_string());
            resolve_in_place(
              &mut resolved_target,
              traits,
              schemas,
              errors,
              schema_id,
              &format!("{}/include/{}", path, inc_name),
              visited,
            );
            visited.remove(inc_name);

            // Merge properties (host overrides trait)
            if let Some(target_props) = resolved_target.get("properties").and_then(|v| v.as_object()) {
              for (k, v) in target_props {
                if !merged_props.contains_key(k) {
                  merged_props.insert(k.clone(), v.clone());
                }
              }
            }

            // Merge patternProperties (host overrides trait)
            if let Some(target_pat_props) = resolved_target.get("patternProperties").and_then(|v| v.as_object()) {
              for (k, v) in target_pat_props {
                if !merged_pattern_props.contains_key(k) {
                  merged_pattern_props.insert(k.clone(), v.clone());
                }
              }
            }

            // Merge required
            if let Some(target_req) = resolved_target.get("required").and_then(|v| v.as_array()) {
              for r in target_req {
                if let Some(s) = r.as_str() {
                  merged_required.insert(s.to_string());
                }
              }
            }

            // Merge display
            if let Some(target_disp) = resolved_target.get("display").and_then(|v| v.as_array()) {
              for d in target_disp {
                if let Some(s) = d.as_str() {
                  merged_display.insert(s.to_string());
                }
              }
            }

            // Merge dependencies
            if let Some(target_deps) = resolved_target.get("dependencies").and_then(|v| v.as_object()) {
              for (dep_prop, dep_val) in target_deps {
                if let Some(existing_val) = merged_dependencies.get_mut(dep_prop) {
                  if let (Some(arr_existing), Some(arr_target)) = (existing_val.as_array_mut(), dep_val.as_array()) {
                    let mut set: HashSet<String> = arr_existing.iter().filter_map(|x| x.as_str().map(String::from)).collect();
                    for x in arr_target {
                      if let Some(s) = x.as_str() {
                        if set.insert(s.to_string()) {
                          arr_existing.push(Value::String(s.to_string()));
                        }
                      }
                    }
                  }
                } else {
                  merged_dependencies.insert(dep_prop.clone(), dep_val.clone());
                }
              }
            }

            // Inherit other non-merged schemas/scalars if not defined in host (type, items, cases, family, format, etc.)
            if let Some(obj) = current.as_object_mut() {
              for (k, v) in resolved_target.as_object().unwrap() {
                if k != "properties" && k != "patternProperties" && k != "required" && k != "display" && k != "dependencies" && k != "include" {
                  if !obj.contains_key(k) {
                    obj.insert(k.clone(), v.clone());
                  }
                }
              }
            }
          } else {
            errors.push(crate::drop::Error {
              code: "TRAIT_NOT_FOUND".to_string(),
              message: format!("Trait or schema '{}' not found for inclusion", inc_name),
              details: crate::drop::ErrorDetails {
                schema: Some(schema_id.to_string()),
                path: Some(path.to_string()),
                ..Default::default()
              },
            });
          }
        }
      }

      if let Some(obj) = current.as_object_mut() {
        if !merged_props.is_empty() {
          obj.insert("properties".to_string(), Value::Object(merged_props));
        }
        if !merged_pattern_props.is_empty() {
          obj.insert("patternProperties".to_string(), Value::Object(merged_pattern_props));
        }
        if !merged_required.is_empty() {
          let mut req_vec: Vec<Value> = merged_required.into_iter().map(Value::String).collect();
          req_vec.sort_by(|a, b| a.as_str().unwrap().cmp(b.as_str().unwrap()));
          obj.insert("required".to_string(), Value::Array(req_vec));
        }
        if !merged_display.is_empty() {
          let mut disp_vec: Vec<Value> = merged_display.into_iter().map(Value::String).collect();
          disp_vec.sort_by(|a, b| a.as_str().unwrap().cmp(b.as_str().unwrap()));
          obj.insert("display".to_string(), Value::Array(disp_vec));
        }
        if !merged_dependencies.is_empty() {
          obj.insert("dependencies".to_string(), Value::Object(merged_dependencies));
        }
      }
    }
  }

  // Recursively process children
  if let Some(obj) = current.as_object_mut() {
    if let Some(props) = obj.get_mut("properties").and_then(|v| v.as_object_mut()) {
      for (k, v) in props {
        resolve_in_place(v, traits, schemas, errors, schema_id, &format!("{}/{}", path, k), visited);
      }
    }
    if let Some(pat_props) = obj.get_mut("patternProperties").and_then(|v| v.as_object_mut()) {
      for (k, v) in pat_props {
        resolve_in_place(v, traits, schemas, errors, schema_id, &format!("{}/{}", path, k), visited);
      }
    }
    if let Some(items) = obj.get_mut("items") {
      resolve_in_place(items, traits, schemas, errors, schema_id, &format!("{}/items", path), visited);
    }
    if let Some(prefix_items) = obj.get_mut("prefixItems").and_then(|v| v.as_array_mut()) {
      for (i, v) in prefix_items.iter_mut().enumerate() {
        resolve_in_place(v, traits, schemas, errors, schema_id, &format!("{}/prefixItems/{}", path, i), visited);
      }
    }
    if let Some(additional_props) = obj.get_mut("additionalProperties") {
      resolve_in_place(additional_props, traits, schemas, errors, schema_id, &format!("{}/additionalProperties", path), visited);
    }
    if let Some(one_of) = obj.get_mut("oneOf").and_then(|v| v.as_array_mut()) {
      for (i, v) in one_of.iter_mut().enumerate() {
        resolve_in_place(v, traits, schemas, errors, schema_id, &format!("{}/oneOf/{}", path, i), visited);
      }
    }
    if let Some(contains) = obj.get_mut("contains") {
      resolve_in_place(contains, traits, schemas, errors, schema_id, &format!("{}/contains", path), visited);
    }
    if let Some(not) = obj.get_mut("not") {
      resolve_in_place(not, traits, schemas, errors, schema_id, &format!("{}/not", path), visited);
    }
    if let Some(cases) = obj.get_mut("cases").and_then(|v| v.as_array_mut()) {
      for (i, c_val) in cases.iter_mut().enumerate() {
        if let Some(c_obj) = c_val.as_object_mut() {
          if let Some(when) = c_obj.get_mut("when") {
            resolve_in_place(when, traits, schemas, errors, schema_id, &format!("{}/cases/{}/when", path, i), visited);
          }
          if let Some(then) = c_obj.get_mut("then") {
            resolve_in_place(then, traits, schemas, errors, schema_id, &format!("{}/cases/{}/then", path, i), visited);
          }
          if let Some(else_) = c_obj.get_mut("else") {
            resolve_in_place(else_, traits, schemas, errors, schema_id, &format!("{}/cases/{}/else", path, i), visited);
          }
        }
      }
    }
  }
}