open Ezjsonm
module JU = Json_utils

let src = Logs.Src.create "claude.message" ~doc:"Claude messages"
module Log = (val Logs.src_log src : Logs.LOG)


module User = struct
  type content = 
    | String of string
    | Blocks of Content_block.t list
  
  type t = { content : content }
  
  let create_string s = { content = String s }
  let create_blocks blocks = { content = Blocks blocks }
  
  let create_with_tool_result ~tool_use_id ~content ?is_error () =
    let tool_result = Content_block.tool_result ~tool_use_id ~content ?is_error () in
    { content = Blocks [tool_result] }
  
  let create_mixed ~text ~tool_results =
    let blocks = 
      let text_blocks = match text with
        | Some t -> [Content_block.text t]
        | None -> []
      in
      let tool_blocks = List.map (fun (tool_use_id, content, is_error) ->
        Content_block.tool_result ~tool_use_id ~content ?is_error ()
      ) tool_results in
      text_blocks @ tool_blocks
    in
    { content = Blocks blocks }
  
  let content t = t.content
  
  let as_text t = match t.content with
    | String s -> Some s
    | Blocks _ -> None
  
  let get_blocks t = match t.content with
    | String s -> [Content_block.text s]
    | Blocks blocks -> blocks
  
  let to_json t =
    let content_json = match t.content with
      | String s -> `String s
      | Blocks blocks -> 
          `A (List.map Content_block.to_json blocks)
    in
    `O [
      ("type", `String "user");
      ("message", `O [
        ("role", `String "user");
        ("content", content_json);
      ]);
    ]
  
  let of_json = function
    | `O fields ->
        let message = List.assoc "message" fields in
        let content = match message with
          | `O msg_fields ->
              (match List.assoc "content" msg_fields with
               | `String s -> String s
               | `A blocks -> Blocks (List.map Content_block.of_json blocks)
               | _ -> raise (Invalid_argument "User.of_json: invalid content"))
          | _ -> raise (Invalid_argument "User.of_json: invalid message")
        in
        { content }
    | _ -> raise (Invalid_argument "User.of_json: expected object")
  
  let pp fmt t =
    match t.content with
    | String s -> 
        if String.length s > 60 then
          let truncated = String.sub s 0 57 in
          Fmt.pf fmt "@[<2>User:@ %s...@]" truncated
        else
          Fmt.pf fmt "@[<2>User:@ %S@]" s
    | Blocks blocks ->
        let text_count = List.length (List.filter (function 
          | Content_block.Text _ -> true | _ -> false) blocks) in
        let tool_result_count = List.length (List.filter (function 
          | Content_block.Tool_result _ -> true | _ -> false) blocks) in
        match text_count, tool_result_count with
        | 1, 0 -> 
            let text = List.find_map (function
              | Content_block.Text t -> Some (Content_block.Text.text t)
              | _ -> None) blocks in
            Fmt.pf fmt "@[<2>User:@ %a@]" Fmt.(option string) text
        | 0, 1 ->
            Fmt.pf fmt "@[<2>User:@ [tool result]@]"
        | 0, n when n > 1 ->
            Fmt.pf fmt "@[<2>User:@ [%d tool results]@]" n
        | _ ->
            Fmt.pf fmt "@[<2>User:@ [%d blocks]@]" (List.length blocks)
end

