(** HomeKit Accessory Protocol (HAP) implementation.

    This module implements the HAP protocol for controlling HomeKit accessories:
    - Discovery via mDNS (_hap._tcp)
    - Pair Setup using SRP-6a
    - Pair Verify using Curve25519
    - Encrypted sessions using ChaCha20-Poly1305 *)

let log_src = Logs.Src.create "hap"

module Log = (val Logs.src_log log_src : Logs.LOG)
open Result.Syntax

(** {1 Errors} *)

let err_pair_setup code =
  Error (`Msg (Fmt.str "Pair setup error: %d" (Char.code code.[0])))

let err_pair_setup_m4 code =
  Error (`Msg (Fmt.str "Pair setup M4 error: %d" (Char.code code.[0])))

let err_pair_setup_m6 code =
  Error (`Msg (Fmt.str "Pair setup M6 error: %d" (Char.code code.[0])))

let err_pair_verify_m2 code =
  Error (`Msg (Fmt.str "Pair verify M2 error: %d" (Char.code code.[0])))

let err_pair_verify_m4 code =
  Error (`Msg (Fmt.str "Pair verify M4 error: %d" (Char.code code.[0])))

(* Helper to convert IP string to Eio address *)
let ipv4_of_string ip =
  Eio.Net.Ipaddr.of_raw (Ipaddr.V4.to_octets (Ipaddr.V4.of_string_exn ip))

(* TLV encoding for HAP *)
module Tlv = struct
  type t = (int * string) list

  let empty = []
  let add typ value tlv = (typ, value) :: tlv

  let get typ tlv =
    List.find_map (fun (t, v) -> if t = typ then Some v else None) tlv

  let get_exn typ tlv =
    match get typ tlv with
    | Some v -> v
    | None -> Fmt.failwith "TLV type %d not found" typ

  (* Encode TLV to bytes - values > 255 bytes are split *)
  let encode tlv =
    let buf = Buffer.create 256 in
    List.iter
      (fun (typ, value) ->
        let len = String.length value in
        let rec write_chunks offset =
          if offset >= len then ()
          else begin
            let chunk_len = min 255 (len - offset) in
            Buffer.add_char buf (Char.chr typ);
            Buffer.add_char buf (Char.chr chunk_len);
            Buffer.add_substring buf value offset chunk_len;
            write_chunks (offset + chunk_len)
          end
        in
        if len = 0 then begin
          Buffer.add_char buf (Char.chr typ);
          Buffer.add_char buf '\x00'
        end
        else write_chunks 0)
      (List.rev tlv);
    Buffer.contents buf

  (* Decode TLV from bytes - concatenate split values.
     Stops gracefully on truncated input, returning entries parsed so far. *)
  let decode data =
    let len = String.length data in
    let rec parse offset acc =
      if offset >= len then List.rev acc
      else if offset + 2 > len then List.rev acc
      else begin
        let typ = Char.code data.[offset] in
        let vlen = Char.code data.[offset + 1] in
        if offset + 2 + vlen > len then List.rev acc
        else begin
          let value = String.sub data (offset + 2) vlen in
          (* Concatenate with previous if same type *)
          let acc =
            match acc with
            | (prev_typ, prev_val) :: rest when prev_typ = typ ->
                (typ, prev_val ^ value) :: rest
            | _ -> (typ, value) :: acc
          in
          parse (offset + 2 + vlen) acc
        end
      end
    in
    parse 0 []
end

(* TLV types for HAP *)
module Tlv_type = struct
  let method_ = 0x00
  let identifier = 0x01
  let salt = 0x02
  let public_key = 0x03
  let proof = 0x04
  let encrypted_data = 0x05
  let state = 0x06
  let error = 0x07
  let retry_delay = 0x08
  let certificate = 0x09
  let signature = 0x0a
  let permissions = 0x0b
  let fragment_data = 0x0c
  let fragment_last = 0x0d
  let separator = 0xff
end

(* HAP errors *)
module Hap_error = struct
  let unknown = 0x01
  let authentication = 0x02
  let backoff = 0x03
  let max_peers = 0x04
  let max_tries = 0x05
  let unavailable = 0x06
  let busy = 0x07
end

(* ChaCha20-Poly1305 encryption *)
let chacha20_poly1305_encrypt ~key ~nonce ~aad data =
  let key = Crypto.Chacha20.of_secret key in
  let encrypted =
    Crypto.Chacha20.authenticate_encrypt ~key ~nonce ~adata:aad data
  in
  encrypted

let chacha20_poly1305_decrypt ~key ~nonce ~aad data =
  let key = Crypto.Chacha20.of_secret key in
  match Crypto.Chacha20.authenticate_decrypt ~key ~nonce ~adata:aad data with
  | Some decrypted -> Ok decrypted
  | None -> Error (`Msg "Decryption failed")

(* HKDF-SHA512 key derivation *)
let hkdf_sha512 ~salt ~ikm ~info ~length =
  let prk = Hkdf.extract ~hash:`SHA512 ~salt ikm in
  Hkdf.expand ~hash:`SHA512 ~prk ~info length

