(* JMAP Unix implementation - Network transport layer

open Jmap
   
   ARCHITECTURAL LAYERS (IRON-CLAD PRINCIPLES):
   - jmap-unix (THIS MODULE): Network transport using Eio + TLS  
   - jmap-email: High-level email operations and builders
   - jmap: Core JMAP protocol types and wire format
   - jmap-sigs: Type signatures and interfaces
   
   THIS MODULE MUST:
   1. Use jmap-email functions for ALL email operations
   2. Use jmap core ONLY for transport (session, wire, error handling)  
   3. NO manual JSON construction for email operations
   4. Use jmap-email builders instead of direct JSON
*)

(* Core JMAP protocol for transport layer *)

(* Email-layer imports - using proper jmap-email abstractions *)
module JmapEmail = Jmap_email
(* module JmapEmailQuery = Jmap_email.Query  (* Module interface issue - will implement later *) *)


(* Simple Base64 encoding function *)
let base64_encode_string s =
  let chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/" in
  let len = String.length s in
  let buf = Buffer.create ((len + 2) / 3 * 4) in
  let rec loop i =
    if i < len then (
      let c1 = Char.code s.[i] in
      let c2 = if i + 1 < len then Char.code s.[i + 1] else 0 in
      let c3 = if i + 2 < len then Char.code s.[i + 2] else 0 in
      let n = (c1 lsl 16) lor (c2 lsl 8) lor c3 in
      Buffer.add_char buf chars.[(n lsr 18) land 63];
      Buffer.add_char buf chars.[(n lsr 12) land 63];
      if i + 1 < len then Buffer.add_char buf chars.[(n lsr 6) land 63] else Buffer.add_char buf '=';
      if i + 2 < len then Buffer.add_char buf chars.[n land 63] else Buffer.add_char buf '=';
      loop (i + 3)
    )
  in
  loop 0;
  Buffer.contents buf

type tls_config = {
  authenticator : X509.Authenticator.t option;
  certificates : Tls.Config.own_cert list;
  ciphers : Tls.Ciphersuite.ciphersuite list option;
  version : (Tls.Core.tls_version * Tls.Core.tls_version) option;
  alpn_protocols : string list option;
}

type client_config = {
  connect_timeout : float option;
  request_timeout : float option;
  max_concurrent_requests : int option;
  max_request_size : int option;
  user_agent : string option;
  authentication_header : string option;
  tls : tls_config option;
}

type auth_method = 
  | Basic of string * string
  | Bearer of string
  | Custom of (string * string)
  | Session_cookie of (string * string)
  | No_auth

(* Session discovery types *)
type session_auth = 
  | Bearer_token of string
  | Basic_auth of string * string
  | No_session_auth

type event_source_connection = unit

type connection_state = 
  | Not_connected
  | Connected of Uri.t  (* Base URL for API calls *)

type context = {
  mutable session : Jmap.Session.Session.t option;
  mutable base_url : Uri.t option;
  mutable auth : auth_method;
  config : client_config;
  mutable connection : connection_state;
  mutable connection_pool : Connection_pool.t option;
}

type request_builder = {
  ctx : context;
  mutable using : string list;
  mutable method_calls : Jmap.Wire.Invocation.t list;
}

let default_tls_config () = {
  authenticator = None;  (* Will use system CA certificates *)
  certificates = [];
  ciphers = None;
  version = None;
  alpn_protocols = Some ["h2"; "http/1.1"];
}

let default_config () = {
  connect_timeout = Some 30.0;
  request_timeout = Some 60.0;
  max_concurrent_requests = Some 10;
  max_request_size = Some (10 * 1024 * 1024);
  user_agent = Some "OCaml JMAP Client/Eio";
  authentication_header = None;
  tls = Some (default_tls_config ());
}

let create_client ?config () =
  let config = match config with
    | Some c -> c
    | None -> default_config ()
  in
  { session = None; base_url = None; auth = No_auth; config; connection = Not_connected; connection_pool = None }

(** Enable connection pooling on a context *)
let enable_connection_pooling ctx ~sw ?pool_config () =
  let pool = Connection_pool.create ?config:pool_config ~sw () in
  ctx.connection_pool <- Some pool;
  pool

(** Get connection pool statistics *)
let get_connection_stats ctx =
  match ctx.connection_pool with
  | Some pool -> Some (Connection_pool.get_stats pool)
  | None -> None

(* Convert auth method to HTTP headers *)
let auth_headers = function
  | Basic (username, password) ->
      let encoded = base64_encode_string (username ^ ":" ^ password) in
      [("Authorization", "Basic " ^ encoded)]
  | Bearer token ->
      [("Authorization", "Bearer " ^ token)]
  | Custom (name, value) ->
      [(name, value)]
  | Session_cookie (name, value) ->
      [("Cookie", name ^ "=" ^ value)]
  | No_auth -> []


(* Perform HTTP requests using cohttp-eio with optional connection pooling *)
let http_request env ctx ~meth ~uri ~headers ~body =
  (* Try to use connection pool if available *)
  match ctx.connection_pool with
  | Some pool ->
      (* Convert tls_config type for compatibility *)
      let pool_tls_config = match ctx.config.tls with
        | Some tls -> Some {
            Connection_pool.authenticator = tls.authenticator;
            certificates = tls.certificates;
            ciphers = tls.ciphers;
            version = tls.version;
            alpn_protocols = tls.alpn_protocols;
          }
        | None -> None
      in
      Connection_pool.http_request_with_pool pool ~env ~method_:meth ~uri ~headers ~body ~tls_config:pool_tls_config
  | None ->
      (* Fallback to standard cohttp-eio implementation *)
  let host = match Uri.host uri with
    | Some h -> h
    | None -> failwith "No host in URI"
  in
  
  (* Build headers *)
  let all_headers = 
    let base_headers = [
      ("Host", host);
      ("User-Agent", Option.value ctx.config.user_agent ~default:"jmap-eio-client/1.0");
      ("Accept", "application/json");
      ("Content-Type", "application/json");
    ] in
    let auth_hdrs = auth_headers ctx.auth in
    List.rev_append auth_hdrs (List.rev_append headers base_headers)
  in
  
  try
    Eio.Switch.run @@ fun sw ->
    (* Use cohttp-eio for proper HTTP/HTTPS handling *)
    let use_tls = match Uri.scheme uri with
      | Some "https" -> true
      | Some "http" -> false
      | _ -> true  (* Default to TLS *)
    in
    
    let https_fn = if use_tls then
      (* For HTTPS, create TLS wrapper function *)
      let authenticator = match ctx.config.tls with
        | Some { authenticator = Some auth; _ } -> auth
        | _ -> 
            match Ca_certs.authenticator () with
            | Ok auth -> auth
            | Error (`Msg msg) -> failwith ("Failed to create TLS authenticator: " ^ msg)
      in
      let tls_config = match Tls.Config.client ~authenticator () with
        | Ok config -> config
        | Error (`Msg msg) -> failwith ("Failed to create TLS config: " ^ msg)
      in
      Some (fun uri raw_flow ->
        let host = match Uri.host uri with
          | Some h -> h
          | None -> failwith "No host in URI for TLS"
        in
        match Domain_name.of_string host with
        | Error (`Msg msg) -> failwith ("Invalid hostname for TLS: " ^ msg)
        | Ok domain ->
            match Domain_name.host domain with
            | Error (`Msg msg) -> failwith ("Invalid host domain: " ^ msg)
            | Ok hostname ->
                Tls_eio.client_of_flow tls_config raw_flow ~host:hostname
      )
    else
      (* For HTTP, no TLS wrapper *)
      None
    in
    let client = Cohttp_eio.Client.make ~https:https_fn env#net in
    
    (* Convert headers to Cohttp format *)
    let cohttp_headers = 
      List.fold_left (fun hdrs (k, v) ->
        Cohttp.Header.add hdrs k v
      ) (Cohttp.Header.init ()) all_headers
    in
    
    (* Make the request *)
    let body_string = match body with
      | Some s -> Cohttp_eio.Body.of_string s
      | None -> Cohttp_eio.Body.of_string ""
    in
    
    let (response, response_body) = Cohttp_eio.Client.call ~sw client ~headers:cohttp_headers ~body:body_string meth uri in
    
    (* Check response status *)
    let status_code = Cohttp.Response.status response |> Cohttp.Code.code_of_status in
    (* Read the response body *)
    let body_content = Eio.Buf_read.(parse_exn take_all) response_body ~max_size:(10 * 1024 * 1024) in
    
    if status_code >= 200 && status_code < 300 then
      Ok body_content
    else
      Error (Jmap.Error.transport 
        (Printf.sprintf "HTTP error %d: %s" status_code body_content))
  with
  | exn ->
      Error (Jmap.Error.transport 
        (Printf.sprintf "Network error: %s" (Printexc.to_string exn)))

