(* Unix-specific JMAP client implementation interface. *)

open Jmap
open Jmap.Types
open Jmap.Error
open Jmap.Session
open Jmap.Wire

(* Configuration options for a JMAP client context *)
type client_config = {
  connect_timeout : float option;        (* Connection timeout in seconds *)
  request_timeout : float option;        (* Request timeout in seconds *)
  max_concurrent_requests : int option;  (* Maximum concurrent requests *)
  max_request_size : int option;         (* Maximum request size in bytes *)
  user_agent : string option;            (* User-Agent header value *)
  authentication_header : string option; (* Custom Authentication header name *)
}

(* Authentication method options *)
type auth_method = 
  | Basic of string * string                  (* Basic auth with username and password *)
  | Bearer of string                          (* Bearer token auth *)
  | Custom of (string * string)               (* Custom header name and value *)
  | Session_cookie of (string * string)       (* Session cookie name and value *)
  | No_auth                                   (* No authentication *)

(* The internal state of a JMAP client connection *)
type context = {
  config: client_config;
  mutable session_url: Uri.t option;
  mutable session: Session.t option;
  mutable auth: auth_method;
}

(* Represents an active EventSource connection *)
type event_source_connection = {
  event_url: Uri.t;
  mutable is_connected: bool;
}

(* A request builder for constructing and sending JMAP requests *)
type request_builder = {
  ctx: context;
  mutable using: string list;
  mutable method_calls: Invocation.t list;
}

(* Create default configuration options *)
let default_config () = {
  connect_timeout = Some 30.0;
  request_timeout = Some 300.0;
  max_concurrent_requests = Some 4;
  max_request_size = Some (1024 * 1024 * 10); (* 10 MB *)
  user_agent = Some "OCaml JMAP Unix Client/1.0";
  authentication_header = None;
}

(* Create a client context with the specified configuration *)
let create_client ?(config = default_config ()) () = {
  config;
  session_url = None;
  session = None;
  auth = No_auth;
}

(* Mock implementation for the Unix connection *)
let connect ctx ?session_url ?username ~host ?port ?auth_method () =
  (* In a real implementation, this would use Unix HTTP functions *)
  let auth = match auth_method with
    | Some auth -> auth
    | None -> No_auth
  in
  
  (* Store the auth method for future requests *)
  ctx.auth <- auth;
  
  (* Set session URL, either directly or after discovery *)
  let session_url = match session_url with
    | Some url -> url
    | None ->
        (* In a real implementation, this would perform RFC 8620 discovery *)
        let proto = "https" in
        let host_with_port = match port with
          | Some p -> host ^ ":" ^ string_of_int p
          | None -> host
        in
        Uri.of_string (proto ^ "://" ^ host_with_port ^ "/.well-known/jmap")
  in
  ctx.session_url <- Some session_url;
  
  (* Create a mock session object for this example *)
  let caps = Hashtbl.create 4 in
  Hashtbl.add caps Jmap.capability_core (`Assoc []);
  
  let accounts = Hashtbl.create 1 in
  let acct = Account.v 
    ~name:"user@example.com" 
    ~is_personal:true 
    ~is_read_only:false 
    ()
  in
  Hashtbl.add accounts "u1" acct;
  
  let primary = Hashtbl.create 1 in
  Hashtbl.add primary Jmap.capability_core "u1";
  
  let api_url = 
    Uri.of_string ("https://" ^ host ^ "/api/jmap")
  in
  
  let session = Session.v
    ~capabilities:caps
    ~accounts
    ~primary_accounts:primary
    ~username:"user@example.com"
    ~api_url
    ~download_url:(Uri.of_string ("https://" ^ host ^ "/download/{accountId}/{blobId}"))
    ~upload_url:(Uri.of_string ("https://" ^ host ^ "/upload/{accountId}"))
    ~event_source_url:(Uri.of_string ("https://" ^ host ^ "/eventsource"))
    ~state:"1"
    ()
  in
  
  ctx.session <- Some session;
  Ok (ctx, session)

(* Create a request builder for constructing a JMAP request *)
let build ctx = {
  ctx;
  using = [Jmap.capability_core]; (* Default to core capability *)
  method_calls = [];
}

(* Set the using capabilities for a request *)
let using builder capabilities = 
  { builder with using = capabilities }

(* Add a method call to a request builder *)
let add_method_call builder name args id =
  let call = Invocation.v 
    ~method_name:name 
    ~arguments:args 
    ~method_call_id:id 
    ()
  in
  { builder with method_calls = builder.method_calls @ [call] }

(* Create a reference to a previous method call result *)
let create_reference result_of name =
  Jmap.Wire.Result_reference.v 
    ~result_of
    ~name
    ~path:"" (* In a real implementation, this would include a JSON pointer *)
    ()

(* Execute a request and return the response *)
let execute builder =
  match builder.ctx.session with
  | None -> Error (protocol_error "No active session")
  | Some session ->
      (* In a real implementation, this would create and send an HTTP request *)
      
      (* Create a mock response for this implementation *)
      let results = List.map (fun call ->
        let method_name = Invocation.method_name call in
        let call_id = Invocation.method_call_id call in
        if method_name = "Core/echo" then
          (* Echo method implementation *)
          Ok call
        else
          (* For other methods, return a method error *)
          Error (
            Method_error.v 
              ~description:(Method_error_description.v 
                ~description:"Method not implemented in mock" 
                ())
              `ServerUnavailable,
            "Mock implementation"
          )
      ) builder.method_calls in
      
      let resp = Response.v
        ~method_responses:results
        ~session_state:(session |> Session.state)
        ()
      in
      Ok resp