(* Ed25519 key pair *)
module Ed25519 = struct
  type keypair = { secret : string; public : string }

  let generate () =
    let priv, pub = Crypto_ec.Ed25519.generate () in
    {
      secret = Crypto_ec.Ed25519.priv_to_octets priv;
      public = Crypto_ec.Ed25519.pub_to_octets pub;
    }

  let sign ~secret data =
    match Crypto_ec.Ed25519.priv_of_octets secret with
    | Ok priv -> Crypto_ec.Ed25519.sign ~key:priv data
    | Error _ -> failwith "Invalid Ed25519 private key"

  let verify ~public ~signature data =
    match Crypto_ec.Ed25519.pub_of_octets public with
    | Ok pub -> Crypto_ec.Ed25519.verify ~key:pub signature ~msg:data
    | Error _ -> false
end

(* X25519 key exchange *)
module X25519 = struct
  type keypair = { secret : string; public : string }

  let generate () =
    let secret, public = Crypto_ec.X25519.gen_key () in
    { secret = Crypto_ec.X25519.secret_to_octets secret; public }

  let shared_secret ~secret ~public =
    match Crypto_ec.X25519.secret_of_octets secret with
    | Ok (secret, _) -> (
        match Crypto_ec.X25519.key_exchange secret public with
        | Ok s -> Ok s
        | Error _ -> Error (`Msg "X25519 key exchange failed"))
    | Error _ -> Error (`Msg "Invalid X25519 secret key")
end

(* Controller pairing data *)
type pairing = {
  accessory_id : string;
  accessory_ltpk : string; (* Long-term public key *)
  controller_id : string;
  controller_ltsk : string; (* Long-term secret key *)
  controller_ltpk : string; (* Long-term public key *)
}

(* Accessory info from discovery *)
type accessory_info = {
  name : string;
  device_id : string;
  ip : string;
  port : int;
  model : string option;
  config_num : int;
  state_num : int;
  category : int;
  paired : bool;
}

(* HAP session state *)
type session = {
  pairing : pairing;
  ip : string;
  port : int;
  encrypt_key : string;
  decrypt_key : string;
  mutable encrypt_count : int64;
  mutable decrypt_count : int64;
}

(* Encrypted frame for HAP sessions *)
let encrypt_frame session data =
  let nonce = Bytes.create 12 in
  Bytes.set_int64_le nonce 4 session.encrypt_count;
  session.encrypt_count <- Int64.succ session.encrypt_count;
  let len = String.length data in
  let len_bytes = Bytes.create 2 in
  Bytes.set_uint16_le len_bytes 0 len;
  let aad = Bytes.to_string len_bytes in
  let encrypted =
    chacha20_poly1305_encrypt ~key:session.encrypt_key
      ~nonce:(Bytes.to_string nonce) ~aad data
  in
  aad ^ encrypted

let decrypt_frame session data =
  if String.length data < 18 then Error (`Msg "Frame too short")
  else begin
    let len = Char.code data.[0] lor (Char.code data.[1] lsl 8) in
    let encrypted = String.sub data 2 (String.length data - 2) in
    if String.length encrypted < len + 16 then Error (`Msg "Frame truncated")
    else begin
      let nonce = Bytes.create 12 in
      Bytes.set_int64_le nonce 4 session.decrypt_count;
      session.decrypt_count <- Int64.succ session.decrypt_count;
      let aad = String.sub data 0 2 in
      chacha20_poly1305_decrypt ~key:session.decrypt_key
        ~nonce:(Bytes.to_string nonce) ~aad encrypted
    end
  end

(* HTTP helpers *)
let http_post ~net ~clock ~sw ~ip ~port ~path ~content_type:ct ~body =
  let url = Fmt.str "http://%s:%d%s" ip port path in
  let headers = Requests.Headers.(empty |> set_string "Content-Type" ct) in
  let body = Requests.Body.of_string Requests.Mime.octet_stream body in
  let timeout = Requests.Timeout.v ~connect:10.0 ~read:10.0 () in
  let response =
    Requests.One.post ~sw ~clock ~net ~headers ~body ~timeout ~verify_tls:false
      url
  in
  Ok (Requests.Response.text response)

(* Build encrypted M5 request with controller credentials *)
let build_m5 ~session_key_bytes ~enc_key (controller_kp : Ed25519.keypair)
    controller_id =
  let controller_x =
    hkdf_sha512 ~salt:"Pair-Setup-Controller-Sign-Salt" ~ikm:session_key_bytes
      ~info:"Pair-Setup-Controller-Sign-Info" ~length:32
  in
  let sign_data = controller_x ^ controller_id ^ controller_kp.public in
  let signature = Ed25519.sign ~secret:controller_kp.secret sign_data in
  let sub_tlv =
    Tlv.(
      empty
      |> add Tlv_type.identifier controller_id
      |> add Tlv_type.public_key controller_kp.public
      |> add Tlv_type.signature signature
      |> encode)
  in
  let encrypted =
    chacha20_poly1305_encrypt ~key:enc_key ~nonce:"\x00\x00\x00\x00PS-Msg05"
      ~aad:"" sub_tlv
  in
  Tlv.(
    empty |> add Tlv_type.state "\x05"
    |> add Tlv_type.encrypted_data encrypted
    |> encode)

(* Verify M6 accessory response and return pairing info *)
let verify_m6 ~session_key_bytes ~enc_key (controller_kp : Ed25519.keypair)
    controller_id m6 =
  let enc_data = Tlv.get_exn Tlv_type.encrypted_data m6 in
  let* decrypted =
    chacha20_poly1305_decrypt ~key:enc_key ~nonce:"\x00\x00\x00\x00PS-Msg06"
      ~aad:"" enc_data
  in
  let sub_tlv = Tlv.decode decrypted in
  let accessory_id = Tlv.get_exn Tlv_type.identifier sub_tlv in
  let accessory_ltpk = Tlv.get_exn Tlv_type.public_key sub_tlv in
  let accessory_sig = Tlv.get_exn Tlv_type.signature sub_tlv in
  let accessory_x =
    hkdf_sha512 ~salt:"Pair-Setup-Accessory-Sign-Salt" ~ikm:session_key_bytes
      ~info:"Pair-Setup-Accessory-Sign-Info" ~length:32
  in
  let verify_data = accessory_x ^ accessory_id ^ accessory_ltpk in
  if
    not
      (Ed25519.verify ~public:accessory_ltpk ~signature:accessory_sig
         verify_data)
  then Error (`Msg "Accessory signature verification failed")
  else begin
    Log.info (fun f -> f "Pair setup complete! Accessory ID: %s" accessory_id);
    Ok
      {
        accessory_id;
        accessory_ltpk;
        controller_id;
        controller_ltsk = controller_kp.secret;
        controller_ltpk = controller_kp.public;
      }
  end

(* Pair Setup M5/M6 exchange - derive keys, sign, encrypt, and verify *)
let pair_setup_exchange ~net ~sw ~clock ~ip ~port ~session_key_bytes ~enc_key =
  let controller_kp = Ed25519.generate () in
  let controller_id = "maison-controller" in
  let m5 = build_m5 ~session_key_bytes ~enc_key controller_kp controller_id in
  let* m6_body =
    http_post ~net ~clock ~sw ~ip ~port ~path:"/pair-setup"
      ~content_type:"application/pairing+tlv8" ~body:m5
  in
  let m6 = Tlv.decode m6_body in
  match Tlv.get Tlv_type.error m6 with
  | Some e -> err_pair_setup_m6 e
  | None -> verify_m6 ~session_key_bytes ~enc_key controller_kp controller_id m6

(* SRP M3/M4 verify exchange and derive encryption key for M5/M6 *)
let srp_verify ~net ~sw ~clock ~ip ~port ~srp_client ~salt ~big_b ~session_key =
  let big_a = Srp.Client.public_key srp_client in
  let m1_proof =
    Srp.Client.compute_proof srp_client ~salt ~big_b ~session_key
  in
  let n_len = (Z.numbits Srp.n + 7) / 8 in
  let m3 =
    Tlv.(
      empty |> add Tlv_type.state "\x03"
      |> add Tlv_type.public_key (Srp.bytes_of_z ~pad:n_len big_a)
      |> add Tlv_type.proof m1_proof
      |> encode)
  in
  let* m4_body =
    http_post ~net ~clock ~sw ~ip ~port ~path:"/pair-setup"
      ~content_type:"application/pairing+tlv8" ~body:m3
  in
  let m4 = Tlv.decode m4_body in
  match Tlv.get Tlv_type.error m4 with
  | Some e -> err_pair_setup_m4 e
  | None ->
      let m2_proof = Tlv.get_exn Tlv_type.proof m4 in
      if
        not
          (Srp.Client.verify_proof srp_client ~m1:m1_proof ~m2:m2_proof
             ~session_key)
      then Error (`Msg "Server proof verification failed")
      else begin
        Log.info (fun f -> f "SRP verification successful");
        let enc_key =
          hkdf_sha512 ~salt:"Pair-Setup-Encrypt-Salt" ~ikm:session_key
            ~info:"Pair-Setup-Encrypt-Info" ~length:32
        in
        pair_setup_exchange ~net ~sw ~clock ~ip ~port
          ~session_key_bytes:session_key ~enc_key
      end

(* Pair Setup - M1 through M6 *)
let pair_setup ~net ~sw ~clock ~ip ~port ~pin =
  Log.info (fun f -> f "Starting pair setup with %s:%d" ip port);
  let m1 =
    Tlv.(
      empty |> add Tlv_type.state "\x01"
      |> add Tlv_type.method_ "\x00"
      |> encode)
  in
  let* m2_body =
    http_post ~net ~clock ~sw ~ip ~port ~path:"/pair-setup"
      ~content_type:"application/pairing+tlv8" ~body:m1
  in
  let m2 = Tlv.decode m2_body in
  match Tlv.get Tlv_type.error m2 with
  | Some e -> err_pair_setup e
  | None ->
      let salt = Tlv.get_exn Tlv_type.salt m2 in
      let big_b_bytes = Tlv.get_exn Tlv_type.public_key m2 in
      let big_b = Srp.z_of_bytes big_b_bytes in
      Log.info (fun f ->
          f "Received M2, salt=%d bytes, B=%d bytes" (String.length salt)
            (String.length big_b_bytes));
      let srp_client = Srp.Client.create ~username:"Pair-Setup" ~password:pin in
      let* session_key =
        Srp.Client.compute_session_key srp_client ~salt ~big_b
      in
      srp_verify ~net ~sw ~clock ~ip ~port ~srp_client ~salt ~big_b ~session_key

(* Build encrypted M3 verify request *)
let build_verify_m3 ~enc_key ~pairing ~kp ~accessory_pk =
  let sign_data = kp.X25519.public ^ pairing.controller_id ^ accessory_pk in
  let signature = Ed25519.sign ~secret:pairing.controller_ltsk sign_data in
  let sub_tlv =
    Tlv.(
      empty
      |> add Tlv_type.identifier pairing.controller_id
      |> add Tlv_type.signature signature
      |> encode)
  in
  let encrypted =
    chacha20_poly1305_encrypt ~key:enc_key ~nonce:"\x00\x00\x00\x00PV-Msg03"
      ~aad:"" sub_tlv
  in
  Tlv.(
    empty |> add Tlv_type.state "\x03"
    |> add Tlv_type.encrypted_data encrypted
    |> encode)

(* Derive session keys from shared secret *)
let derive_session_keys ~pairing ~ip ~port ~shared =
  let enc_key =
    hkdf_sha512 ~salt:"Control-Salt" ~ikm:shared
      ~info:"Control-Write-Encryption-Key" ~length:32
  in
  let dec_key =
    hkdf_sha512 ~salt:"Control-Salt" ~ikm:shared
      ~info:"Control-Read-Encryption-Key" ~length:32
  in
  Log.info (fun f -> f "Pair verify successful, session established");
  Ok
    {
      pairing;
      ip;
      port;
      encrypt_key = enc_key;
      decrypt_key = dec_key;
      encrypt_count = 0L;
      decrypt_count = 0L;
    }

(* Pair Verify M3/M4 - send M3, handle M4, derive session keys *)
let pair_verify_response ~net ~sw ~clock ~ip ~port ~pairing ~kp ~accessory_pk
    ~enc_key ~shared =
  let m3 = build_verify_m3 ~enc_key ~pairing ~kp ~accessory_pk in
  let* m4_body =
    http_post ~net ~clock ~sw ~ip ~port ~path:"/pair-verify"
      ~content_type:"application/pairing+tlv8" ~body:m3
  in
  let m4 = Tlv.decode m4_body in
  match Tlv.get Tlv_type.error m4 with
  | Some e -> err_pair_verify_m4 e
  | None -> derive_session_keys ~pairing ~ip ~port ~shared

(* Verify M2 response: decrypt, check identity and signature *)
let verify_m2 ~pairing ~kp ~enc_key ~enc_data ~accessory_pk =
  let* decrypted =
    chacha20_poly1305_decrypt ~key:enc_key ~nonce:"\x00\x00\x00\x00PV-Msg02"
      ~aad:"" enc_data
  in
  let sub_tlv = Tlv.decode decrypted in
  let accessory_id = Tlv.get_exn Tlv_type.identifier sub_tlv in
  let accessory_sig = Tlv.get_exn Tlv_type.signature sub_tlv in
  if accessory_id <> pairing.accessory_id then
    Error (`Msg "Accessory ID mismatch")
  else
    let verify_data = accessory_pk ^ accessory_id ^ kp.X25519.public in
    if
      not
        (Ed25519.verify ~public:pairing.accessory_ltpk ~signature:accessory_sig
           verify_data)
    then Error (`Msg "Accessory signature verification failed")
    else Ok ()

(* Pair Verify - establish encrypted session *)
let pair_verify ~net ~sw ~clock ~ip ~port ~pairing =
  Log.info (fun f -> f "Starting pair verify with %s:%d" ip port);
  let kp = X25519.generate () in
  let m1 =
    Tlv.(
      empty |> add Tlv_type.state "\x01"
      |> add Tlv_type.public_key kp.public
      |> encode)
  in
  let* m2_body =
    http_post ~net ~clock ~sw ~ip ~port ~path:"/pair-verify"
      ~content_type:"application/pairing+tlv8" ~body:m1
  in
  let m2 = Tlv.decode m2_body in
  match Tlv.get Tlv_type.error m2 with
  | Some e -> err_pair_verify_m2 e
  | None ->
      let accessory_pk = Tlv.get_exn Tlv_type.public_key m2 in
      let enc_data = Tlv.get_exn Tlv_type.encrypted_data m2 in
      let* shared =
        X25519.shared_secret ~secret:kp.secret ~public:accessory_pk
      in
      let enc_key =
        hkdf_sha512 ~salt:"Pair-Verify-Encrypt-Salt" ~ikm:shared
          ~info:"Pair-Verify-Encrypt-Info" ~length:32
      in
      let* () = verify_m2 ~pairing ~kp ~enc_key ~enc_data ~accessory_pk in
      pair_verify_response ~net ~sw ~clock ~ip ~port ~pairing ~kp ~accessory_pk
        ~enc_key ~shared

(* Send encrypted request and read response *)
let request ~net ~sw session req =
  let encrypted = encrypt_frame session req in
  let addr = `Tcp (ipv4_of_string session.ip, session.port) in
  let flow = Eio.Net.connect ~sw net addr in
  Eio.Flow.copy_string encrypted flow;
  let buf = Buffer.create 4096 in
  let rec read () =
    let chunk = Cstruct.create 1024 in
    match Eio.Flow.single_read flow chunk with
    | n ->
        Buffer.add_string buf (Cstruct.to_string (Cstruct.sub chunk 0 n));
        read ()
    | exception End_of_file -> ()
  in
  read ();
  Eio.Flow.close flow;
  decrypt_frame session (Buffer.contents buf)

(* Parse HTTP response body as JSON *)
let parse_json_response decrypted =
  match Re.(exec_opt (compile (str "\r\n\r\n")) decrypted) with
  | None -> Error (`Msg "Invalid response")
  | Some g ->
      let pos = Re.Group.stop g 0 in
      let body = String.sub decrypted pos (String.length decrypted - pos) in
      Result.map_error
        (fun e -> `Msg e)
        (Jsont_bytesrw.decode_string Jsont.json body)

(* Get accessories from a session *)
let accessories ~net ~sw session =
  let path = "/accessories" in
  let req =
    Fmt.str "GET %s HTTP/1.1\r\nHost: %s:%d\r\n\r\n" path session.ip
      session.port
  in
  let* decrypted = request ~net ~sw session req in
  parse_json_response decrypted

(* Characteristic write request codec *)
type char_write = { cw_aid : int; cw_iid : int; cw_value : Jsont.json }

let char_write_codec =
  Jsont.Object.map ~kind:"char_write" (fun aid iid value ->
      { cw_aid = aid; cw_iid = iid; cw_value = value })
  |> Jsont.Object.mem "aid" Jsont.int ~enc:(fun c -> c.cw_aid)
  |> Jsont.Object.mem "iid" Jsont.int ~enc:(fun c -> c.cw_iid)
  |> Jsont.Object.mem "value" Jsont.json ~enc:(fun c -> c.cw_value)
  |> Jsont.Object.finish

type char_write_request = { characteristics : char_write list }

let char_write_request_codec =
  Jsont.Object.map ~kind:"char_write_request" (fun characteristics ->
      { characteristics })
  |> Jsont.Object.mem "characteristics" (Jsont.list char_write_codec)
       ~enc:(fun r -> r.characteristics)
  |> Jsont.Object.finish

(* Write a characteristic *)
let put_characteristic ~net ~sw session ~aid ~iid value =
  let req =
    { characteristics = [ { cw_aid = aid; cw_iid = iid; cw_value = value } ] }
  in
  let body =
    match Jsont_bytesrw.encode_string char_write_request_codec req with
    | Ok s -> s
    | Error _ -> "{}"
  in
  let path = "/characteristics" in
  let req =
    Fmt.str
      "PUT %s HTTP/1.1\r\n\
       Host: %s:%d\r\n\
       Content-Type: application/hap+json\r\n\
       Content-Length: %d\r\n\
       \r\n\
       %s"
      path session.ip session.port (String.length body) body
  in
  let* _decrypted = request ~net ~sw session req in
  Ok ()

(* Read characteristics *)
let characteristics ~net ~sw session ~ids =
  let ids_str =
    String.concat "," (List.map (fun (aid, iid) -> Fmt.str "%d.%d" aid iid) ids)
  in
  let path = Fmt.str "/characteristics?id=%s" ids_str in
  let req =
    Fmt.str "GET %s HTTP/1.1\r\nHost: %s:%d\r\n\r\n" path session.ip
      session.port
  in
  let* decrypted = request ~net ~sw session req in
  parse_json_response decrypted

(* Pairing storage directory *)
let pairings_dir = ".hap/pairings"

let ensure_pairings_dir ~fs =
  let hap_path = Eio.Path.(fs / ".hap") in
  let pairings_path = Eio.Path.(fs / pairings_dir) in
  (try Eio.Path.mkdir ~perm:0o700 hap_path with Eio.Exn.Io _ -> ());
  try Eio.Path.mkdir ~perm:0o700 pairings_path with Eio.Exn.Io _ -> ()

(* Sanitize device_id for filename (replace colons with dashes) *)
let sanitize_id id = String.map (fun c -> if c = ':' then '-' else c) id

let pairing_path_for_id device_id =
  Fmt.str "%s/hap-%s.json" pairings_dir (sanitize_id device_id)

(** Jsont codec for pairing storage *)
module Pairing_json = struct
  type stored = {
    accessory_id : string;
    accessory_ltpk : string; (* Base64-encoded *)
    controller_id : string;
    controller_ltsk : string; (* Base64-encoded *)
    controller_ltpk : string; (* Base64-encoded *)
  }

  let stored =
    Jsont.Object.map ~kind:"hap.pairing"
      (fun
        accessory_id
        accessory_ltpk
        controller_id
        controller_ltsk
        controller_ltpk
      ->
        {
          accessory_id;
          accessory_ltpk;
          controller_id;
          controller_ltsk;
          controller_ltpk;
        })
    |> Jsont.Object.mem "accessory_id" Jsont.string ~enc:(fun p ->
        p.accessory_id)
    |> Jsont.Object.mem "accessory_ltpk" Jsont.string ~enc:(fun p ->
        p.accessory_ltpk)
    |> Jsont.Object.mem "controller_id" Jsont.string ~enc:(fun p ->
        p.controller_id)
    |> Jsont.Object.mem "controller_ltsk" Jsont.string ~enc:(fun p ->
        p.controller_ltsk)
    |> Jsont.Object.mem "controller_ltpk" Jsont.string ~enc:(fun p ->
        p.controller_ltpk)
    |> Jsont.Object.finish

  let of_pairing (p : pairing) : stored =
    {
      accessory_id = p.accessory_id;
      accessory_ltpk = Base64.encode_string p.accessory_ltpk;
      controller_id = p.controller_id;
      controller_ltsk = Base64.encode_string p.controller_ltsk;
      controller_ltpk = Base64.encode_string p.controller_ltpk;
    }

  let to_pairing (s : stored) : pairing =
    {
      accessory_id = s.accessory_id;
      accessory_ltpk = Base64.decode_exn s.accessory_ltpk;
      controller_id = s.controller_id;
      controller_ltsk = Base64.decode_exn s.controller_ltsk;
      controller_ltpk = Base64.decode_exn s.controller_ltpk;
    }
end

(* Save/load pairing to file *)
let save_pairing ~fs ~path (pairing : pairing) =
  let stored = Pairing_json.of_pairing pairing in
  match
    Jsont_bytesrw.encode_string ~format:Jsont.Indent Pairing_json.stored stored
  with
  | Ok json ->
      Eio.Path.save ~create:(`Or_truncate 0o600) Eio.Path.(fs / path) json
  | Error _ -> ()