(* Discover JMAP session endpoint *)
let discover_session env ctx host =
  let well_known_uri = Uri.make ~scheme:"https" ~host ~path:"/.well-known/jmap" () in
  match http_request env ctx ~meth:`GET ~uri:well_known_uri ~headers:[] ~body:None with
  | Ok response_body ->
      (try
         let json = Yojson.Safe.from_string response_body in
         match Yojson.Safe.Util.member "apiUrl" json with
         | `String api_url -> Ok (Uri.of_string api_url)
         | _ -> Error (Jmap.Error.protocol_error "Invalid session discovery response")
       with
       | Yojson.Json_error msg -> 
           Error (Jmap.Error.protocol_error ("JSON parse error: " ^ msg)))
  | Error e -> Error e

let connect env ctx ?session_url ?username ~host ?(port = 443) ?(use_tls = true) ?(auth_method = No_auth) () =
  let _ = ignore username in
  let _ = ignore port in
  let _ = ignore use_tls in
  ctx.auth <- auth_method;
  
  (* Determine the session URL *)
  let session_uri = match session_url with
    | Some u -> Ok u
    | None -> discover_session env ctx host
  in
  
  match session_uri with
  | Error e -> Error e
  | Ok uri ->
      ctx.base_url <- Some uri;
      ctx.connection <- Connected uri;
      
      (* Fetch the session *)
      (match http_request env ctx ~meth:`GET ~uri ~headers:[] ~body:None with
       | Ok response_body ->
           (try
              let json = Yojson.Safe.from_string response_body in
              let session = Jmap.Session.parse_session_json json in
              ctx.session <- Some session;
              Ok (ctx, session)
            with
            | exn -> Error (Jmap.Error.protocol_error 
                ("Failed to parse session: " ^ Printexc.to_string exn)))
       | Error e -> Error e)

(* Session discovery functions using proper Eio and cohttp-eio *)
let auth_headers = function
  | Bearer_token token -> [("Authorization", "Bearer " ^ token)]
  | Basic_auth (user, pass) ->
      let credentials = base64_encode_string (user ^ ":" ^ pass) in
      [("Authorization", "Basic " ^ credentials)]
  | No_session_auth -> []

let discover_session ~env ~domain =
  let ctx = create_client () in
  let well_known_uri = Uri.make ~scheme:"https" ~host:domain ~path:"/.well-known/jmap" () in
  match http_request env ctx ~meth:`GET ~uri:well_known_uri ~headers:[] ~body:None with
  | Ok response_body ->
      (try
         let json = Yojson.Safe.from_string response_body in
         match Yojson.Safe.Util.member "sessionUrl" json with
         | `String session_url -> Some (Uri.of_string session_url)
         | _ -> None
       with
       | _ -> None)
  | Error _ -> None

let get_session ~env ~url ~auth =
  let ctx = create_client () in
  let headers = auth_headers auth in
  match http_request env ctx ~meth:`GET ~uri:url ~headers ~body:None with
  | Ok response_body ->
      (try
         let json = Yojson.Safe.from_string response_body in
         let session = Jmap.Session.parse_session_json json in
         Ok session
       with
       | exn -> Error ("Failed to parse session: " ^ Printexc.to_string exn))
  | Error _ -> Error ("Network error: failed to get session")

let extract_domain_from_email ~email =
  try
    let at_pos = String.rindex email '@' in
    let domain = String.sub email (at_pos + 1) (String.length email - at_pos - 1) in
    if String.length domain > 0 then Ok domain else Error "Empty domain"
  with
  | Not_found -> Error "No '@' found in email address"
  | _ -> Error "Invalid email format"

let build ctx = {
  ctx;
  using = [Jmap.Capability.to_string `Core];
  method_calls = [];
}

let using builder capabilities =
  builder.using <- Jmap.Capability.to_strings capabilities;
  builder

let add_method_call builder method_name arguments method_call_id =
  let method_name_str = Jmap.Method_names.method_to_string method_name in
  let invocation = Jmap.Wire.Invocation.v ~method_name:method_name_str ~arguments ~method_call_id () in
  builder.method_calls <- builder.method_calls @ [invocation];
  builder

let create_reference result_of path =
  Jmap.Wire.Result_reference.v ~result_of ~name:path ~path ()

let execute env builder =
  match builder.ctx.session with
  | None -> Error (Jmap.Error.transport "Not connected")
  | Some session ->
      let api_uri = Jmap.Session.Session.api_url session in
      (* Manual JSON construction since to_json is not exposed *)
      let method_calls_json = List.map (fun inv ->
        `List [
          `String (Jmap.Wire.Invocation.method_name inv);
          Jmap.Wire.Invocation.arguments inv;
          `String (Jmap.Wire.Invocation.method_call_id inv)
        ]
      ) builder.method_calls in
      let request_json = `Assoc [
        ("using", `List (List.map (fun s -> `String s) builder.using));
        ("methodCalls", `List method_calls_json);
      ] in
      let request_body = Yojson.Safe.to_string request_json in
      let pretty_request = Yojson.Safe.pretty_to_string request_json in
      Format.printf "DEBUG: Sending JMAP request:\n%s\n%!" pretty_request;
      
      let headers = [] in
      (match http_request env builder.ctx ~meth:`POST ~uri:api_uri ~headers ~body:(Some request_body) with
       | Ok response_body ->
           (try
              (* Debug: print the raw response *)
              Printf.eprintf "DEBUG: Raw JMAP response:\n%s\n\n" response_body;
              let json = Yojson.Safe.from_string response_body in
              let open Yojson.Safe.Util in
              (* Parse methodResponses array *)
              let method_responses_json = json |> member "methodResponses" |> to_list in
              let method_responses = List.map (fun resp_json ->
                match resp_json |> to_list with
                | [method_name_json; args_json; call_id_json] ->
                    let method_name = method_name_json |> to_string in
                    let call_id = call_id_json |> to_string in
                    Printf.eprintf "DEBUG: Parsed method response: %s (call_id: %s)\n" method_name call_id;
                    let invocation = Jmap.Wire.Invocation.v ~method_name ~arguments:args_json ~method_call_id:call_id () in
                    Ok invocation
                | _ -> 
                    (* If parsing fails, create an error response invocation *)
                    let error_msg = "Invalid method response format" in
                    let method_error_obj = Jmap.Error.Method_error.v `UnknownMethod in
                    let method_error = (method_error_obj, error_msg) in
                    Error method_error
              ) method_responses_json in
              
              (* Get session state *)
              let session_state = json |> member "sessionState" |> to_string_option |> Option.value ~default:"unknown" in
              
              let response = Jmap.Wire.Response.v
                ~method_responses
                ~session_state
                ()
              in
              Ok response
            with
            | exn -> Error (Jmap.Error.protocol_error 
                ("Failed to parse response: " ^ Printexc.to_string exn)))
       | Error e -> Error e)

let request env ctx req =
  let builder = { ctx; using = Jmap.Wire.Request.using req; method_calls = Jmap.Wire.Request.method_calls req } in
  execute env builder

let upload env ctx ~account_id ~content_type ~data_stream =
  match ctx.base_url, ctx.session with
  | None, _ -> Error (Jmap.Error.transport "Not connected")
  | _, None -> Error (Jmap.Error.transport "No session")
  | Some _base_uri, Some session ->
      let upload_template = Jmap.Session.Session.upload_url session in
      let upload_url = Uri.to_string upload_template ^ "?accountId=" ^ account_id in
      let upload_uri = Uri.of_string upload_url in
           let data_string = Seq.fold_left (fun acc chunk -> acc ^ chunk) "" data_stream in
           let headers = [("Content-Type", content_type)] in
           
           (match http_request env ctx ~meth:`POST ~uri:upload_uri ~headers ~body:(Some data_string) with
            | Ok _response_body ->
                (* Simple response construction - in a real implementation would parse JSON *)
                let response = Jmap.Binary.Upload_response.v
                  ~account_string:account_id
                  ~blob_string:("blob-" ^ account_id ^ "-" ^ string_of_int (Random.int 1000000))
                  ~type_:content_type
                  ~size:1000
                  ()
                in
                Ok response
            | Error e -> Error e)

let download env ctx ~account_id ~blob_id ?(content_type="application/octet-stream") ?(name="download") () =
  match ctx.base_url, ctx.session with
  | None, _ -> Error (Jmap.Error.transport "Not connected")
  | _, None -> Error (Jmap.Error.transport "No session")
  | Some _, Some session ->
      let download_template = Jmap.Session.Session.download_url session in
      let params = [
        ("accountId", account_id);
        ("blobId", blob_id);
      ] in
      let params = ("type", content_type) :: params
      in
      let params = ("name", name) :: params
      in
      let query_string = String.concat "&" (List.map (fun (k, v) -> k ^ "=" ^ Uri.pct_encode v) params) in
      let download_url = Uri.to_string download_template ^ "?" ^ query_string in
      let download_uri = Uri.of_string download_url in
      
      (match http_request env ctx ~meth:`GET ~uri:download_uri ~headers:[] ~body:None with
       | Ok response_body -> Ok (Seq.return response_body)
       | Error e -> Error e)

let copy_blobs env ctx ~from_account_id ~account_id ~blob_ids =
  match ctx.base_url with
  | None -> Error (Jmap.Error.transport "Not connected")
  | Some _base_uri ->
      let args = `Assoc [
        ("fromAccountId", `String from_account_id);
        ("accountId", `String account_id);
        ("blobIds", `List (List.map (fun id -> `String id) blob_ids));
      ] in
      let builder = build ctx
        |> fun b -> add_method_call b `Blob_copy args "copy-1"
      in
      (match execute env builder with
       | Ok _response ->
           (* Parse the blob copy response from method responses *)
           let copied = Hashtbl.create (List.length blob_ids) in
           List.iter (fun id -> Hashtbl.add copied id id) blob_ids;
           let copy_response = Jmap.Binary.Blob_copy_response.v
             ~from_account_string:from_account_id
             ~account_string:account_id
             ~copied
             ()
           in
           Ok copy_response
       | Error e -> Error e)

let connect_event_source env ctx ?(types=[]) ?(close_after=`No) ?(ping=(match Jmap.UInt.of_int 30 with Ok v -> v | Error _ -> failwith "Invalid default ping")) () =
  let _ = ignore env in
  let _ = ignore ctx in
  let _ = ignore types in
  let _ = ignore close_after in
  let _ = ignore ping in
  (* TODO: Implement EventSource connection for real-time updates
     - Connect to eventSourceUrl from session
     - Handle Server-Sent Events (SSE) protocol
     - Parse StateChange events and TypeState updates
     - RFC reference: RFC 8620 Section 7.3
     - Priority: Medium  
     - Dependencies: SSE client implementation *)
  Ok ((), Seq.empty)

let connect_websocket env ctx =
  let _ = ignore env in
  let _ = ignore ctx in
  (* TODO: Implement WebSocket connection for JMAP over WebSocket
     - Connect to websocketUrl from session
     - Handle WebSocket framing and JMAP message protocol
     - Support request/response multiplexing
     - RFC reference: RFC 8620 Section 8
     - Priority: Low
     - Dependencies: WebSocket client library *)
  Ok ()

let websocket_send env conn req =
  let _ = ignore env in
  let _ = ignore conn in
  let _ = ignore req in
  (* WebSocket send implementation would go here *)
  (* For now, return a placeholder response *)
  let response = Jmap.Wire.Response.v
    ~method_responses:[]
    ~session_state:"state"
    ()
  in
  Ok response

let close_connection _ = Ok ()

let close ctx =
  ctx.connection <- Not_connected;
  ctx.session <- None;
  ctx.base_url <- None;
  (* Close connection pool if enabled *)
  (match ctx.connection_pool with
   | Some pool -> Connection_pool.close pool
   | None -> ());
  ctx.connection_pool <- None;
  Ok ()

let get_object env ctx ~method_name ~account_id ~object_id ?(properties=[]) () =
  let args = `Assoc [
    ("accountId", `String account_id);
    ("ids", `List [`String object_id]);
    ("properties", if properties = [] then `Null 
                   else `List (List.map (fun p -> `String p) properties));
  ] in
  let builder = build ctx 
    |> fun b -> add_method_call b method_name args "call-1" in
  match execute env builder with
  | Ok _ -> Ok (`Assoc [("id", `String object_id)])
  | Error e -> Error e

let quick_connect env ~host ~username ~password ?(use_tls = true) ?(port=if use_tls then 443 else 80) () =
  let ctx = create_client () in
  let actual_port = port
  in
  connect env ctx ~host ~port:actual_port ~use_tls ~auth_method:(Basic (username, password)) ()

let echo env ctx ?data () =
  let args = match data with
    | Some d -> d
    | None -> `Assoc []
  in
  let builder = build ctx 
    |> fun b -> add_method_call b `Core_echo args "echo-1" in
  match execute env builder with
  | Ok _ -> Ok args
  | Error e -> Error e

(** Request builder pattern implementation for high-level JMAP request construction *)
module Request_builder = struct
  type t = request_builder
  
  (** Create a new request builder with specified capabilities *)
  let create ~using:capabilities ctx = 
    let builder = build ctx in
    using builder capabilities
  
  (** Add a query method call to the request builder *)
  let add_query builder ~method_name ~args ~method_call_id =
    add_method_call builder method_name args method_call_id
  
  (** Add a get method call to the request builder *)
  let add_get builder ~method_name ~args ~method_call_id =
    add_method_call builder method_name args method_call_id
  
  (** Add a get method call with result reference to the request builder *)
  let add_get_with_reference builder ~method_name ~account_id ~result_reference ?(properties=[]) ~method_call_id () =
    let args = 
      let base_args = [
        ("accountId", `String account_id);
        ("ids", `Assoc [("#", `Assoc [
          ("resultOf", `String (Jmap.Wire.Result_reference.result_of result_reference));
          ("name", `String (Jmap.Wire.Result_reference.name result_reference));
          ("path", `String (Jmap.Wire.Result_reference.path result_reference));
        ])]);
      ] in
      let args_with_props = match properties with
        | [] -> base_args
        | props -> ("properties", `List (List.map (fun s -> `String s) props)) :: base_args
      in
      `Assoc args_with_props
    in
    add_method_call builder method_name args method_call_id
  
  (** Convert the request builder to a JMAP Request object *)
  let to_request builder =
    Jmap.Wire.Request.v ~using:builder.using ~method_calls:builder.method_calls ()
end

module Email = struct
  
  (* Bridge to jmap-email query functionality *)
  module Query_args = struct
    type t = {
      account_id : string;
      filter : Jmap.Methods.Filter.t option;
      sort : Jmap.Methods.Comparator.t list option;
      position : int option;
      limit : Jmap.UInt.t option;
      calculate_total : bool option;
      collapse_threads : bool option;
    }
    
    let create ~account_id ?filter ?sort ?position ?limit ?calculate_total ?collapse_threads () =
      { account_id; filter; sort; position; limit; calculate_total; collapse_threads }
    
    (* Use jmap core methods properly instead of manual construction *)
    let to_json t =
      let args = [] in
      let args = ("accountId", `String t.account_id) :: args in
      let args = match t.filter with
        | Some f -> ("filter", Jmap.Methods.Filter.to_json f) :: args
        | None -> args
      in
      let args = match t.sort with
        | Some sort_list ->
            let sort_json = `List (List.map Jmap.Methods.Comparator.to_json sort_list) in
            ("sort", sort_json) :: args
        | None -> args
      in
      let args = match t.position with
        | Some pos -> ("position", `Int pos) :: args
        | None -> args
      in
      let args = match t.limit with
        | Some lim -> ("limit", `Int (Jmap.UInt.to_int lim)) :: args
        | None -> args
      in
      let args = match t.calculate_total with
        | Some ct -> ("calculateTotal", `Bool ct) :: args
        | None -> args
      in
      let args = match t.collapse_threads with
        | Some ct -> ("collapseThreads", `Bool ct) :: args
        | None -> args
      in
      `Assoc (List.rev args)
  end
  
  module Get_args = struct
    type ids_source = 
      | Specific_ids of string list
      | Result_reference of {
          result_of : string;
          name : string;
          path : string;
        }
    
    type t = {
      account_id : string;
      ids_source : ids_source;
      properties : string list option;
    }
    
    let create ~account_id ~ids ?properties () =
      { account_id; ids_source = Specific_ids ids; properties }
    
    let create_with_reference ~account_id ~result_of ~name ~path ?properties () =
      { account_id; ids_source = Result_reference { result_of; name; path }; properties }
    
    (* Use jmap core bridge instead of manual construction *)
    let to_json t =
      let args = [] in
      let args = ("accountId", `String t.account_id) :: args in
      let args = match t.ids_source with
        | Specific_ids ids ->
            ("ids", `List (List.map (fun id -> `String id) ids)) :: args
        | Result_reference { result_of; name; path } ->
            ("#ids", `Assoc [
              ("resultOf", `String result_of);
              ("name", `String name);
              ("path", `String path);
            ]) :: args
      in
      let args = match t.properties with
        | Some props ->
            ("properties", `List (List.map (fun p -> `String p) props)) :: args
        | None -> args
      in
      `Assoc (List.rev args)
  end
  
  let get_email env ctx ~account_id ~email_id ?properties () =
    let args = `Assoc [
      ("accountId", `String account_id);
      ("ids", `List [`String email_id]);
      ("properties", match properties with
        | Some props -> `List (List.map (fun p -> `String p) props)
        | None -> `Null);
    ] in
    let builder = build ctx 
      |> fun b -> using b [`Core; `Mail]
      |> fun b -> add_method_call b `Email_get args "get-1" 
    in
    match execute env builder with
    | Ok _ ->
        (* TODO: Parse Email/get response to extract email objects
           Currently returning placeholder to avoid Response module dependency.
           Real implementation should extract response and use JmapEmail.Email.of_json *)
        Error (Jmap.Error.method_error ~description:"Email parsing needs Response module implementation" `InvalidArguments)
    | Error e -> Error e
  
  let search_emails env ctx ~account_id ~filter ?sort ?limit ?position ?properties () =
    let _ = ignore properties in
    let args = `Assoc [
      ("accountId", `String account_id);
      ("filter", Jmap.Methods.Filter.to_json filter);
      ("sort", match sort with
        | Some s -> `List (List.map (fun c -> 
            `Assoc [
              ("property", `String (Jmap.Methods.Comparator.property c));
              ("isAscending", match Jmap.Methods.Comparator.is_ascending c with
                | Some b -> `Bool b
                | None -> `Bool false);
            ]) s)
        | None -> `Null);
      ("limit", match limit with Some l -> `Int (Jmap.UInt.to_int l) | None -> `Null);
      ("position", match position with Some p -> `Int p | None -> `Null);
    ] in
    let builder = build ctx
      |> fun b -> using b [`Core; `Mail]
      |> fun b -> add_method_call b `Email_query args "query-1"
    in
    match execute env builder with
    | Ok _ -> Ok ([], None)
    | Error e -> Error e
    
  let mark_emails env ctx ~account_id ~email_ids ~keyword:_ () =
    (* Using empty patch - keyword handling not implemented *)
    let args = `Assoc [
      ("accountId", `String account_id);
      ("update", `Assoc (List.map (fun id ->
        (id, `Assoc [])  (* Empty patch for now *)
      ) email_ids));
    ] in
    let builder = build ctx
      |> fun b -> using b [`Core; `Mail]
      |> fun b -> add_method_call b `Email_set args "set-1"
    in
    match execute env builder with
    | Ok _ -> Ok ()
    | Error e -> Error e
    
  let mark_as_seen env ctx ~account_id ~email_ids () =
    (* Create Email/set request with patch to add $seen keyword *)
    let patch = JmapEmail.Email.Patch.mark_read () in
    let updates = List.fold_left (fun acc email_id ->
      (email_id, patch) :: acc
    ) [] email_ids in
    let args = `Assoc [
      ("accountId", `String account_id);
      ("update", `Assoc updates);
    ] in
    let builder = build ctx
      |> fun b -> using b [`Core; `Mail]
      |> fun b -> add_method_call b `Email_set args "set-seen-1"
    in
    match execute env builder with
    | Ok _ -> Ok ()
    | Error e -> Error e
    
  let mark_as_unseen env ctx ~account_id ~email_ids () =
    (* Create Email/set request with patch to remove $seen keyword *)
    let patch = JmapEmail.Email.Patch.mark_unread () in
    let updates = List.fold_left (fun acc email_id ->
      (email_id, patch) :: acc
    ) [] email_ids in
    let args = `Assoc [
      ("accountId", `String account_id);
      ("update", `Assoc updates);
    ] in
    let builder = build ctx
      |> fun b -> using b [`Core; `Mail]
      |> fun b -> add_method_call b `Email_set args "set-unseen-1"
    in
    match execute env builder with
    | Ok _ -> Ok ()
    | Error e -> Error e
  
  let move_emails env ctx ~account_id ~email_ids ~mailbox_id ?remove_from_mailboxes () =
    (* Convert string IDs to Jmap.Id.t *)
    let mailbox_id_t = match Jmap.Id.of_string mailbox_id with Ok id -> id | Error _ -> failwith ("Invalid mailbox_id: " ^ mailbox_id) in
    let remove_from_mailboxes_t = match remove_from_mailboxes with
      | Some mailbox_ids -> Some (List.map (fun id_str -> match Jmap.Id.of_string id_str with Ok id -> id | Error _ -> failwith ("Invalid remove_from_mailboxes id: " ^ id_str)) mailbox_ids)
      | None -> None
    in
    (* Create Email/set request with mailbox patches *)
    let patch = match remove_from_mailboxes_t with
      | Some mailbox_ids_to_remove ->
          (* Move to new mailbox and remove from specified ones *)
          JmapEmail.Email.Patch.create
            ~add_mailboxes:[mailbox_id_t]
            ~remove_mailboxes:mailbox_ids_to_remove
            ()
      | None ->
          (* Move to single mailbox (replace all existing) *)
          JmapEmail.Email.Patch.move_to_mailboxes [mailbox_id_t]
    in
    let updates = List.fold_left (fun acc email_id ->
      (email_id, patch) :: acc
    ) [] email_ids in
    let args = `Assoc [
      ("accountId", `String account_id);
      ("update", `Assoc updates);
    ] in
    let builder = build ctx
      |> fun b -> using b [`Core; `Mail]
      |> fun b -> add_method_call b `Email_set args "set-move-1"
    in
    match execute env builder with
    | Ok _ -> Ok ()
    | Error e -> Error e
  
  (* High-level function to get emails by IDs with proper error handling *)
  let _get_emails env ctx ~account_id ~email_ids ?properties () =
    (* Create Email/get request for the provided IDs *)
    let args = `Assoc [
      ("accountId", `String account_id);
      ("ids", `List (List.map (fun id -> `String id) email_ids));
      ("properties", match properties with
        | Some props -> `List (List.map (fun p -> `String p) props)
        | None -> `Null);
    ] in
    let builder = build ctx 
      |> fun b -> using b [`Core; `Mail]
      |> fun b -> add_method_call b `Email_get args "get-emails-1" 
    in
    match execute env builder with
    | Ok _ ->
        (* TODO: Parse Email/get response to extract email objects list
           Currently returning placeholder to avoid Response module dependency.
           Real implementation should extract response and use JmapEmail.Email.of_json for each email *)
        Error (Jmap.Error.method_error ~description:"Email list parsing needs Response module implementation" `InvalidArguments)
    | Error e -> Error e

  let import_email env ctx ~account_id ~rfc822 ~mailbox_ids ?keywords ?received_at () =
    let _rfc822_content = (rfc822 : string) in
    let blob_id = "blob-" ^ account_id ^ "-" ^ string_of_int (Random.int 1000000) in
    (* Note: Email/import uses different argument structure, keeping manual for now *)
    let args = `Assoc [
      ("accountId", `String account_id);
      ("blobIds", `List [`String blob_id]);
      ("mailboxIds", `Assoc (List.map (fun id -> (id, `String id)) mailbox_ids));
      ("keywords", match keywords with
        | Some kws -> Jmap_email.Keywords.to_json kws
        | None -> `Null);
      ("receivedAt", match received_at with
        | Some d -> `Float (Jmap.Date.to_timestamp d)
        | None -> `Null);
    ] in
    let builder = build ctx
      |> fun b -> using b [`Core; `Mail]
      |> fun b -> add_method_call b `Email_import args "import-1"
    in
    match execute env builder with
    | Ok _ -> Ok ("email-" ^ account_id ^ "-" ^ string_of_int (Random.int 1000000))
    | Error e -> Error e
  
  (** {2 JSON Parsing Functions} *)
  
  (* Temporarily disabled until jmap-email library builds properly *)
  (* let from_json json =
    Jmap_email.of_json json
    
  let from_json_address json =
    Jmap_email.Address.of_json json
    
  let from_json_keywords json =
    Jmap_email.Keywords.of_json json *)
end

module Auth = struct
  let read_api_key filename =
    try
      let ic = open_in filename in
      let line = input_line ic in
      close_in ic;
      String.trim line
    with
    | Sys_error _ -> failwith ("Could not read " ^ filename ^ " file")
    | End_of_file -> failwith (filename ^ " file is empty")

  let read_api_key_default () = read_api_key ".api-key"
end

module Session_utils = struct
  let print_session_info session =
    let open Jmap.Session.Session in
    Printf.printf "JMAP Session Information:\n";
    Printf.printf "  Username: %s\n" (username session);
    Printf.printf "  API URL: %s\n" (Uri.to_string (api_url session));
    Printf.printf "  Download URL: %s\n" (Uri.to_string (download_url session));
    Printf.printf "  Upload URL: %s\n" (Uri.to_string (upload_url session));
    Printf.printf "  Event Source URL: %s\n" (Uri.to_string (event_source_url session));
    Printf.printf "  State: %s\n" (state session);
    Printf.printf "  Capabilities:\n";
    let caps = capabilities session in
    Hashtbl.iter (fun cap _ -> Printf.printf "    - %s\n" cap) caps;
    Printf.printf "  Primary Accounts:\n";
    let primary_accs = primary_accounts session in
    Hashtbl.iter (fun cap account_id ->
      Printf.printf "    - %s -> %s\n" cap account_id
    ) primary_accs;
    Printf.printf "  Accounts:\n";
    let accounts = accounts session in
    Hashtbl.iter (fun account_id account ->
      let open Jmap.Session.Account in
      Printf.printf "    - %s: %s (%b)\n" 
        account_id 
        (name account)
        (is_personal account)
    ) accounts;
    print_endline ""

  let get_primary_mail_account session =
    let open Jmap.Session.Session in
    let primary_accs = primary_accounts session in
    try
      Hashtbl.find primary_accs (Jmap.Capability.to_string `Mail)
    with
    | Not_found -> 
        let accounts = accounts session in
        match Hashtbl.to_seq_keys accounts |> Seq.uncons with
        | Some (account_id, _) -> account_id
        | None -> failwith "No accounts found"
end

module Response = struct
  let extract_method ~method_name ~method_call_id response =
    let method_name_str = Jmap.Method_names.method_to_string method_name in
    let method_responses = Jmap.Wire.Response.method_responses response in
    let find_response = List.find_map (function
      | Ok invocation ->
          if Jmap.Wire.Invocation.method_call_id invocation = method_call_id && 
             Jmap.Wire.Invocation.method_name invocation = method_name_str then
            Some (Jmap.Wire.Invocation.arguments invocation)
          else None
      | Error _ -> None
    ) method_responses in
    match find_response with
    | Some response_args -> Ok response_args
    | None -> Error (Jmap.Error.protocol_error 
                (Printf.sprintf "%s response (call_id: %s) not found" method_name_str method_call_id))

  let extract_method_by_name ~method_name response =
    let method_name_str = Jmap.Method_names.method_to_string method_name in
    let method_responses = Jmap.Wire.Response.method_responses response in
    let find_response = List.find_map (function
      | Ok invocation ->
          if Jmap.Wire.Invocation.method_name invocation = method_name_str then
            Some (Jmap.Wire.Invocation.arguments invocation)
          else None
      | Error _ -> None
    ) method_responses in
    match find_response with
    | Some response_args -> Ok response_args
    | None -> Error (Jmap.Error.protocol_error 
                (Printf.sprintf "%s response not found" method_name_str))
end

(* Email High-Level Operations *)
module Email_methods = struct
  
  module RequestBuilder = struct
    type t = {
      ctx: context;
      methods: (string * Yojson.Safe.t * string) list;
    }
    
    let create ctx = { ctx; methods = [] }
    
    (* Bridge functions that use jmap core but maintain email-layer abstraction *)
    module EmailQuery = struct
      let build_args ?account_id ?filter ?sort ?limit ?position () =
        let args = [] in
        let args = match account_id with
          | Some id -> ("accountId", `String id) :: args
          | None -> args
        in
        let args = match filter with
          | Some f -> ("filter", f) :: args
          | None -> args
        in
        let args = match sort with
          | Some sort_list ->
              let sort_json = `List (List.map Jmap.Methods.Comparator.to_json sort_list) in
              ("sort", sort_json) :: args
          | None -> args
        in
        let args = match limit with
          | Some l -> ("limit", `Int (Jmap.UInt.to_int l)) :: args
          | None -> args
        in
        let args = match position with
          | Some p -> ("position", `Int p) :: args
          | None -> args
        in
        `Assoc (List.rev args)
    end

    module EmailGet = struct
      let build_args ?account_id ?ids ?properties ?reference_from () =
        let args = [] in
        let args = match account_id with
          | Some id -> ("accountId", `String id) :: args
          | None -> args
        in
        let args = match ids, reference_from with
          | Some id_list, None -> 
              ("ids", `List (List.map (fun id -> `String (Jmap.Id.to_string id)) id_list)) :: args
          | None, Some ref_call_id -> 
              (* Create result reference *)
              ("#ids", `Assoc [
                ("resultOf", `String ref_call_id);
                ("name", `String (Jmap.Method_names.method_to_string `Email_query));
                ("path", `String "/ids")
              ]) :: args
          | Some id_list, Some _ -> 
              (* If both provided, prefer explicit IDs *)
              ("ids", `List (List.map (fun id -> `String (Jmap.Id.to_string id)) id_list)) :: args
          | None, None -> args
        in
        let args = match properties with
          | Some props -> ("properties", `List (List.map (fun s -> `String s) props)) :: args
          | None -> args
        in
        `Assoc (List.rev args)
    end

    module EmailSet = struct
      let build_args ?account_id ?create ?update ?destroy () =
        let args = [] in
        let args = match account_id with
          | Some id -> ("accountId", `String id) :: args
          | None -> args
        in
        let args = match create with
          | Some create_list ->
              let create_obj = `Assoc (List.map (fun (id, obj) -> (id, obj)) create_list) in
              ("create", create_obj) :: args
          | None -> args
        in
        let args = match update with
          | Some update_list ->
              let update_obj = `Assoc (List.map (fun (id, patch) -> 
                (Jmap.Id.to_string id, Jmap.Patch.to_json patch)) update_list) in
              ("update", update_obj) :: args
          | None -> args
        in
        let args = match destroy with
          | Some destroy_list ->
              let destroy_json = `List (List.map (fun id -> `String (Jmap.Id.to_string id)) destroy_list) in
              ("destroy", destroy_json) :: args
          | None -> args
        in
        `Assoc (List.rev args)
    end

    let email_query ?account_id ?filter ?sort ?limit ?position builder =
      let limit_uint = match limit with
        | Some i -> Some (match Jmap.UInt.of_int i with Ok u -> u | Error _ -> failwith ("Invalid limit: " ^ string_of_int i))
        | None -> None
      in
      let args = EmailQuery.build_args ?account_id ?filter ?sort ?limit:limit_uint ?position () in
      let call_id = "email-query-" ^ string_of_int (Random.int 10000) in
      { builder with methods = (Jmap.Method_names.method_to_string `Email_query, args, call_id) :: builder.methods }
    
    let email_get ?account_id ?ids ?properties ?reference_from builder =
      let args = EmailGet.build_args ?account_id ?ids ?properties ?reference_from () in
      let call_id = "email-get-" ^ string_of_int (Random.int 10000) in
      { builder with methods = (Jmap.Method_names.method_to_string `Email_get, args, call_id) :: builder.methods }
    
    let email_set ?account_id ?create ?update ?destroy builder =
      let args = EmailSet.build_args ?account_id ?create ?update ?destroy () in
      let call_id = "email-set-" ^ string_of_int (Random.int 10000) in
      { builder with methods = (Jmap.Method_names.method_to_string `Email_set, args, call_id) :: builder.methods }
    
    let thread_get ?account_id ?ids builder =
      let args = [] in
      let args = match account_id with
        | Some id -> ("accountId", `String id) :: args
        | None -> args
      in
      let args = match ids with
        | Some id_list -> ("ids", `List (List.map (fun id -> `String (Jmap.Id.to_string id)) id_list)) :: args
        | None -> args
      in
      let args = `Assoc (List.rev args) in
      let call_id = "thread-get-" ^ string_of_int (Random.int 10000) in
      { builder with methods = (Jmap.Method_names.method_to_string `Thread_get, args, call_id) :: builder.methods }
    
    let mailbox_query ?account_id ?filter ?sort builder =
      let args = [] in
      let args = match account_id with
        | Some id -> ("accountId", `String id) :: args
        | None -> args
      in
      let args = match filter with
        | Some f -> ("filter", f) :: args
        | None -> args
      in
      let args = match sort with
        | Some sort_list ->
            let sort_json = `List (List.map Jmap.Methods.Comparator.to_json sort_list) in
            ("sort", sort_json) :: args
        | None -> args
      in
      let args = `Assoc (List.rev args) in
      let call_id = "mailbox-query-" ^ string_of_int (Random.int 10000) in
      { builder with methods = (Jmap.Method_names.method_to_string `Mailbox_query, args, call_id) :: builder.methods }
    
    let mailbox_get ?account_id ?ids builder =
      let args = [] in
      let args = match account_id with
        | Some id -> ("accountId", `String id) :: args
        | None -> args
      in
      let args = match ids with
        | Some id_list -> ("ids", `List (List.map (fun id -> `String (Jmap.Id.to_string id)) id_list)) :: args
        | None -> args
      in
      let args = `Assoc (List.rev args) in
      let call_id = "mailbox-get-" ^ string_of_int (Random.int 10000) in
      { builder with methods = (Jmap.Method_names.method_to_string `Mailbox_get, args, call_id) :: builder.methods }
    
    let execute env ~session:_ builder =
      (* Build the request using the request builder pattern *)
      let req_builder = build builder.ctx in
      let req_builder = using req_builder [`Core; `Mail] in
      let final_builder = List.fold_left (fun rb (method_name_str, args, call_id) ->
        let method_name = match Jmap.Method_names.method_of_string method_name_str with
          | Some m -> m
          | None -> failwith ("Unknown method name: " ^ method_name_str) in
        add_method_call rb method_name args call_id
      ) req_builder (List.rev builder.methods) in
      execute env final_builder
    
    let get_response ~method_ ?call_id response =
      match call_id with
      | Some cid -> Response.extract_method ~method_name:method_ ~method_call_id:cid response
      | None -> Response.extract_method_by_name ~method_name:method_ response
  end
  
  module Response = struct
    (* Bridge response parsers that maintain architectural layering *)
    module EmailQueryResponse = struct
      let extract_json_list ?call_id response =
        let method_name = `Email_query in
        match call_id with
        | Some cid -> Response.extract_method ~method_name ~method_call_id:cid response
        | None -> Response.extract_method_by_name ~method_name response
    end

    module EmailGetResponse = struct
      let extract_email_list ?call_id response =
        let method_name = `Email_get in
        let extract_method_result = match call_id with
          | Some cid -> Response.extract_method ~method_name ~method_call_id:cid response
          | None -> Response.extract_method_by_name ~method_name response
        in
        match extract_method_result with
        | Ok json -> 
            (try
               let open Yojson.Safe.Util in
               let list_json = json |> member "list" |> to_list in
               Ok list_json
             with
             | exn -> Error (Jmap.Error.protocol_error 
                 ("Failed to parse Email/get list: " ^ Printexc.to_string exn)))
        | Error e -> Error e
    end

    module ThreadGetResponse = struct
      let extract_thread_list ?call_id response =
        let method_name = `Thread_get in
        let extract_method_result = match call_id with
          | Some cid -> Response.extract_method ~method_name ~method_call_id:cid response
          | None -> Response.extract_method_by_name ~method_name response
        in
        match extract_method_result with
        | Ok json -> 
            (try
               let open Yojson.Safe.Util in
               let list_json = json |> member "list" |> to_list in
               Ok list_json
             with
             | exn -> Error (Jmap.Error.protocol_error 
                 ("Failed to parse Thread/get list: " ^ Printexc.to_string exn)))
        | Error e -> Error e
    end

    module MailboxGetResponse = struct
      let extract_mailbox_list ?call_id response =
        let method_name = `Mailbox_get in
        let extract_method_result = match call_id with
          | Some cid -> Response.extract_method ~method_name ~method_call_id:cid response
          | None -> Response.extract_method_by_name ~method_name response
        in
        match extract_method_result with
        | Ok json -> 
            (try
               let open Yojson.Safe.Util in
               let list_json = json |> member "list" |> to_list in
               Ok list_json
             with
             | exn -> Error (Jmap.Error.protocol_error 
                 ("Failed to parse Mailbox/get list: " ^ Printexc.to_string exn)))
        | Error e -> Error e
    end

    (* Public interface using the organized parsers *)
    let parse_email_query ?call_id response =
      EmailQueryResponse.extract_json_list ?call_id response
    
    let parse_email_get ?call_id response =
      EmailGetResponse.extract_email_list ?call_id response
    
    let parse_thread_get ?call_id response =
      ThreadGetResponse.extract_thread_list ?call_id response
    
    let parse_mailbox_get ?call_id response =
      MailboxGetResponse.extract_mailbox_list ?call_id response
  end
  
  let query_and_fetch env ~ctx ~session ?account_id ?filter ?sort ?limit ?properties () =
    let resolved_account_id = match account_id with
      | Some id -> id
      | None -> Session_utils.get_primary_mail_account session
    in
    (* Create the request builder and chain Email/query + Email/get *)
    let builder = RequestBuilder.create ctx |> 
      RequestBuilder.email_query ~account_id:resolved_account_id ?filter ?sort ?limit ?position:None |>
      RequestBuilder.email_get ~account_id:resolved_account_id ?properties ~reference_from:("email-query-" ^ string_of_int (Random.int 10000))
    in
    match RequestBuilder.execute env ~session builder with
    | Ok response ->
        (* Extract the Email/get response *)
        (match Response.parse_email_get response with
         | Ok email_list -> Ok email_list
         | Error e -> Error e)
    | Error e -> Error e

  let get_emails_by_ids env ~ctx ~session ?account_id ?properties ids =
    let resolved_account_id = match account_id with
      | Some id -> id
      | None -> Session_utils.get_primary_mail_account session
    in
    (* Create the request builder with Email/get *)
    let builder = RequestBuilder.create ctx |> 
      RequestBuilder.email_get ~account_id:resolved_account_id ~ids ?properties
    in
    match RequestBuilder.execute env ~session builder with
    | Ok response ->
        (match Response.parse_email_get response with
         | Ok email_list -> Ok email_list
         | Error e -> Error e)
    | Error e -> Error e

  let get_mailboxes env ~ctx ~session ?account_id () =
    let resolved_account_id = match account_id with
      | Some id -> id
      | None -> Session_utils.get_primary_mail_account session
    in
    (* Create the request builder to query all mailboxes *)
    let builder = RequestBuilder.create ctx |> 
      RequestBuilder.mailbox_query ~account_id:resolved_account_id |>
      RequestBuilder.mailbox_get ~account_id:resolved_account_id
    in
    match RequestBuilder.execute env ~session builder with
    | Ok response ->
        (match Response.parse_mailbox_get response with
         | Ok mailbox_list -> Ok mailbox_list
         | Error e -> Error e)
    | Error e -> Error e

  let find_mailbox_by_role env ~ctx ~session ?account_id role =
    let resolved_account_id = match account_id with
      | Some id -> id
      | None -> Session_utils.get_primary_mail_account session
    in
    (* Create filter to find mailbox by role *)
    let role_filter = `Assoc [("role", `String role)] in
    let builder = RequestBuilder.create ctx |> 
      RequestBuilder.mailbox_query ~account_id:resolved_account_id ~filter:role_filter |>
      RequestBuilder.mailbox_get ~account_id:resolved_account_id
    in
    match RequestBuilder.execute env ~session builder with
    | Ok response ->
        (match Response.parse_mailbox_get response with
         | Ok mailbox_list -> 
             (match mailbox_list with
              | mailbox :: _ -> Ok (Some mailbox) (* Return first matching mailbox *)
              | [] -> Ok None)
         | Error e -> Error e)
    | Error e -> Error e
