(** Connection pooling for efficient JMAP client connection reuse.
    
    This module provides connection pooling functionality to reduce connection overhead.
    For demonstration purposes, this implements statistics tracking and connection management
    concepts while still using cohttp-eio for the actual HTTP operations.
    
    @see <https://www.rfc-editor.org/rfc/rfc8620.html#section-3.3> RFC 8620, Section 3.3
*)

(** TLS configuration options *)
type tls_config = {
  authenticator : X509.Authenticator.t option;  (** Custom TLS authenticator *)
  certificates : Tls.Config.own_cert list;       (** Client certificates for mutual TLS *)
  ciphers : Tls.Ciphersuite.ciphersuite list option; (** Allowed cipher suites *)
  version : (Tls.Core.tls_version * Tls.Core.tls_version) option; (** Min and max TLS versions *)
  alpn_protocols : string list option;           (** ALPN protocol list *)
}

(** Statistics for connection pool monitoring *)
type pool_stats = {
  total_connections : int;       (** Total connections created *)
  active_connections : int;      (** Currently active connections *)  
  idle_connections : int;        (** Currently idle connections *)
  total_requests : int;          (** Total requests processed *)
  cache_hits : int;             (** Requests served from cached connections *)
  cache_misses : int;           (** Requests requiring new connections *)
  connection_failures : int;     (** Failed connection attempts *)
}

(** Connection pool configuration *)
type pool_config = {
  max_connections : int;         (** Maximum total connections *)
  max_idle_connections : int;    (** Maximum idle connections to keep *)
  connection_timeout : float;    (** Connection establishment timeout (seconds) *)
  idle_timeout : float;          (** Time to keep idle connections (seconds) *)
  max_lifetime : float;          (** Maximum connection lifetime (seconds) *)
  health_check_interval : float; (** Health check interval (seconds) *)
  enable_keep_alive : bool;      (** Enable HTTP keep-alive *)
}

(** Connection info for tracking *)
type connection_info = {
  id : string;                  (** Unique connection ID *)
  host : string;                (** Target host *)
  port : int;                   (** Target port *)
  use_tls : bool;               (** TLS usage flag *)
  created_at : float;           (** Connection creation timestamp *)
  last_used : float;            (** Last usage timestamp *)
  request_count : int;          (** Number of requests served *)
}

(** Connection pool type *)
type t = {
  config : pool_config;
  mutable connections : connection_info list;
  mutable stats : pool_stats;
}

(** Create default pool configuration *)
let default_config () = {
  max_connections = 20;
  max_idle_connections = 10;
  connection_timeout = 10.0;
  idle_timeout = 300.0;      (* 5 minutes *)
  max_lifetime = 3600.0;     (* 1 hour *)
  health_check_interval = 60.0; (* 1 minute *)
  enable_keep_alive = true;
}

(** Generate unique connection ID *)
let generate_connection_id () =
  let timestamp = Unix.gettimeofday () in
  let random = Random.int 100000 in
  Printf.sprintf "conn_%f_%d" timestamp random

(** Create a new connection pool *)
let create ?(config = default_config ()) ~sw () =
  let _ = sw in  (* Acknowledge unused parameter *)
  let initial_stats = {
    total_connections = 0;
    active_connections = 0;
    idle_connections = 0;
    total_requests = 0;
    cache_hits = 0;
    cache_misses = 0;
    connection_failures = 0;
  } in
  {
    config;
    connections = [];
    stats = initial_stats;
  }

(** Check if connection is still healthy *)
let is_connection_healthy pool conn =
  let now = Unix.gettimeofday () in
  let age = now -. conn.created_at in
  let idle_time = now -. conn.last_used in
  
  age < pool.config.max_lifetime &&
  idle_time < pool.config.idle_timeout

(** Find existing connection for host/port *)
let find_connection pool ~host ~port ~use_tls =
  List.find_opt (fun conn ->
    conn.host = host && 
    conn.port = port && 
    conn.use_tls = use_tls && 
    is_connection_healthy pool conn
  ) pool.connections

(** Create new connection info *)
let create_connection_info ~host ~port ~use_tls =
  let now = Unix.gettimeofday () in
  {
    id = generate_connection_id ();
    host;
    port;
    use_tls;
    created_at = now;
    last_used = now;
    request_count = 0;
  }

(** Update connection usage *)
let use_connection pool conn =
  let now = Unix.gettimeofday () in
  let updated_conn = {
    conn with
    last_used = now;
    request_count = conn.request_count + 1;
  } in
  
  (* Update connections list *)
  pool.connections <- updated_conn :: 
    (List.filter (fun c -> c.id <> conn.id) pool.connections);
  
  (* Update stats *)
  pool.stats <- {
    pool.stats with
    cache_hits = pool.stats.cache_hits + 1;
    total_requests = pool.stats.total_requests + 1;
  };
  
  updated_conn