let load_pairing ~fs ~path =
  let full_path = Eio.Path.(fs / path) in
  if not (Eio.Path.is_file full_path) then None
  else begin
    try
      let content = Eio.Path.load full_path in
      match Jsont_bytesrw.decode_string Pairing_json.stored content with
      | Ok stored -> Some (Pairing_json.to_pairing stored)
      | Error _ -> None
    with Eio.Io _ -> None
  end

(* Save pairing by device_id *)
let save_pairing_by_id ~fs pairing =
  ensure_pairings_dir ~fs;
  let path = pairing_path_for_id pairing.accessory_id in
  save_pairing ~fs ~path pairing;
  path

(* Find pairing for a device by its HAP device_id *)
let pairing_by_id ~fs device_id =
  let path = pairing_path_for_id device_id in
  load_pairing ~fs ~path

(* HAP category codes *)
let category_name = function
  | 1 -> "Other"
  | 2 -> "Bridge"
  | 3 -> "Fan"
  | 4 -> "Garage Door Opener"
  | 5 -> "Lightbulb"
  | 6 -> "Door Lock"
  | 7 -> "Outlet"
  | 8 -> "Switch"
  | 9 -> "Thermostat"
  | 10 -> "Sensor"
  | 11 -> "Security System"
  | 12 -> "Door"
  | 13 -> "Window"
  | 14 -> "Window Covering"
  | 15 -> "Programmable Switch"
  | 16 -> "Range Extender"
  | 17 -> "IP Camera"
  | 18 -> "Video Doorbell"
  | 19 -> "Air Purifier"
  | 20 -> "Heater"
  | 21 -> "Air Conditioner"
  | 22 -> "Humidifier"
  | 23 -> "Dehumidifier"
  | 24 -> "Apple TV"
  | 25 -> "HomePod"
  | 26 -> "Speaker"
  | 27 -> "AirPort"
  | 28 -> "Sprinkler"
  | 29 -> "Faucet"
  | 30 -> "Shower Head"
  | 31 -> "Television"
  | 32 -> "Target Controller"
  | 33 -> "WiFi Router"
  | 34 -> "Audio Receiver"
  | 35 -> "TV Set Top Box"
  | 36 -> "TV Streaming Stick"
  | _ -> "Unknown"