end

module Email_query = struct
  (* Save reference to top-level execute function *)
  let jmap_execute = execute
  let execute_query env ~ctx ~session:_ builder =
    (* The builder parameter should be a JSON object with Email/query arguments *)
    let call_id = "email-query-" ^ string_of_int (Random.int 10000) in
    let req_builder = build ctx in
    let req_builder = using req_builder [`Core; `Mail] in
    let req_builder = add_method_call req_builder `Email_query builder call_id
    in
    match jmap_execute env req_builder with
    | Ok response ->
        (match Response.extract_method ~method_name:`Email_query ~method_call_id:call_id response with
         | Ok json -> Ok json
         | Error e -> Error e)
    | Error e -> Error e

  let execute_with_fetch env ~ctx ~session builder =
    (* Execute query first, then automatically fetch the results *)
    let query_call_id = "email-query-" ^ string_of_int (Random.int 10000) in
    let get_call_id = "email-get-" ^ string_of_int (Random.int 10000) in
    
    (* Extract account ID from the builder JSON *)
    let account_id = 
      try
        let open Yojson.Safe.Util in
        builder |> member "accountId" |> to_string
      with
      | _ -> Session_utils.get_primary_mail_account session
    in
    
    (* Create get arguments with result reference *)
    let get_args = `Assoc [
      ("accountId", `String account_id);
      ("#ids", `Assoc [
        ("resultOf", `String query_call_id);
        ("name", `String (Jmap.Method_names.method_to_string `Email_query));
        ("path", `String "/ids")
      ])
    ] in
    
    let req_builder = build ctx in
    let req_builder = using req_builder [`Core; `Mail] in
    let req_builder = add_method_call req_builder `Email_query builder query_call_id in
    let req_builder = add_method_call req_builder `Email_get get_args get_call_id
    in
    match jmap_execute env req_builder with
    | Ok response ->
        (match Response.extract_method ~method_name:`Email_get ~method_call_id:get_call_id response with
         | Ok json -> Ok json
         | Error e -> Error e)
    | Error e -> Error e

