open Cohttp
open Cohttp_lwt_unix
open Lwt
open Yojson.Basic.Util

(* Read from stdin *)
let read_stdin_to_string () =
  let rec read_lines acc =
    try
      let line = input_line stdin in
      read_lines (acc ^ line ^ "\n")
    with End_of_file -> acc
  in
  read_lines ""

(* Get env variable *)
let get_env api () =
  try Sys.getenv api
  with Not_found -> failwith api ^ " environment variable not found"

(* Get api keys *)
let ant_api_key = get_env "ANTHROPIC_API_KEY" ()
let memos_api_key = get_env "MEMOS_API_KEY" ()

(* Make a request to Anthropic *)
let get_from_claude prompt () =
  let model = "claude-3-haiku-20240307" in
  let max_tokens = 1024 in
  let ant_uri = Uri.of_string "https://api.anthropic.com/v1/messages" in
  let headers =
    Header.init () |> fun h ->
    Header.add h "x-api-key" ant_api_key |> fun h ->
    Header.add h "anthropic-version" "2023-06-01" |> fun h ->
    Header.add h "content-type" "application/json"
  in
  let body =
    `Assoc
      [
        ("model", `String model);
        ("max_tokens", `Int max_tokens);
        ( "messages",
          `List
            [ `Assoc [ ("role", `String "user"); ("content", `String prompt) ] ]
        );
      ]
    |> Yojson.Basic.to_string
  in
  Client.post ~headers ~body:(Cohttp_lwt.Body.of_string body) ant_uri
  >>= fun (resp, body) ->
  let status = resp |> Response.status in
  if Code.is_success (Code.code_of_status status) then
    body |> Cohttp_lwt.Body.to_string >|= fun body_str ->
    let json = Yojson.Basic.from_string body_str in
    let content = json |> member "content" in
    index 0 content |> member "text" |> to_string
  else Lwt.fail (Failure "Error fetching data from Anthropic")

(* Post data to Memos *)
let post_to_memos prompt text =
  let memo = prompt ^ "\n\n" ^ text ^ "\n\n#claude" in
  let body =
    `Assoc [ ("content", `String memo); ("visibility", `String "PRIVATE") ]
    |> Yojson.Basic.to_string
  in
  let api_key = "Bearer " ^ memos_api_key in
  let headers =
    Header.init () |> fun h ->
    Header.add h "Authorization" api_key |> fun h ->
    Header.add h "Content-Type" "application/json"
  in
  let uri = Uri.of_string "https://cdecary-memos.fly.dev/api/v2/memos" in
  Client.post ~headers ~body:(`String body) uri >>= fun (resp, _) ->
  let status = resp |> Response.status in
  if Code.is_success (Code.code_of_status status) then Lwt.return_unit
  else Lwt.fail (Failure "Error posting data to Memos")

(* Main function that orchestrates the process *)
let main () =
  let prompt = read_stdin_to_string () in
  Printf.printf "(* %s *)\n\n" prompt;
  get_from_claude prompt () >>= fun text ->
  Printf.printf "(* %s *)\n\n" text;
  post_to_memos prompt text >>= fun () -> Lwt.return_unit

let () = Lwt_main.run (main ())