(* Parse HAP TXT record *)
let parse_hap_txt txt =
  (* TXT record contains key=value pairs *)
  let pairs = String.split_on_char ' ' txt in
  let find key =
    List.find_map
      (fun p ->
        match String.split_on_char '=' p with
        | [ k; v ] when k = key -> Some v
        | _ -> None)
      pairs
  in
  let find_int key = Option.bind (find key) int_of_string_opt in
  let device_id = find "id" in
  let model = find "md" in
  let config_num = find_int "c#" in
  let state_num = find_int "s#" in
  let category = find_int "ci" in
  let paired = find_int "sf" = Some 0 in
  (device_id, model, config_num, state_num, category, paired)

(* Build accessory_info from mDNS instance, SRV, TXT, and address records *)
let build_device_info (r : Mdns.response) instance =
  match
    List.find_opt (fun (n, _, _) -> Domain_name.equal n instance) r.srvs
  with
  | None -> None
  | Some (_, port, target) ->
      let txt =
        List.find_map
          (fun (n, t) -> if Domain_name.equal n instance then Some t else None)
          r.txts
        |> Option.value ~default:[] |> String.concat " "
      in
      let ip =
        List.find_map
          (fun (n, ip) ->
            if Domain_name.equal n target then Some (Ipaddr.V4.to_string ip)
            else None)
          r.addrs
        |> Option.value ~default:(Domain_name.to_string target)
      in
      let device_id, model, config_num, state_num, category, paired =
        parse_hap_txt txt
      in
      let name =
        match Domain_name.get_label instance 0 with
        | Ok label -> label
        | Error _ -> Domain_name.to_string instance
      in
      Some
        {
          name;
          device_id = Option.value ~default:"" device_id;
          ip;
          port;
          model;
          config_num = Option.value ~default:0 config_num;
          state_num = Option.value ~default:0 state_num;
          category = Option.value ~default:0 category;
          paired;
        }