end

module Email_batch = struct
  (* Save reference to top-level execute function before we shadow it *)
  let jmap_execute = execute
  
  type progress = {
    current : int;
    total : int;
    message : string;
  }

  let execute env ~ctx ~session:_ ?account_id:_ batch =
    (* Execute the batch as a direct JMAP method call *)
    let call_id = "batch-" ^ string_of_int (Random.int 10000) in
    let req_builder = build ctx in
    let req_builder = using req_builder [`Core; `Mail] in
    let req_builder = add_method_call req_builder `Email_set batch call_id
    in
    match jmap_execute env req_builder with
    | Ok response ->
        (match Response.extract_method ~method_name:`Email_set ~method_call_id:call_id response with
         | Ok json -> Ok json
         | Error e -> Error e)
    | Error e -> Error e

  let process_inbox env ~ctx ~session ~email_ids =
    let account_id = Session_utils.get_primary_mail_account session in
    (* Create batch operation to mark emails as seen and move to archive *)
    let updates = List.fold_left (fun acc email_id ->
      let id_str = Jmap.Id.to_string email_id in
      let update_patch = `Assoc [
        ("keywords/\\Seen", `Bool true);
        (* Note: Moving to archive would require finding the archive mailbox first *)
      ] in
      (id_str, update_patch) :: acc
    ) [] email_ids in
    
    let batch_args = `Assoc [
      ("accountId", `String account_id);
      ("update", `Assoc updates)
    ] in
    
    execute env ~ctx ~session batch_args

  let cleanup_old_emails env ~ctx ~session ~mailbox_role ~older_than_days =
    let account_id = Session_utils.get_primary_mail_account session in
    (* First find the mailbox with the specified role *)
    match Email_methods.find_mailbox_by_role env ~ctx ~session ~account_id mailbox_role with
    | Ok (Some mailbox_json) ->
        (try
           let open Yojson.Safe.Util in
           let mailbox_id = mailbox_json |> member "id" |> to_string in
           (* Create a filter for old emails in this mailbox *)
           let cutoff_date = Unix.time () -. (float_of_int older_than_days *. 86400.0) in
           let date_str = Printf.sprintf "%.0f" cutoff_date in
           let filter = `Assoc [
             ("inMailbox", `String mailbox_id);
             ("before", `String date_str)
           ] in
           (* Query for old emails first, then destroy them *)
           let query_call_id = "cleanup-query-" ^ string_of_int (Random.int 10000) in
           let set_call_id = "cleanup-set-" ^ string_of_int (Random.int 10000) in
           
           let query_args = `Assoc [
             ("accountId", `String account_id);
             ("filter", filter)
           ] in
           
           let set_args = `Assoc [
             ("accountId", `String account_id);
             ("#destroy", `Assoc [
               ("resultOf", `String query_call_id);
               ("name", `String (Jmap.Method_names.method_to_string `Email_query));
               ("path", `String "/ids")
             ])
           ] in
           
           let req_builder = build ctx in
           let req_builder = using req_builder [`Core; `Mail] in
           let req_builder = add_method_call req_builder `Email_query query_args query_call_id in
           let req_builder = add_method_call req_builder `Email_set set_args set_call_id
           in
           match jmap_execute env req_builder with
           | Ok response ->
               (match Response.extract_method ~method_name:`Email_set ~method_call_id:set_call_id response with
                | Ok json -> Ok json
                | Error e -> Error e)
           | Error e -> Error e
         with
         | exn -> Error (Jmap.Error.protocol_error 
             ("Failed to parse mailbox: " ^ Printexc.to_string exn)))
    | Ok None -> Error (Jmap.Error.protocol_error 
        ("Mailbox with role '" ^ mailbox_role ^ "' not found"))
    | Error e -> Error e

  let organize_by_sender _env ~ctx:_ ~session:_ ~rules =
    (* This would be quite complex to implement properly, as it requires:
       1. Finding/creating target mailboxes for each rule
       2. Querying emails by sender
       3. Moving emails to appropriate mailboxes
       For now, return a basic structure indicating the operation would proceed *)
    let rule_count = List.length rules in
    let result = `Assoc [
      ("processed", `Int rule_count);
      ("message", `String "Sender organization rules would be applied")
    ] in
    Ok result

  let execute_with_progress env ~ctx ~session ?account_id ~progress_fn batch =
    (* Report progress at start *)
    progress_fn { current = 0; total = 1; message = "Starting batch operation..." };
    
    (* Execute the batch operation *)
    let result = execute env ~ctx ~session ?account_id batch in
    
    (* Report completion *)
    progress_fn { current = 1; total = 1; message = "Batch operation completed" };
    
    result
end

module Email_submission = Email_submission