///|
/// A small, type-safe builder for JSON Schema fragments.
///
/// LLM tool/function parameters are described with JSON Schema. Writing that
/// schema by hand as a raw JSON string is error-prone; `Schema` lets you build
/// it with typed constructors and render it to a `Json` value via `to_json`.
///
/// Example:
/// ```
/// let params = Schema::object(
///   fields=[
///     ("city", Schema::string(description="City name")),
///     ("units", Schema::enum_(["celsius", "fahrenheit"])),
///   ],
///   required=["city"],
/// )
/// ```
pub(all) enum Schema {
  /// A string, optionally with a description.
  StrSchema(description~ : String?)
  /// A number (floating point).
  NumSchema(description~ : String?)
  /// An integer.
  IntSchema(description~ : String?)
  /// A boolean.
  BoolSchema(description~ : String?)
  /// A string constrained to a fixed set of values.
  EnumSchema(values~ : Array[String], description~ : String?)
  /// An array whose items follow the given schema.
  ArraySchema(items~ : Schema, description~ : String?)
  /// An object with named properties and a list of required property names.
  ObjectSchema(
    fields~ : Array[(String, Schema)],
    required~ : Array[String],
    description~ : String?
  )
}

///|
/// A string schema.
pub fn Schema::string(description? : String) -> Schema {
  StrSchema(description~)
}

///|
/// A number (float) schema.
pub fn Schema::number(description? : String) -> Schema {
  NumSchema(description~)
}

///|
/// An integer schema.
pub fn Schema::integer(description? : String) -> Schema {
  IntSchema(description~)
}

///|
/// A boolean schema.
pub fn Schema::boolean(description? : String) -> Schema {
  BoolSchema(description~)
}

///|
/// A string-enum schema.
pub fn Schema::enum_(values : Array[String], description? : String) -> Schema {
  EnumSchema(values~, description~)
}

///|
/// An array schema over `items`.
pub fn Schema::array(items : Schema, description? : String) -> Schema {
  ArraySchema(items~, description~)
}

///|
/// An object schema. `fields` maps property names to sub-schemas; `required`
/// lists the names that must be present.
pub fn Schema::object(
  fields~ : Array[(String, Schema)],
  required? : Array[String] = [],
  description? : String,
) -> Schema {
  ObjectSchema(fields~, required~, description~)
}

///|
/// Attach the optional `description` field to a schema object, if present.
fn add_description(obj : Map[String, Json], description : String?) -> Unit {
  if description is Some(d) {
    obj["description"] = Json::string(d)
  }
}

///|
pub impl ToJson for Schema with fn to_json(self : Schema) -> Json {
  match self {
    StrSchema(description~) => {
      let obj : Map[String, Json] = { "type": Json::string("string") }
      add_description(obj, description)
      Json::object(obj)
    }
    NumSchema(description~) => {
      let obj : Map[String, Json] = { "type": Json::string("number") }
      add_description(obj, description)
      Json::object(obj)
    }
    IntSchema(description~) => {
      let obj : Map[String, Json] = { "type": Json::string("integer") }
      add_description(obj, description)
      Json::object(obj)
    }
    BoolSchema(description~) => {
      let obj : Map[String, Json] = { "type": Json::string("boolean") }
      add_description(obj, description)
      Json::object(obj)
    }
    EnumSchema(values~, description~) => {
      let vals = []
      for v in values {
        vals.push(Json::string(v))
      }
      let obj : Map[String, Json] = {
        "type": Json::string("string"),
        "enum": Json::array(vals),
      }
      add_description(obj, description)
      Json::object(obj)
    }
    ArraySchema(items~, description~) => {
      let obj : Map[String, Json] = {
        "type": Json::string("array"),
        "items": items.to_json(),
      }
      add_description(obj, description)
      Json::object(obj)
    }
    ObjectSchema(fields~, required~, description~) => {
      let props : Map[String, Json] = {}
      for field in fields {
        let (name, sub) = field
        props[name] = sub.to_json()
      }
      let req = []
      for r in required {
        req.push(Json::string(r))
      }
      let obj : Map[String, Json] = {
        "type": Json::string("object"),
        "properties": Json::object(props),
      }
      if req.length() > 0 {
        obj["required"] = Json::array(req)
      }
      add_description(obj, description)
      Json::object(obj)
    }
  }
}

///|
/// Build a `Tool` from a name, description, and a typed `Schema` for its
/// parameters — a type-safe alternative to writing the JSON Schema by hand.
pub fn Tool::with_schema(
  name : String,
  description : String,
  parameters : Schema,
) -> Tool {
  { name, description, parameters: parameters.to_json() }
}