(* Perform a JMAP API request *)
let request ctx req =
  match ctx.session_url, ctx.session with
  | None, _ -> Error (protocol_error "No session URL configured")
  | _, None -> Error (protocol_error "No active session")
  | Some url, Some session ->
      (* In a real implementation, this would serialize the request and send it *)
      
      (* Mock response implementation *)
      let method_calls = Request.method_calls req in
      let results = List.map (fun call ->
        let method_name = Invocation.method_name call in
        let call_id = Invocation.method_call_id call in
        if method_name = "Core/echo" then
          (* Echo method implementation *)
          Ok call
        else
          (* For other methods, return a method error *)
          Error (
            Method_error.v 
              ~description:(Method_error_description.v 
                ~description:"Method not implemented in mock" 
                ())
              `ServerUnavailable,
            "Mock implementation"
          )
      ) method_calls in
      
      let resp = Response.v
        ~method_responses:results
        ~session_state:(session |> Session.state)
        ()
      in
      Ok resp

(* Upload binary data *)
let upload ctx ~account_id ~content_type ~data_stream =
  match ctx.session with
  | None -> Error (protocol_error "No active session")
  | Some session ->
      (* In a real implementation, would upload the data stream *)
      
      (* Mock success response *)
      let response = Jmap.Binary.Upload_response.v
        ~account_id
        ~blob_id:"b123456"
        ~type_:content_type
        ~size:1024 (* Mock size *)
        ()
      in
      Ok response

(* Download binary data *)
let download ctx ~account_id ~blob_id ?content_type ?name =
  match ctx.session with
  | None -> Error (protocol_error "No active session")
  | Some session ->
      (* In a real implementation, would download the data and return a stream *)
      
      (* Mock data stream - in real code, this would be read from the HTTP response *)
      let mock_data = "This is mock downloaded data for blob " ^ blob_id in
      let seq = Seq.cons mock_data Seq.empty in
      Ok seq

(* Copy blobs between accounts *)
let copy_blobs ctx ~from_account_id ~account_id ~blob_ids =
  match ctx.session with
  | None -> Error (protocol_error "No active session")
  | Some session ->
      (* In a real implementation, would perform server-side copy *)
      
      (* Mock success response with first blob copied and second failed *)
      let copied = Hashtbl.create 1 in
      Hashtbl.add copied (List.hd blob_ids) "b999999";
      
      let response = Jmap.Binary.Blob_copy_response.v
        ~from_account_id
        ~account_id
        ~copied
        ()
      in
      Ok response
      
(* Connect to the EventSource for push notifications *)
let connect_event_source ctx ?types ?close_after ?ping =
  match ctx.session with
  | None -> Error (protocol_error "No active session")
  | Some session ->
      (* In a real implementation, would connect to EventSource URL *)
      
      (* Create mock connection *)
      let event_url = Session.event_source_url session in
      let conn = { event_url; is_connected = true } in
      
      (* Create a mock event sequence *)
      let mock_state_change = 
        let changed = Hashtbl.create 1 in
        let account_id = "u1" in
        let state_map = Hashtbl.create 2 in
        Hashtbl.add state_map "Email" "s123";
        Hashtbl.add state_map "Mailbox" "s456";
        Hashtbl.add changed account_id state_map;
        
        Push.State_change.v ~changed ()
      in
      
      let ping_data = 
        Push.Event_source_ping_data.v ~interval:30 ()
      in
      
      (* Create a sequence with one state event and one ping event *)
      let events = Seq.cons (`State mock_state_change) 
                     (Seq.cons (`Ping ping_data) Seq.empty) in
      
      Ok (conn, events)

(* Create a websocket connection for JMAP over WebSocket *)
let connect_websocket ctx =
  match ctx.session with
  | None -> Error (protocol_error "No active session")
  | Some session ->
      (* In a real implementation, would connect via WebSocket *)
      
      (* Mock connection *)
      let event_url = Session.api_url session in
      let conn = { event_url; is_connected = true } in
      Ok conn

(* Send a message over a websocket connection *)
let websocket_send conn req =
  if not conn.is_connected then
    Error (protocol_error "WebSocket not connected")
  else
    (* In a real implementation, would send over WebSocket *)
    
    (* Mock response (same as request function) *)
    let method_calls = Request.method_calls req in
    let results = List.map (fun call ->
      let method_name = Invocation.method_name call in
      let call_id = Invocation.method_call_id call in
      if method_name = "Core/echo" then
        Ok call
      else
        Error (
          Method_error.v 
            ~description:(Method_error_description.v 
              ~description:"Method not implemented in mock" 
              ())
            `ServerUnavailable,
          "Mock implementation"
        )
    ) method_calls in
    
    let resp = Response.v
      ~method_responses:results
      ~session_state:"1"
      ()
    in
    Ok resp

(* Close an EventSource or WebSocket connection *)
let close_connection conn =
  if not conn.is_connected then
    Error (protocol_error "Connection already closed")
  else begin
    conn.is_connected <- false;
    Ok ()
  end

(* Close the JMAP connection context *)
let close ctx =
  ctx.session <- None;
  ctx.session_url <- None;
  Ok ()

(* Helper functions for common tasks *)

(* Helper to get a single object by ID *)
let get_object ctx ~method_name ~account_id ~object_id ?properties =
  let properties_param = match properties with
    | Some props -> `List (List.map (fun p -> `String p) props)
    | None -> `Null
  in
  
  let args = `Assoc [
    ("accountId", `String account_id);
    ("ids", `List [`String object_id]);
    ("properties", properties_param);
  ] in
  
  let request_builder = build ctx
    |> add_method_call method_name args "r1"
  in
  
  match execute request_builder with
  | Error e -> Error e
  | Ok response ->
      (* Find the method response and extract the list with the object *)
      match response |> Response.method_responses with
      | [Ok invocation] when Invocation.method_name invocation = method_name ^ "/get" -> 
          let args = Invocation.arguments invocation in
          begin match Yojson.Safe.Util.member "list" args with
          | `List [obj] -> Ok obj
          | _ -> Error (protocol_error "Object not found or invalid response")
          end
      | _ ->
          Error (protocol_error "Method response not found")

(* Helper to set up the connection with minimal options *)
let quick_connect ~host ~username ~password =
  let ctx = create_client () in
  connect ctx ~host ~auth_method:(Basic(username, password)) ()

(* Perform a Core/echo request to test connectivity *)
let echo ctx ?data () =
  let data = match data with
    | Some d -> d
    | None -> `Assoc [("hello", `String "world")]
  in
  
  let request_builder = build ctx
    |> add_method_call "Core/echo" data "echo1"
  in
  
  match execute request_builder with
  | Error e -> Error e
  | Ok response ->
      (* Find the Core/echo response and extract the echoed data *)
      match response |> Response.method_responses with
      | [Ok invocation] when Invocation.method_name invocation = "Core/echo" -> 
          Ok (Invocation.arguments invocation)
      | _ ->
          Error (protocol_error "Echo response not found")

(* High-level email operations *)
module Email = struct
  open Jmap_email.Types
  
  (* Get an email by ID *)
  let get_email ctx ~account_id ~email_id ?properties () =
    let props = match properties with
      | Some p -> p
      | None -> List.map email_property_to_string detailed_email_properties
    in
    
    match get_object ctx ~method_name:"Email/get" ~account_id ~object_id:email_id ~properties:props with
    | Error e -> Error e
    | Ok json ->
        (* In a real implementation, would parse the JSON into an Email.t structure *)
        let mock_email = Email.create
          ~id:email_id
          ~thread_id:"t12345"
          ~mailbox_ids:(let h = Hashtbl.create 1 in Hashtbl.add h "inbox" true; h)
          ~keywords:(Keywords.of_list [Keywords.Seen])
          ~subject:"Mock Email Subject"
          ~preview:"This is a mock email..."
          ~from:[Email_address.v ~name:"Sender Name" ~email:"sender@example.com" ()]
          ~to_:[Email_address.v ~email:"recipient@example.com" ()]
          ()
        in
        Ok mock_email
  
  (* Search for emails using a filter *)
  let search_emails ctx ~account_id ~filter ?sort ?limit ?position ?properties () =
    (* Create the query args *)
    let args = `Assoc [
      ("accountId", `String account_id);
      ("filter", Jmap.Methods.Filter.to_json filter);
      ("sort", match sort with
        | Some s -> `List [] (* Would convert sort params *)
        | None -> `List [`Assoc [("property", `String "receivedAt"); ("isAscending", `Bool false)]]);
      ("limit", match limit with
        | Some l -> `Int l
        | None -> `Int 20);
      ("position", match position with
        | Some p -> `Int p
        | None -> `Int 0);
    ] in
    
    let request_builder = build ctx
      |> add_method_call "Email/query" args "q1"
    in
    
    (* If properties were provided, add a Email/get method call as well *)
    let request_builder = match properties with
      | Some _ -> 
          let get_args = `Assoc [
            ("accountId", `String account_id);
            ("#ids", `Assoc [
              ("resultOf", `String "q1");
              ("name", `String "Email/query");
              ("path", `String "/ids")
            ]);
            ("properties", match properties with
              | Some p -> `List (List.map (fun prop -> `String prop) p)
              | None -> `Null);
          ] in
          add_method_call request_builder "Email/get" get_args "g1"
      | None -> request_builder
    in
    
    match execute request_builder with
    | Error e -> Error e
    | Ok response ->
        (* Find the query response and extract the IDs *)
        match Response.method_responses response with
        | [Ok q_inv; Ok g_inv] 
          when Invocation.method_name q_inv = "Email/query" 
          && Invocation.method_name g_inv = "Email/get" ->
            
            (* Extract IDs from query response *)
            let q_args = Invocation.arguments q_inv in
            let ids = match Yojson.Safe.Util.member "ids" q_args with
              | `List l -> List.map Yojson.Safe.Util.to_string l
              | _ -> []
            in
            
            (* Extract emails from get response *)
            let g_args = Invocation.arguments g_inv in
            (* In a real implementation, would parse each email in the list *)
            let emails = List.map (fun id ->
              Email.create
                ~id
                ~thread_id:("t" ^ id)
                ~subject:(Printf.sprintf "Mock Email %s" id)
                ()
            ) ids in
            
            Ok (ids, Some emails)
            
        | [Ok q_inv] when Invocation.method_name q_inv = "Email/query" ->
            (* If only query was performed (no properties requested) *)
            let q_args = Invocation.arguments q_inv in
            let ids = match Yojson.Safe.Util.member "ids" q_args with
              | `List l -> List.map Yojson.Safe.Util.to_string l
              | _ -> []
            in
            
            Ok (ids, None)
            
        | _ ->
            Error (protocol_error "Query response not found")
  
  (* Mark multiple emails with a keyword *)
  let mark_emails ctx ~account_id ~email_ids ~keyword () =
    (* Create the set args with a patch to add the keyword *)
    let keyword_patch = Jmap_email.Keyword_ops.add_keyword_patch keyword in
    
    (* Create patches map for each email *)
    let update = Hashtbl.create (List.length email_ids) in
    List.iter (fun id ->
      Hashtbl.add update id keyword_patch
    ) email_ids;
    
    let args = `Assoc [
      ("accountId", `String account_id);
      ("update", `Assoc (
        List.map (fun id ->
          (id, `Assoc (List.map (fun (path, value) ->
            (path, value)
          ) keyword_patch))
        ) email_ids
      ));
    ] in
    
    let request_builder = build ctx
      |> add_method_call "Email/set" args "s1"
    in
    
    match execute request_builder with
    | Error e -> Error e
    | Ok response ->
        (* In a real implementation, would check for errors *)
        Ok ()
  
  (* Mark emails as seen/read *)
  let mark_as_seen ctx ~account_id ~email_ids () =
    mark_emails ctx ~account_id ~email_ids ~keyword:Keywords.Seen ()
  
  (* Mark emails as unseen/unread *)
  let mark_as_unseen ctx ~account_id ~email_ids () =
    let keyword_patch = Jmap_email.Keyword_ops.mark_unseen_patch () in
    
    (* Create patches map for each email *)
    let update = Hashtbl.create (List.length email_ids) in
    List.iter (fun id ->
      Hashtbl.add update id keyword_patch
    ) email_ids;
    
    let args = `Assoc [
      ("accountId", `String account_id);
      ("update", `Assoc (
        List.map (fun id ->
          (id, `Assoc (List.map (fun (path, value) ->
            (path, value)
          ) keyword_patch))
        ) email_ids
      ));
    ] in
    
    let request_builder = build ctx
      |> add_method_call "Email/set" args "s1"
    in
    
    match execute request_builder with
    | Error e -> Error e
    | Ok _response -> Ok ()
  
  (* Move emails to a different mailbox *)
  let move_emails ctx ~account_id ~email_ids ~mailbox_id ?remove_from_mailboxes () =
    (* Create patch to add to destination mailbox *)
    let add_patch = [("mailboxIds/" ^ mailbox_id, `Bool true)] in
    
    (* If remove_from_mailboxes is specified, add patches to remove *)
    let remove_patch = match remove_from_mailboxes with
      | Some mailboxes -> 
          List.map (fun mbx -> ("mailboxIds/" ^ mbx, `Null)) mailboxes
      | None -> []
    in
    
    (* Combine patches *)
    let patches = add_patch @ remove_patch in
    
    (* Create patches map for each email *)
    let update = Hashtbl.create (List.length email_ids) in
    List.iter (fun id ->
      Hashtbl.add update id patches
    ) email_ids;
    
    let args = `Assoc [
      ("accountId", `String account_id);
      ("update", `Assoc (
        List.map (fun id ->
          (id, `Assoc (List.map (fun (path, value) ->
            (path, value)
          ) patches))
        ) email_ids
      ));
    ] in
    
    let request_builder = build ctx
      |> add_method_call "Email/set" args "s1"
    in
    
    match execute request_builder with
    | Error e -> Error e
    | Ok _response -> Ok ()
  
  (* Import an RFC822 message *)
  let import_email ctx ~account_id ~rfc822 ~mailbox_ids ?keywords ?received_at () =
    (* In a real implementation, would first upload the message as a blob *)
    let mock_blob_id = "b9876" in
    
    (* Create the Email/import call *)
    let args = `Assoc [
      ("accountId", `String account_id);
      ("emails", `Assoc [
        ("msg1", `Assoc [
          ("blobId", `String mock_blob_id);
          ("mailboxIds", `Assoc (
            List.map (fun id -> (id, `Bool true)) mailbox_ids
          ));
          ("keywords", match keywords with
            | Some kws -> 
                `Assoc (List.map (fun k -> 
                  (Types.Keywords.to_string k, `Bool true)) kws)
            | None -> `Null);
          ("receivedAt", match received_at with
            | Some d -> `String (string_of_float d) (* Would format as RFC3339 *)
            | None -> `Null);
        ])
      ]);
    ] in
    
    let request_builder = build ctx
      |> add_method_call "Email/import" args "i1"
    in
    
    match execute request_builder with
    | Error e -> Error e
    | Ok response ->
        (* In a real implementation, would extract the created ID *)
        Ok "e12345"
end