(* Discover HAP devices using mDNS *)
let discover ~sw ~net ~clock ?(timeout = 3.0) () =
  let service_name = Domain_name.of_string_exn "_hap._tcp.local" in
  let r = Mdns.merge (Mdns.query ~sw ~net ~clock ~timeout service_name) in
  (* Get unique service instances from PTR records *)
  let instances =
    List.filter_map
      (fun (service, instance) ->
        if Domain_name.equal service service_name then Some instance else None)
      r.ptrs
    |> List.sort_uniq Domain_name.compare
  in
  (* Build device info for each instance *)
  List.filter_map (build_device_info r) instances

(* Find pairing for an IP by discovering the device first *)
let pairing_for_ip ~sw ~net ~clock ~fs ip =
  let devices = discover ~sw ~net ~clock ~timeout:2.0 () in
  match List.find_opt (fun (d : accessory_info) -> d.ip = ip) devices with
  | None -> None
  | Some info ->
      if info.device_id = "" then None else pairing_by_id ~fs info.device_id

(* Get accessory info for an IP *)
let accessory_info ~sw ~net ~clock ip =
  let devices = discover ~sw ~net ~clock ~timeout:2.0 () in
  List.find_opt (fun (d : accessory_info) -> d.ip = ip) devices

(* Pretty print accessory info *)
let pp_accessory_info ppf info =
  let cat = category_name info.category in
  let status = if info.paired then "paired" else "unpaired" in
  Fmt.pf ppf "@[<v 0>%s@," info.name;
  Fmt.pf ppf "  Type:      %s@," cat;
  Fmt.pf ppf "  Device ID: %s@," info.device_id;
  Fmt.pf ppf "  Address:   %s:%d@," info.ip info.port;
  Option.iter (fun m -> Fmt.pf ppf "  Model:     %s@," m) info.model;
  Fmt.pf ppf "  Status:    %s@," status;
  Fmt.pf ppf "  Config:    #%d, State: #%d@]" info.config_num info.state_num