module Assistant = struct
  type error = [
    | `Authentication_failed
    | `Billing_error
    | `Rate_limit
    | `Invalid_request
    | `Server_error
    | `Unknown
  ]

  let error_to_string = function
    | `Authentication_failed -> "authentication_failed"
    | `Billing_error -> "billing_error"
    | `Rate_limit -> "rate_limit"
    | `Invalid_request -> "invalid_request"
    | `Server_error -> "server_error"
    | `Unknown -> "unknown"

  let error_of_string = function
    | "authentication_failed" -> `Authentication_failed
    | "billing_error" -> `Billing_error
    | "rate_limit" -> `Rate_limit
    | "invalid_request" -> `Invalid_request
    | "server_error" -> `Server_error
    | "unknown" | _ -> `Unknown

  type t = {
    content : Content_block.t list;
    model : string;
    error : error option;
  }

  let create ~content ~model ?error () = { content; model; error }
  let content t = t.content
  let model t = t.model
  let error t = t.error
  
  let get_text_blocks t =
    List.filter_map (function
      | Content_block.Text text -> Some (Content_block.Text.text text)
      | _ -> None
    ) t.content
  
  let get_tool_uses t =
    List.filter_map (function
      | Content_block.Tool_use tool -> Some tool
      | _ -> None
    ) t.content
  
  let get_thinking t =
    List.filter_map (function
      | Content_block.Thinking thinking -> Some thinking
      | _ -> None
    ) t.content
  
  let has_tool_use t =
    List.exists (function
      | Content_block.Tool_use _ -> true
      | _ -> false
    ) t.content
  
  let combined_text t =
    String.concat "\n" (get_text_blocks t)
  
  let to_json t =
    let msg_fields = [
      ("content", `A (List.map Content_block.to_json t.content));
      ("model", `String t.model);
    ] in
    let msg_fields = match t.error with
      | Some err -> ("error", `String (error_to_string err)) :: msg_fields
      | None -> msg_fields
    in
    `O [
      ("type", `String "assistant");
      ("message", `O msg_fields);
    ]
  
  let of_json = function
    | `O fields ->
        let message = List.assoc "message" fields in
        let content, model, error = match message with
          | `O msg_fields ->
              let content =
                match List.assoc "content" msg_fields with
                | `A blocks -> List.map Content_block.of_json blocks
                | _ -> raise (Invalid_argument "Assistant.of_json: invalid content")
              in
              let model = JU.assoc_string "model" msg_fields in
              let error =
                match JU.assoc_string_opt "error" msg_fields with
                | Some err_str -> Some (error_of_string err_str)
                | None -> None
              in
              content, model, error
          | _ -> raise (Invalid_argument "Assistant.of_json: invalid message")
        in
        { content; model; error }
    | _ -> raise (Invalid_argument "Assistant.of_json: expected object")
  
  let pp fmt t =
    let text_count = List.length (get_text_blocks t) in
    let tool_count = List.length (get_tool_uses t) in
    let thinking_count = List.length (get_thinking t) in
    match text_count, tool_count, thinking_count with
    | 1, 0, 0 ->
        (* Simple text response *)
        Fmt.pf fmt "@[<2>Assistant@ [%s]:@ %S@]"
          t.model (combined_text t)
    | _, 0, 0 when text_count > 0 ->
        (* Multiple text blocks *)
        Fmt.pf fmt "@[<2>Assistant@ [%s]:@ %d text blocks@]"
          t.model text_count
    | 0, _, 0 when tool_count > 0 ->
        (* Only tool uses *)
        let tools = get_tool_uses t in
        let tool_names = List.map Content_block.Tool_use.name tools in
        Fmt.pf fmt "@[<2>Assistant@ [%s]:@ tools(%a)@]"
          t.model Fmt.(list ~sep:comma string) tool_names
    | _ ->
        (* Mixed content *)
        let parts = [] in
        let parts = if text_count > 0 then 
          Printf.sprintf "%d text" text_count :: parts else parts in
        let parts = if tool_count > 0 then 
          Printf.sprintf "%d tools" tool_count :: parts else parts in
        let parts = if thinking_count > 0 then 
          Printf.sprintf "%d thinking" thinking_count :: parts else parts in
        Fmt.pf fmt "@[<2>Assistant@ [%s]:@ %s@]"
          t.model (String.concat ", " (List.rev parts))
end

module System = struct
  module Data = struct
    (* Store both the raw JSON and provide typed accessors *)
    type t = value  (* The full JSON data *)
    
    let empty = `O []
    
    let of_assoc assoc = `O assoc
    
    let get_string t key = JU.get_field_string_opt t key
    
    let get_int t key = JU.get_field_int_opt t key
    
    let get_bool t key = JU.get_field_bool_opt t key
    
    let get_float t key = JU.get_field_float_opt t key
    
    let get_list t key = 
      match t with
      | `O fields -> 
          (match List.assoc_opt key fields with
          | Some (`A lst) -> Some lst
          | _ -> None)
      | _ -> None
    
    let get_field t key =
      match t with
      | `O fields -> List.assoc_opt key fields
      | _ -> None
    
    let raw_json t = t
    
    let to_json t = t
    let of_json json = json
  end
  
  type t = {
    subtype : string;
    data : Data.t;
  }
  
  let create ~subtype ~data = { subtype; data }
  let subtype t = t.subtype
  let data t = t.data
  
  let to_json t =
    `O [
      ("type", `String "system");
      ("subtype", `String t.subtype);
      ("data", Data.to_json t.data);
    ]
  
  let of_json = function
    | `O fields ->
        let subtype = JU.assoc_string "subtype" fields in
        let data = Data.of_json (try List.assoc "data" fields with Not_found -> `O fields) in
        { subtype; data }
    | _ -> raise (Invalid_argument "System.of_json: expected object")
  
  let pp fmt t =
    match t.subtype with
    | "init" ->
        let session_id = Data.get_string t.data "session_id" in
        let model = Data.get_string t.data "model" in
        let cwd = Data.get_string t.data "cwd" in
        Fmt.pf fmt "@[<2>System.init@ { session_id = %a;@ model = %a;@ cwd = %a }@]"
          Fmt.(option string) session_id
          Fmt.(option string) model
          Fmt.(option string) cwd
    | "error" ->
        let error = Data.get_string t.data "error" in
        Fmt.pf fmt "@[<2>System.error@ { error = %a }@]"
          Fmt.(option string) error
    | _ ->
        Fmt.pf fmt "@[<2>System.%s@ { ... }@]" t.subtype
end

module Result = struct
  module Usage = struct
    type t = value
    
    let create ?input_tokens ?output_tokens ?total_tokens 
               ?cache_creation_input_tokens ?cache_read_input_tokens () =
      let fields = [] in
      let fields = match input_tokens with
        | Some n -> ("input_tokens", `Float (float_of_int n)) :: fields
        | None -> fields in
      let fields = match output_tokens with
        | Some n -> ("output_tokens", `Float (float_of_int n)) :: fields
        | None -> fields in
      let fields = match total_tokens with
        | Some n -> ("total_tokens", `Float (float_of_int n)) :: fields
        | None -> fields in
      let fields = match cache_creation_input_tokens with
        | Some n -> ("cache_creation_input_tokens", `Float (float_of_int n)) :: fields
        | None -> fields in
      let fields = match cache_read_input_tokens with
        | Some n -> ("cache_read_input_tokens", `Float (float_of_int n)) :: fields
        | None -> fields in
      `O fields
    
    let input_tokens t = JU.get_field_int_opt t "input_tokens"
    
    let output_tokens t = JU.get_field_int_opt t "output_tokens"
    
    let total_tokens t = JU.get_field_int_opt t "total_tokens"
    
    let cache_creation_input_tokens t = JU.get_field_int_opt t "cache_creation_input_tokens"
    
    let cache_read_input_tokens t = JU.get_field_int_opt t "cache_read_input_tokens"
    
    let effective_input_tokens t =
      match input_tokens t with
      | None -> 0
      | Some input ->
          let cached = Option.value (cache_read_input_tokens t) ~default:0 in
          max 0 (input - cached)
    
    let total_cost_estimate t ~input_price ~output_price =
      match input_tokens t, output_tokens t with
      | Some input, Some output ->
          let input_cost = float_of_int input *. input_price /. 1_000_000. in
          let output_cost = float_of_int output *. output_price /. 1_000_000. in
          Some (input_cost +. output_cost)
      | _ -> None
    
    let pp fmt t =
      Fmt.pf fmt "@[<2>Usage@ { input = %a;@ output = %a;@ total = %a;@ \
                  cache_creation = %a;@ cache_read = %a }@]"
        Fmt.(option int) (input_tokens t)
        Fmt.(option int) (output_tokens t)
        Fmt.(option int) (total_tokens t)
        Fmt.(option int) (cache_creation_input_tokens t)
        Fmt.(option int) (cache_read_input_tokens t)
    
    let to_json t = t
    let of_json json = json
  end
  
  type t = {
    subtype : string;
    duration_ms : int;
    duration_api_ms : int;
    is_error : bool;
    num_turns : int;
    session_id : string;
    total_cost_usd : float option;
    usage : Usage.t option;
    result : string option;
    structured_output : value option;
  }
  
  let create ~subtype ~duration_ms ~duration_api_ms ~is_error ~num_turns
      ~session_id ?total_cost_usd ?usage ?result ?structured_output () =
    { subtype; duration_ms; duration_api_ms; is_error; num_turns;
      session_id; total_cost_usd; usage; result; structured_output }
  
  let subtype t = t.subtype
  let duration_ms t = t.duration_ms
  let duration_api_ms t = t.duration_api_ms
  let is_error t = t.is_error
  let num_turns t = t.num_turns
  let session_id t = t.session_id
  let total_cost_usd t = t.total_cost_usd
  let usage t = t.usage
  let result t = t.result
  let structured_output t = t.structured_output
  
  let to_json t =
    let fields = [
      ("type", `String "result");
      ("subtype", `String t.subtype);
      ("duration_ms", `Float (float_of_int t.duration_ms));
      ("duration_api_ms", `Float (float_of_int t.duration_api_ms));
      ("is_error", `Bool t.is_error);
      ("num_turns", `Float (float_of_int t.num_turns));
      ("session_id", `String t.session_id);
    ] in
    let fields = match t.total_cost_usd with
      | Some cost -> ("total_cost_usd", `Float cost) :: fields
      | None -> fields
    in
    let fields = match t.usage with
      | Some usage -> ("usage", Usage.to_json usage) :: fields
      | None -> fields
    in
    let fields = match t.result with
      | Some result -> ("result", `String result) :: fields
      | None -> fields
    in
    let fields = match t.structured_output with
      | Some output -> ("structured_output", output) :: fields
      | None -> fields
    in
    `O fields
  
  let of_json = function
    | `O fields ->
        let subtype = JU.assoc_string "subtype" fields in
        let duration_ms = int_of_float (JU.assoc_float "duration_ms" fields) in
        let duration_api_ms = int_of_float (JU.assoc_float "duration_api_ms" fields) in
        let is_error = JU.assoc_bool "is_error" fields in
        let num_turns = int_of_float (JU.assoc_float "num_turns" fields) in
        let session_id = JU.assoc_string "session_id" fields in
        let total_cost_usd = JU.assoc_float_opt "total_cost_usd" fields in
        let usage = Option.map Usage.of_json (List.assoc_opt "usage" fields) in
        let result = JU.assoc_string_opt "result" fields in
        let structured_output = List.assoc_opt "structured_output" fields in
        { subtype; duration_ms; duration_api_ms; is_error; num_turns;
          session_id; total_cost_usd; usage; result; structured_output }
    | _ -> raise (Invalid_argument "Result.of_json: expected object")
  
  let pp fmt t =
    if t.is_error then
      Fmt.pf fmt "@[<2>Result.error@ { session = %S;@ result = %a }@]"
        t.session_id
        Fmt.(option string) t.result
    else
      let tokens_info = match t.usage with
        | Some u ->
            let input = Usage.input_tokens u in
            let output = Usage.output_tokens u in
            let cached = Usage.cache_read_input_tokens u in
            (match input, output, cached with
            | Some i, Some o, Some c when c > 0 ->
                Printf.sprintf " (tokens: %d+%d, cached: %d)" i o c
            | Some i, Some o, _ ->
                Printf.sprintf " (tokens: %d+%d)" i o
            | _ -> "")
        | None -> ""
      in
      Fmt.pf fmt "@[<2>Result.%s@ { duration = %dms;@ cost = $%.4f%s }@]"
        t.subtype 
        t.duration_ms
        (Option.value t.total_cost_usd ~default:0.0)
        tokens_info
end

type t =
  | User of User.t
  | Assistant of Assistant.t
  | System of System.t
  | Result of Result.t

let user_string s = User (User.create_string s)
let user_blocks blocks = User (User.create_blocks blocks)
let user_with_tool_result ~tool_use_id ~content ?is_error () =
  User (User.create_with_tool_result ~tool_use_id ~content ?is_error ())

let assistant ~content ~model ?error () = Assistant (Assistant.create ~content ~model ?error ())
let assistant_text ~text ~model ?error () =
  Assistant (Assistant.create ~content:[Content_block.text text] ~model ?error ())

let system ~subtype ~data = System (System.create ~subtype ~data)
let system_init ~session_id =
  let data = System.Data.of_assoc [("session_id", `String session_id)] in
  System (System.create ~subtype:"init" ~data)
let system_error ~error =
  let data = System.Data.of_assoc [("error", `String error)] in
  System (System.create ~subtype:"error" ~data)

let result ~subtype ~duration_ms ~duration_api_ms ~is_error ~num_turns
    ~session_id ?total_cost_usd ?usage ?result ?structured_output () =
  Result (Result.create ~subtype ~duration_ms ~duration_api_ms ~is_error
    ~num_turns ~session_id ?total_cost_usd ?usage ?result ?structured_output ())

let to_json = function
  | User t -> User.to_json t
  | Assistant t -> Assistant.to_json t
  | System t -> System.to_json t
  | Result t -> Result.to_json t

let of_json json =
  match json with
  | `O fields -> (
      match List.assoc_opt "type" fields with
      | Some (`String "user") -> User (User.of_json json)
      | Some (`String "assistant") -> Assistant (Assistant.of_json json)
      | Some (`String "system") -> System (System.of_json json)
      | Some (`String "result") -> Result (Result.of_json json)
      | _ -> raise (Invalid_argument "Message.of_json: unknown type")
    )
  | _ -> raise (Invalid_argument "Message.of_json: expected object")

let pp fmt = function
  | User t -> User.pp fmt t
  | Assistant t -> Assistant.pp fmt t
  | System t -> System.pp fmt t
  | Result t -> Result.pp fmt t

let is_user = function User _ -> true | _ -> false
let is_assistant = function Assistant _ -> true | _ -> false
let is_system = function System _ -> true | _ -> false
let is_result = function Result _ -> true | _ -> false

let is_error = function
  | Result r -> Result.is_error r
  | System s -> System.subtype s = "error"
  | _ -> false

let extract_text = function
  | User u -> User.as_text u
  | Assistant a -> 
      let text = Assistant.combined_text a in
      if text = "" then None else Some text
  | _ -> None

let extract_tool_uses = function
  | Assistant a -> Assistant.get_tool_uses a
  | _ -> []

let get_session_id = function
  | System s when System.subtype s = "init" ->
      System.Data.get_string (System.data s) "session_id"
  | Result r -> Some (Result.session_id r)
  | _ -> None

let log_received t =
  Log.info (fun m -> m "← %a" pp t)

let log_sending t =
  Log.info (fun m -> m "→ %a" pp t)

let log_error msg t =
  Log.err (fun m -> m "%s: %a" msg pp t)