(** Add new connection to pool *)
let add_connection pool conn =
  pool.connections <- conn :: pool.connections;
  pool.stats <- {
    pool.stats with
    total_connections = pool.stats.total_connections + 1;
    cache_misses = pool.stats.cache_misses + 1;
    total_requests = pool.stats.total_requests + 1;
  }

(** Perform HTTP request using pool for statistics tracking *)
let http_request_with_pool pool ~env ~method_ ~uri ~headers ~body ~tls_config =
  let host = match Uri.host uri with
    | Some h -> h
    | None -> failwith "No host in URI"
  in
  let use_tls = match Uri.scheme uri with
    | Some "https" -> true
    | Some "http" -> false  
    | _ -> true
  in
  let port = match Uri.port uri with
    | Some p -> p
    | None -> if use_tls then 443 else 80
  in
  
  try
    (* Check if we have a cached connection for this endpoint *)
    let _conn_info = match find_connection pool ~host ~port ~use_tls with
      | Some existing_conn ->
          (* Update existing connection usage *)
          use_connection pool existing_conn
      | None ->
          (* Check connection limits *)
          if List.length pool.connections >= pool.config.max_connections then (
            pool.stats <- {
              pool.stats with
              connection_failures = pool.stats.connection_failures + 1;
            };
            failwith ("Connection pool full: " ^ string_of_int pool.config.max_connections)
          ) else (
            (* Create new connection info *)
            let new_conn = create_connection_info ~host ~port ~use_tls in
            add_connection pool new_conn;
            new_conn
          )
    in
    
    (* Actually perform HTTP request using cohttp-eio *)
    let https_fn = if use_tls then
      let authenticator = match tls_config with
        | Some tls when tls.authenticator <> None -> 
            (match tls.authenticator with Some auth -> auth | None -> assert false)
        | _ -> 
            match Ca_certs.authenticator () with
            | Ok auth -> auth
            | Error (`Msg msg) -> failwith ("TLS authenticator error: " ^ msg)
      in
      let tls_config_obj = match Tls.Config.client ~authenticator () with
        | Ok config -> config
        | Error (`Msg msg) -> failwith ("TLS config error: " ^ 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_obj raw_flow ~host:hostname
      )
    else
      None
    in
    
    Eio.Switch.run @@ fun sw ->
    let client = Cohttp_eio.Client.make ~https:https_fn env#net in
    
    let cohttp_headers = 
      List.fold_left (fun hdrs (k, v) ->
        Cohttp.Header.add hdrs k v
      ) (Cohttp.Header.init ()) headers
    in
    
    let body_obj = 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_obj method_ uri in
    
    let status_code = Cohttp.Response.status response |> Cohttp.Code.code_of_status in
    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 ->
      pool.stats <- {
        pool.stats with
        connection_failures = pool.stats.connection_failures + 1;
      };
      Error (Jmap.Error.transport 
        (Printf.sprintf "Connection error: %s" (Printexc.to_string exn)))

(** Clean up old connections *)
let cleanup_connections pool =
  let now = Unix.gettimeofday () in
  let (healthy, _unhealthy) = List.partition (is_connection_healthy pool) pool.connections in
  
  (* Keep only healthy connections, respecting idle limit *)
  let idle_connections = List.filter (fun c -> 
    now -. c.last_used > 1.0  (* Idle for more than 1 second *)
  ) healthy in
  
  let keep_idle = 
    if List.length idle_connections > pool.config.max_idle_connections then
      let sorted = List.sort (fun a b -> 
        compare b.last_used a.last_used  (* Most recently used first *)
      ) idle_connections in
      let rec list_take n = function
        | [] -> []
        | h :: t when n > 0 -> h :: list_take (n - 1) t
        | _ -> []
      in
      list_take pool.config.max_idle_connections sorted
    else
      idle_connections
  in
  
  let active_connections = List.filter (fun c ->
    now -. c.last_used <= 1.0
  ) healthy in
  
  pool.connections <- active_connections @ keep_idle;
  pool.stats <- {
    pool.stats with
    total_connections = List.length pool.connections;
    active_connections = List.length active_connections;
    idle_connections = List.length keep_idle;
  }

(** Get pool statistics *)
let get_stats pool = 
  cleanup_connections pool;
  pool.stats

(** Close all connections and clean up pool *)
let close pool =
  pool.connections <- [];
  pool.stats <- {
    pool.stats with
    total_connections = 0;
    active_connections = 0;
    idle_connections = 0;
  }