(** {1 HAP JSON Codecs} *)

module Hap_json = struct
  type characteristic = { iid : int; type_ : string; value : Jsont.json option }
  (** HAP characteristic *)

  let characteristic =
    Jsont.Object.map ~kind:"hap.characteristic" (fun iid type_ value ->
        { iid; type_; value })
    |> Jsont.Object.mem "iid" Jsont.int ~enc:(fun c -> c.iid)
    |> Jsont.Object.mem "type" Jsont.string ~enc:(fun c -> c.type_)
    |> Jsont.Object.opt_mem "value" Jsont.json ~enc:(fun c -> c.value)
    |> Jsont.Object.finish

  type service = {
    iid : int;
    type_ : string;
    characteristics : characteristic list;
  }
  (** HAP service *)

  let service =
    Jsont.Object.map ~kind:"hap.service" (fun iid type_ characteristics ->
        { iid; type_; characteristics })
    |> Jsont.Object.mem "iid" Jsont.int ~enc:(fun s -> s.iid)
    |> Jsont.Object.mem "type" Jsont.string ~enc:(fun s -> s.type_)
    |> Jsont.Object.mem "characteristics" (Jsont.list characteristic)
         ~enc:(fun s -> s.characteristics)
    |> Jsont.Object.finish

  type accessory = { aid : int; services : service list }
  (** HAP accessory *)

  let accessory =
    Jsont.Object.map ~kind:"hap.accessory" (fun aid services ->
        { aid; services })
    |> Jsont.Object.mem "aid" Jsont.int ~enc:(fun a -> a.aid)
    |> Jsont.Object.mem "services" (Jsont.list service) ~enc:(fun a ->
        a.services)
    |> Jsont.Object.finish

  type accessories_response = { accessories : accessory list }
  (** HAP accessories response *)

  let accessories_response =
    Jsont.Object.map ~kind:"hap.accessories_response" (fun accessories ->
        { accessories })
    |> Jsont.Object.mem "accessories" (Jsont.list accessory) ~enc:(fun r ->
        r.accessories)
    |> Jsont.Object.finish

  type char_value = { aid : int; iid : int; value : Jsont.json option }
  (** HAP characteristics value *)

  let char_value =
    Jsont.Object.map ~kind:"hap.char_value" (fun aid iid value ->
        { aid; iid; value })
    |> Jsont.Object.mem "aid" Jsont.int ~enc:(fun c -> c.aid)
    |> Jsont.Object.mem "iid" Jsont.int ~enc:(fun c -> c.iid)
    |> Jsont.Object.opt_mem "value" Jsont.json ~enc:(fun c -> c.value)
    |> Jsont.Object.finish

  type characteristics_response = { characteristics : char_value list }

  let characteristics_response =
    Jsont.Object.map ~kind:"hap.characteristics_response"
      (fun characteristics -> { characteristics })
    |> Jsont.Object.mem "characteristics" (Jsont.list char_value) ~enc:(fun r ->
        r.characteristics)
    |> Jsont.Object.finish
