open Ezjsonm

let src = Logs.Src.create "claude.sdk_control" ~doc:"Claude SDK control protocol"
module Log = (val Logs.src_log src : Logs.LOG)

module JU = Json_utils

module Request = struct
  type interrupt = {
    subtype : [`Interrupt];
  }
  
  type permission = {
    subtype : [`Can_use_tool];
    tool_name : string;
    input : value;
    permission_suggestions : Permissions.Update.t list option;
    blocked_path : string option;
  }
  
  type initialize = {
    subtype : [`Initialize];
    hooks : (string * value) list option;
  }
  
  type set_permission_mode = {
    subtype : [`Set_permission_mode];
    mode : Permissions.Mode.t;
  }
  
  type hook_callback = {
    subtype : [`Hook_callback];
    callback_id : string;
    input : value;
    tool_use_id : string option;
  }
  
  type mcp_message = {
    subtype : [`Mcp_message];
    server_name : string;
    message : value;
  }

  type set_model = {
    subtype : [`Set_model];
    model : string;
  }

  type get_server_info = {
    subtype : [`Get_server_info];
  }

  type t =
    | Interrupt of interrupt
    | Permission of permission
    | Initialize of initialize
    | Set_permission_mode of set_permission_mode
    | Hook_callback of hook_callback
    | Mcp_message of mcp_message
    | Set_model of set_model
    | Get_server_info of get_server_info
  
  let interrupt () = Interrupt { subtype = `Interrupt }
  
  let permission ~tool_name ~input ?permission_suggestions ?blocked_path () =
    Permission {
      subtype = `Can_use_tool;
      tool_name;
      input;
      permission_suggestions;
      blocked_path;
    }
  
  let initialize ?hooks () =
    Initialize { subtype = `Initialize; hooks }
  
  let set_permission_mode ~mode =
    Set_permission_mode { subtype = `Set_permission_mode; mode }
  
  let hook_callback ~callback_id ~input ?tool_use_id () =
    Hook_callback {
      subtype = `Hook_callback;
      callback_id;
      input;
      tool_use_id;
    }
  
  let mcp_message ~server_name ~message =
    Mcp_message {
      subtype = `Mcp_message;
      server_name;
      message;
    }

  let set_model ~model =
    Set_model { subtype = `Set_model; model }

  let get_server_info () =
    Get_server_info { subtype = `Get_server_info }

  let to_json = function
    | Interrupt _ ->
        `O [("subtype", `String "interrupt")]
    | Permission p ->
        let fields = [
          ("subtype", `String "can_use_tool");
          ("tool_name", `String p.tool_name);
          ("input", p.input);
        ] in
        let fields = match p.permission_suggestions with
          | Some suggestions ->
              ("permission_suggestions", 
               `A (List.map Permissions.Update.to_json suggestions)) :: fields
          | None -> fields
        in
        let fields = match p.blocked_path with
          | Some path -> ("blocked_path", `String path) :: fields
          | None -> fields
        in
        `O fields
    | Initialize i ->
        let fields = [("subtype", `String "initialize")] in
        let fields = match i.hooks with
          | Some hooks ->
              ("hooks", `O hooks) :: fields
          | None -> fields
        in
        `O fields
    | Set_permission_mode s ->
        `O [
          ("subtype", `String "set_permission_mode");
          ("mode", Permissions.Mode.to_json s.mode);
        ]
    | Hook_callback h ->
        let fields = [
          ("subtype", `String "hook_callback");
          ("callback_id", `String h.callback_id);
          ("input", h.input);
        ] in
        let fields = match h.tool_use_id with
          | Some id -> ("tool_use_id", `String id) :: fields
          | None -> fields
        in
        `O fields
    | Mcp_message m ->
        `O [
          ("subtype", `String "mcp_message");
          ("server_name", `String m.server_name);
          ("message", m.message);
        ]
    | Set_model s ->
        `O [
          ("subtype", `String "set_model");
          ("model", `String s.model);
        ]
    | Get_server_info _ ->
        `O [("subtype", `String "get_server_info")]

  let of_json = function
    | `O fields ->
        let subtype = JU.assoc_string "subtype" fields in
        (match subtype with
        | "interrupt" ->
            Interrupt { subtype = `Interrupt }
        | "can_use_tool" ->
            let tool_name = JU.assoc_string "tool_name" fields in
            let input = List.assoc "input" fields in
            let permission_suggestions =
              match List.assoc_opt "permission_suggestions" fields with
              | Some (`A lst) -> 
                  Some (List.map Permissions.Update.of_json lst)
              | _ -> None
            in
            let blocked_path = JU.assoc_string_opt "blocked_path" fields in
            Permission {
              subtype = `Can_use_tool;
              tool_name;
              input;
              permission_suggestions;
              blocked_path;
            }
        | "initialize" ->
            let hooks =
              match List.assoc_opt "hooks" fields with
              | Some (`O hooks) -> Some hooks
              | _ -> None
            in
            Initialize { subtype = `Initialize; hooks }
        | "set_permission_mode" ->
            let mode = List.assoc "mode" fields |> Permissions.Mode.of_json in
            Set_permission_mode { subtype = `Set_permission_mode; mode }
        | "hook_callback" ->
            let callback_id = JU.assoc_string "callback_id" fields in
            let input = List.assoc "input" fields in
            let tool_use_id = JU.assoc_string_opt "tool_use_id" fields in
            Hook_callback {
              subtype = `Hook_callback;
              callback_id;
              input;
              tool_use_id;
            }
        | "mcp_message" ->
            let server_name = JU.assoc_string "server_name" fields in
            let message = List.assoc "message" fields in
            Mcp_message {
              subtype = `Mcp_message;
              server_name;
              message;
            }
        | "set_model" ->
            let model = JU.assoc_string "model" fields in
            Set_model { subtype = `Set_model; model }
        | "get_server_info" ->
            Get_server_info { subtype = `Get_server_info }
        | _ -> raise (Invalid_argument ("Unknown request subtype: " ^ subtype)))
    | _ -> raise (Invalid_argument "Request.of_json: expected object")
  
  let pp fmt = function
    | Interrupt _ ->
        Fmt.pf fmt "@[<2>Interrupt@]"
    | Permission p ->
        Fmt.pf fmt "@[<2>Permission@ { tool = %S;@ blocked_path = %a }@]"
          p.tool_name Fmt.(option string) p.blocked_path
    | Initialize i ->
        Fmt.pf fmt "@[<2>Initialize@ { hooks = %s }@]"
          (if Option.is_some i.hooks then "present" else "none")
    | Set_permission_mode s ->
        Fmt.pf fmt "@[<2>SetPermissionMode@ { mode = %a }@]"
          Permissions.Mode.pp s.mode
    | Hook_callback h ->
        Fmt.pf fmt "@[<2>HookCallback@ { id = %S;@ tool_use_id = %a }@]"
          h.callback_id Fmt.(option string) h.tool_use_id
    | Mcp_message m ->
        Fmt.pf fmt "@[<2>McpMessage@ { server = %S }@]"
          m.server_name
    | Set_model s ->
        Fmt.pf fmt "@[<2>SetModel@ { model = %S }@]" s.model
    | Get_server_info _ ->
        Fmt.pf fmt "@[<2>GetServerInfo@]"
end

module Response = struct
  type success = {
    subtype : [`Success];
    request_id : string;
    response : value option;
  }
  
  type error = {
    subtype : [`Error];
    request_id : string;
    error : string;
  }
  
  type t =
    | Success of success
    | Error of error
  
  let success ~request_id ?response () =
    Success {
      subtype = `Success;
      request_id;
      response;
    }
  
  let error ~request_id ~error =
    Error {
      subtype = `Error;
      request_id;
      error;
    }
  
  let to_json = function
    | Success s ->
        let fields = [
          ("subtype", `String "success");
          ("request_id", `String s.request_id);
        ] in
        let fields = match s.response with
          | Some resp -> ("response", resp) :: fields
          | None -> fields
        in
        `O fields
    | Error e ->
        `O [
          ("subtype", `String "error");
          ("request_id", `String e.request_id);
          ("error", `String e.error);
        ]
  
  let of_json = function
    | `O fields ->
        let subtype = JU.assoc_string "subtype" fields in
        let request_id = JU.assoc_string "request_id" fields in
        (match subtype with
        | "success" ->
            let response = List.assoc_opt "response" fields in
            Success {
              subtype = `Success;
              request_id;
              response;
            }
        | "error" ->
            let error = JU.assoc_string "error" fields in
            Error {
              subtype = `Error;
              request_id;
              error;
            }
        | _ -> raise (Invalid_argument ("Unknown response subtype: " ^ subtype)))
    | _ -> raise (Invalid_argument "Response.of_json: expected object")
  
  let pp fmt = function
    | Success s ->
        Fmt.pf fmt "@[<2>Success@ { request_id = %S;@ response = %s }@]"
          s.request_id (if Option.is_some s.response then "present" else "none")
    | Error e ->
        Fmt.pf fmt "@[<2>Error@ { request_id = %S;@ error = %S }@]"
          e.request_id e.error
end

type control_request = {
  type_ : [`Control_request];
  request_id : string;
  request : Request.t;
}