end

(** {1 High-level control} *)

(* HAP characteristic type UUIDs (short form) *)
module Char_type = struct
  let on = "25" (* 00000025-0000-1000-8000-0026BB765291 *)
end

(* Decode Jsont.json via codec *)
let decode codec json =
  match Jsont_bytesrw.encode_string Jsont.json json with
  | Error e -> Error e
  | Ok str -> (
      match Jsont_bytesrw.decode_string codec str with
      | Ok v -> Ok v
      | Error e -> Error e)

(* Find the On characteristic IID from accessories JSON *)
let on_characteristic_iid json =
  match decode Hap_json.accessories_response json with
  | Error _ -> None
  | Ok resp ->
      List.find_map
        (fun (acc : Hap_json.accessory) ->
          List.find_map
            (fun (svc : Hap_json.service) ->
              List.find_map
                (fun (chr : Hap_json.characteristic) ->
                  if String.lowercase_ascii chr.type_ = Char_type.on then
                    Some (acc.aid, chr.iid)
                  else None)
                svc.characteristics)
            acc.services)
        resp.accessories

(* Control an accessory by IP - establishes session, finds characteristic, sets value *)
let control_outlet ~net ~sw ~clock ~fs ~ip ~value =
  (* 1. Discover to get device info *)
  let devices = discover ~sw ~net ~clock ~timeout:2.0 () in
  match List.find_opt (fun (d : accessory_info) -> d.ip = ip) devices with
  | None -> Error (`Msg "Device not found via HAP discovery")
  | Some info -> (
      if info.device_id = "" then Error (`Msg "Device has no device_id")
      else
        (* 2. Find pairing *)
        match pairing_by_id ~fs info.device_id with
        | None ->
            Error (`Msg "No pairing found for device - run 'plug pair' first")
        | Some pairing -> (
            (* 3. Establish session *)
            let* session =
              pair_verify ~net ~sw ~clock ~ip ~port:info.port ~pairing
            in
            (* 4. Get accessories to find On characteristic *)
            let* accessories_json = accessories ~net ~sw session in
            match on_characteristic_iid accessories_json with
            | None -> Error (`Msg "Could not find On characteristic")
            | Some (aid, iid) ->
                (* 5. Set value *)
                put_characteristic ~net ~sw session ~aid ~iid
                  (Jsont.Bool (value, Jsont.Meta.none))))

let turn_on_outlet ~net ~sw ~clock ~fs ip =
  control_outlet ~net ~sw ~clock ~fs ~ip ~value:true

let turn_off_outlet ~net ~sw ~clock ~fs ip =
  control_outlet ~net ~sw ~clock ~fs ~ip ~value:false

(* Extract bool value from characteristics response *)
let bool_value json =
  match decode Hap_json.characteristics_response json with
  | Error _ -> None
  | Ok (resp : Hap_json.characteristics_response) -> (
      match resp.characteristics with
      | [ (c : Hap_json.char_value) ] -> (
          match c.value with Some (Jsont.Bool (b, _)) -> Some b | _ -> None)
      | _ -> None)

let toggle_outlet ~net ~sw ~clock ~fs ip =
  (* For toggle, we need to read current state first *)
  let devices = discover ~sw ~net ~clock ~timeout:2.0 () in
  match List.find_opt (fun (d : accessory_info) -> d.ip = ip) devices with
  | None -> Error (`Msg "Device not found via HAP discovery")
  | Some info -> (
      if info.device_id = "" then Error (`Msg "Device has no device_id")
      else
        match pairing_by_id ~fs info.device_id with
        | None -> Error (`Msg "No pairing found for device")
        | Some pairing -> (
            let* session =
              pair_verify ~net ~sw ~clock ~ip ~port:info.port ~pairing
            in
            let* accessories_json = accessories ~net ~sw session in
            match on_characteristic_iid accessories_json with
            | None -> Error (`Msg "Could not find On characteristic")
            | Some (aid, iid) -> (
                let* chars_json =
                  characteristics ~net ~sw session ~ids:[ (aid, iid) ]
                in
                match bool_value chars_json with
                | None -> Error (`Msg "Could not read current state")
                | Some v ->
                    put_characteristic ~net ~sw session ~aid ~iid
                      (Jsont.Bool (not v, Jsont.Meta.none)))))