type control_response = {
  type_ : [`Control_response];
  response : Response.t;
}

type t =
  | Request of control_request
  | Response of control_response

let create_request ~request_id ~request =
  Request {
    type_ = `Control_request;
    request_id;
    request;
  }

let create_response ~response =
  Response {
    type_ = `Control_response;
    response;
  }

let to_json = function
  | Request r ->
      `O [
        ("type", `String "control_request");
        ("request_id", `String r.request_id);
        ("request", Request.to_json r.request);
      ]
  | Response r ->
      `O [
        ("type", `String "control_response");
        ("response", Response.to_json r.response);
      ]

let of_json = function
  | `O fields ->
      let type_ = JU.assoc_string "type" fields in
      (match type_ with
      | "control_request" ->
          let request_id = JU.assoc_string "request_id" fields in
          let request = List.assoc "request" fields |> Request.of_json in
          Request {
            type_ = `Control_request;
            request_id;
            request;
          }
      | "control_response" ->
          let response = List.assoc "response" fields |> Response.of_json in
          Response {
            type_ = `Control_response;
            response;
          }
      | _ -> raise (Invalid_argument ("Unknown control type: " ^ type_)))
  | _ -> raise (Invalid_argument "of_json: expected object")

let pp fmt = function
  | Request r ->
      Fmt.pf fmt "@[<2>ControlRequest@ { id = %S;@ request = %a }@]"
        r.request_id Request.pp r.request
  | Response r ->
      Fmt.pf fmt "@[<2>ControlResponse@ { %a }@]"
        Response.pp r.response

let log_request req =
  Log.debug (fun m -> m "SDK control request: %a" Request.pp req)

let log_response resp =
  Log.debug (fun m -> m "SDK control response: %a" Response.pp resp)

(** Server information *)
module Server_info = struct
  type t = {
    version : string;
    capabilities : string list;
    commands : string list;
    output_styles : string list;
  }

  let create ~version ~capabilities ~commands ~output_styles =
    { version; capabilities; commands; output_styles }

  let version t = t.version
  let capabilities t = t.capabilities
  let commands t = t.commands
  let output_styles t = t.output_styles

  let of_json = function
    | `O fields ->
        let version = JU.assoc_string "version" fields in
        let capabilities =
          match List.assoc_opt "capabilities" fields with
          | Some (`A lst) -> List.map Ezjsonm.get_string lst
          | _ -> []
        in
        let commands =
          match List.assoc_opt "commands" fields with
          | Some (`A lst) -> List.map Ezjsonm.get_string lst
          | _ -> []
        in
        let output_styles =
          match List.assoc_opt "outputStyles" fields with
          | Some (`A lst) -> List.map Ezjsonm.get_string lst
          | _ -> []
        in
        { version; capabilities; commands; output_styles }
    | _ -> raise (Invalid_argument "Server_info.of_json: expected object")

  let to_json t =
    `O [
      ("version", `String t.version);
      ("capabilities", `A (List.map (fun s -> `String s) t.capabilities));
      ("commands", `A (List.map (fun s -> `String s) t.commands));
      ("outputStyles", `A (List.map (fun s -> `String s) t.output_styles));
    ]

  let pp fmt t =
    Fmt.pf fmt "@[<2>ServerInfo@ { version = %S;@ capabilities = [%a];@ commands = [%a];@ output_styles = [%a] }@]"
      t.version
      Fmt.(list ~sep:(any ", ") (quote string)) t.capabilities
      Fmt.(list ~sep:(any ", ") (quote string)) t.commands
      Fmt.(list ~sep:(any ", ") (quote string)) t.output_styles
end