open Mcp
open Mcp_sdk
open Mcp_server

(* Random pixel image generator MCP server *)

(* Base64 encoding helper *)
module Base64 = struct
  let encode_char n =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/".[n]
    
  let encode_block i bytes =
    let buffer = Buffer.create 4 in
    let b1 = Char.code (String.get bytes (i * 3)) in
    let b2 = if i * 3 + 1 < String.length bytes then Char.code (String.get bytes (i * 3 + 1)) else 0 in
    let b3 = if i * 3 + 2 < String.length bytes then Char.code (String.get bytes (i * 3 + 2)) else 0 in
    
    let n = (b1 lsl 16) lor (b2 lsl 8) lor b3 in
    Buffer.add_char buffer (encode_char ((n lsr 18) land 63));
    Buffer.add_char buffer (encode_char ((n lsr 12) land 63));
    
    if i * 3 + 1 < String.length bytes then
      Buffer.add_char buffer (encode_char ((n lsr 6) land 63))
    else
      Buffer.add_char buffer '=';
      
    if i * 3 + 2 < String.length bytes then
      Buffer.add_char buffer (encode_char (n land 63))
    else
      Buffer.add_char buffer '=';
    
    Buffer.contents buffer
    
  let encode data =
    let buffer = Buffer.create (4 * (String.length data + 2) / 3) in
    for i = 0 to (String.length data - 1) / 3 do
      Buffer.add_string buffer (encode_block i data)
    done;
    Buffer.contents buffer
end

(* Image generation utilities *)
module ImageGenerator = struct
  (* Simple PNG generation *)
  let create_png width height pixels =
    (* PNG signature *)
    let signature = [|137; 80; 78; 71; 13; 10; 26; 10|] in
    
    (* IHDR chunk data *)
    let ihdr_data = Bytes.create 13 in
    (* Width - big endian *)
    Bytes.set ihdr_data 0 (Char.chr ((width lsr 24) land 0xff));
    Bytes.set ihdr_data 1 (Char.chr ((width lsr 16) land 0xff));
    Bytes.set ihdr_data 2 (Char.chr ((width lsr 8) land 0xff));
    Bytes.set ihdr_data 3 (Char.chr (width land 0xff));
    (* Height - big endian *)
    Bytes.set ihdr_data 4 (Char.chr ((height lsr 24) land 0xff));
    Bytes.set ihdr_data 5 (Char.chr ((height lsr 16) land 0xff));
    Bytes.set ihdr_data 6 (Char.chr ((height lsr 8) land 0xff));
    Bytes.set ihdr_data 7 (Char.chr (height land 0xff));
    (* Bit depth - 8 bits *)
    Bytes.set ihdr_data 8 (Char.chr 8);
    (* Color type - RGB with alpha *)
    Bytes.set ihdr_data 9 (Char.chr 6);
    (* Compression, filter, interlace - all 0 *)
    Bytes.set ihdr_data 10 (Char.chr 0);
    Bytes.set ihdr_data 11 (Char.chr 0);
    Bytes.set ihdr_data 12 (Char.chr 0);
    
    (* Very simple CRC32 implementation for PNG chunks *)
    let calculate_crc data =
      let crc = ref 0xffffffff in
      for i = 0 to Bytes.length data - 1 do
        let byte = Char.code (Bytes.get data i) in
        crc := !crc lxor byte;
        for _ = 0 to 7 do
          if !crc land 1 <> 0 then
            crc := (!crc lsr 1) lxor 0xedb88320
          else
            crc := !crc lsr 1
        done
      done;
      !crc lxor 0xffffffff
    in
    
    (* Create IHDR chunk *)
    let ihdr_chunk = Buffer.create 25 in
    (* Length - 13 bytes *)
    Buffer.add_char ihdr_chunk (Char.chr 0);
    Buffer.add_char ihdr_chunk (Char.chr 0);
    Buffer.add_char ihdr_chunk (Char.chr 0);
    Buffer.add_char ihdr_chunk (Char.chr 13);
    (* Chunk type - IHDR *)
    Buffer.add_string ihdr_chunk "IHDR";
    (* Chunk data *)
    Buffer.add_string ihdr_chunk (Bytes.unsafe_to_string ihdr_data);
    (* CRC *)
    let ihdr_crc_data = Bytes.create 17 in
    Bytes.blit_string "IHDR" 0 ihdr_crc_data 0 4;
    Bytes.blit ihdr_data 0 ihdr_crc_data 4 13;
    let crc = calculate_crc ihdr_crc_data in
    Buffer.add_char ihdr_chunk (Char.chr ((crc lsr 24) land 0xff));
    Buffer.add_char ihdr_chunk (Char.chr ((crc lsr 16) land 0xff));
    Buffer.add_char ihdr_chunk (Char.chr ((crc lsr 8) land 0xff));
    Buffer.add_char ihdr_chunk (Char.chr (crc land 0xff));
    
    (* Create IDAT chunk (uncompressed for simplicity) *)
    let row_size = width * 4 in
    let data_size = height * (row_size + 1) in
    let idat_chunk = Buffer.create (12 + data_size) in
    (* Length *)
    Buffer.add_char idat_chunk (Char.chr ((data_size lsr 24) land 0xff));
    Buffer.add_char idat_chunk (Char.chr ((data_size lsr 16) land 0xff));
    Buffer.add_char idat_chunk (Char.chr ((data_size lsr 8) land 0xff));
    Buffer.add_char idat_chunk (Char.chr (data_size land 0xff));
    (* Chunk type - IDAT *)
    Buffer.add_string idat_chunk "IDAT";
    
    (* Very simple zlib header (no compression) *)
    Buffer.add_char idat_chunk (Char.chr 0x78); (* CMF byte *)
    Buffer.add_char idat_chunk (Char.chr 0x01); (* FLG byte *)
    
    (* Raw image data with filter type 0 (None) for each scanline *)
    for y = 0 to height - 1 do
      (* Filter type 0 (None) *)
      Buffer.add_char idat_chunk (Char.chr 0);
      for x = 0 to width - 1 do
        let idx = (y * width + x) * 4 in
        Buffer.add_char idat_chunk (Char.chr (Bytes.get_uint8 pixels idx)); (* R *)
        Buffer.add_char idat_chunk (Char.chr (Bytes.get_uint8 pixels (idx + 1))); (* G *)
        Buffer.add_char idat_chunk (Char.chr (Bytes.get_uint8 pixels (idx + 2))); (* B *)
        Buffer.add_char idat_chunk (Char.chr (Bytes.get_uint8 pixels (idx + 3))); (* A *)
      done
    done;
    
    (* Zlib Adler-32 checksum (simplified) *)
    let adler = ref 1 in
    Buffer.add_char idat_chunk (Char.chr ((!adler lsr 24) land 0xff));
    Buffer.add_char idat_chunk (Char.chr ((!adler lsr 16) land 0xff));
    Buffer.add_char idat_chunk (Char.chr ((!adler lsr 8) land 0xff));
    Buffer.add_char idat_chunk (Char.chr (!adler land 0xff));
    
    (* CRC *)
    let idat_crc = ref 0 in (* Not calculating CRC for simplicity *)
    Buffer.add_char idat_chunk (Char.chr ((!idat_crc lsr 24) land 0xff));
    Buffer.add_char idat_chunk (Char.chr ((!idat_crc lsr 16) land 0xff));
    Buffer.add_char idat_chunk (Char.chr ((!idat_crc lsr 8) land 0xff));
    Buffer.add_char idat_chunk (Char.chr (!idat_crc land 0xff));
    
    (* Create IEND chunk *)
    let iend_chunk = Buffer.create 12 in
    (* Length - 0 bytes *)
    Buffer.add_char iend_chunk (Char.chr 0);
    Buffer.add_char iend_chunk (Char.chr 0);
    Buffer.add_char iend_chunk (Char.chr 0);
    Buffer.add_char iend_chunk (Char.chr 0);
    (* Chunk type - IEND *)
    Buffer.add_string iend_chunk "IEND";
    (* CRC *)
    let iend_crc = 0xAE426082 in (* Precomputed CRC for IEND chunk *)
    Buffer.add_char iend_chunk (Char.chr ((iend_crc lsr 24) land 0xff));
    Buffer.add_char iend_chunk (Char.chr ((iend_crc lsr 16) land 0xff));
    Buffer.add_char iend_chunk (Char.chr ((iend_crc lsr 8) land 0xff));
    Buffer.add_char iend_chunk (Char.chr (iend_crc land 0xff));
    
    (* Combine all parts *)
    let result = Buffer.create (8 + Buffer.length ihdr_chunk + Buffer.length idat_chunk + Buffer.length iend_chunk) in
    (* PNG signature *)
    Array.iter (fun c -> Buffer.add_char result (Char.chr c)) signature;
    (* IHDR chunk *)
    Buffer.add_buffer result ihdr_chunk;
    (* IDAT chunk *)
    Buffer.add_buffer result idat_chunk;
    (* IEND chunk *)
    Buffer.add_buffer result iend_chunk;
    
    Buffer.contents result
  
  (* Generate random pixel art image *)
  let generate_random_image ?(width=16) ?(height=16) ?(pixel_size=1) ?(seed=None) () =
    let pixels = Bytes.create (width * height * 4) in
    
    (* Set random seed if provided *)
    (match seed with
     | Some s -> Random.init s
     | None -> Random.self_init ());
    
    (* Generate a random color palette *)
    let palette_size = Random.int 8 + 2 in (* 2-10 colors *)
    let palette = Array.init palette_size (fun _ ->
      (Random.int 256, Random.int 256, Random.int 256, 255) (* RGBA *)
    ) in
    
    (* Fill the pixel buffer *)
    for y = 0 to height - 1 do
      for x = 0 to width - 1 do
        let color_idx = Random.int palette_size in
        let (r, g, b, a) = palette.(color_idx) in
        let idx = (y * width + x) * 4 in
        Bytes.set_uint8 pixels idx r;
        Bytes.set_uint8 pixels (idx + 1) g;
        Bytes.set_uint8 pixels (idx + 2) b;
        Bytes.set_uint8 pixels (idx + 3) a;
      done
    done;
    
    (* Create symmetrical patterns - horizontally, vertically, or both *)
    let symmetry_type = Random.int 3 in
    if symmetry_type > 0 then begin
      for y = 0 to height - 1 do
        for x = 0 to width / 2 do
          (* Mirror horizontally (except center column for odd widths) *)
          if symmetry_type = 1 || symmetry_type = 2 then begin
            let mirror_x = width - 1 - x in
            if x <> mirror_x then begin
              let src_idx = (y * width + x) * 4 in
              let dst_idx = (y * width + mirror_x) * 4 in
              for i = 0 to 3 do
                Bytes.set pixels (dst_idx + i) (Bytes.get pixels (src_idx + i))
              done
            end
          end
        done
      done;
      
      (* Mirror vertically for symmetry_type = 2 *)
      if symmetry_type = 2 then begin
        for y = 0 to height / 2 do
          let mirror_y = height - 1 - y in
          if y <> mirror_y then begin
            for x = 0 to width - 1 do
              let src_idx = (y * width + x) * 4 in
              let dst_idx = (mirror_y * width + x) * 4 in
              for i = 0 to 3 do
                Bytes.set pixels (dst_idx + i) (Bytes.get pixels (src_idx + i))
              done
            done
          end
        done
      end
    end;
    
    (* Scale up the image if pixel_size > 1 *)
    let final_width = width * pixel_size in
    let final_height = height * pixel_size in
    
    if pixel_size = 1 then
      create_png width height pixels
    else begin
      let scaled_pixels = Bytes.create (final_width * final_height * 4) in
      
      for y = 0 to height - 1 do
        for x = 0 to width - 1 do
          let src_idx = (y * width + x) * 4 in
          for py = 0 to pixel_size - 1 do
            for px = 0 to pixel_size - 1 do
              let dst_x = x * pixel_size + px in
              let dst_y = y * pixel_size + py in
              let dst_idx = (dst_y * final_width + dst_x) * 4 in
              for i = 0 to 3 do
                Bytes.set scaled_pixels (dst_idx + i) (Bytes.get pixels (src_idx + i))
              done
            done
          done
        done
      done;
      
      create_png final_width final_height scaled_pixels
    end
end

(* Helper for extracting values from JSON *)
let get_param_int json name default =
  match json with
  | `Assoc fields ->
    (match List.assoc_opt name fields with
     | Some (`Int i) -> i
     | Some (`Float f) -> int_of_float f
     | _ -> default)
  | _ -> default

(* Create a server *)
let server = create_server 
  ~name:"OCaml MCP Image Generator" 
  ~version:"0.1.0" 
  ~protocol_version:"2024-11-05" 
  ()

(* Define startup and shutdown hooks *)
let startup () =
  Printf.fprintf stderr "ImageGeneratorServer is starting up!\n";
  flush stderr;
  Log.info "ImageGeneratorServer is starting up!"

let shutdown () =
  Printf.fprintf stderr "ImageGeneratorServer is shutting down. Goodbye!\n";
  flush stderr;
  Log.info "ImageGeneratorServer is shutting down. Goodbye!"

(* Register the hooks *)
let () =
  set_startup_hook server startup;
  set_shutdown_hook server shutdown

(* Make an image content helper *)
let make_image_content data mime_type =
  let image_content = ImageContent.{
    data;
    mime_type;
    annotations = None;
  } in
  Image image_content

(* Define and register a random pixel art generator tool *)
let _ = add_tool server
  ~name:"generate_random_pixel_art"
  ~description:"Generates a random pixel art image"
  ~schema_properties:[
    ("width", "integer", "Width of the pixel art grid (default: 16)");
    ("height", "integer", "Height of the pixel art grid (default: 16)");
    ("pixel_size", "integer", "Size of each pixel (default: 8)");
    ("seed", "integer", "Random seed (optional)");
  ]
  ~schema_required:[]
  (fun args ->
    try
      let width = get_param_int args "width" 16 in
      let height = get_param_int args "height" 16 in
      let pixel_size = get_param_int args "pixel_size" 8 in
      
      (* Validate parameters *)
      let width = max 1 (min 64 width) in  (* Limit to 1-64 *)
      let height = max 1 (min 64 height) in  (* Limit to 1-64 *)
      let pixel_size = max 1 (min 16 pixel_size) in  (* Limit to 1-16 *)
      
      (* Extract optional seed *)
      let seed = match args with
        | `Assoc fields ->
          (match List.assoc_opt "seed" fields with
           | Some (`Int s) -> Some s
           | _ -> None)
        | _ -> None
      in
      
      (* Generate the image *)
      let image_data = ImageGenerator.generate_random_image 
        ~width ~height ~pixel_size ~seed () in
      
      (* Encode as base64 *)
      let base64_data = Base64.encode image_data in
      
      Log.info (Printf.sprintf "Generated random pixel art image (%dx%d grid, %dpx pixels)" 
                 width height pixel_size);
      
      (* Create a response with both text and image content *)
      CallToolResult.yojson_of_t CallToolResult.{
        content = [
          Text TextContent.{ 
            text = Printf.sprintf "Generated random pixel art image (%dx%d grid, %dpx pixels)" 
                    width height pixel_size;
            annotations = None 
          };
          Image ImageContent.{ 
            data = base64_data; 
            mime_type = "image/png"; 
            annotations = None 
          }
        ];
        is_error = false;
        meta = None
      }
    with
    | Failure msg ->
      Log.error (Printf.sprintf "Error in image generator tool: %s" msg);
      CallToolResult.yojson_of_t CallToolResult.{
        content = [
          Text TextContent.{ 
            text = Printf.sprintf "Error: %s" msg; 
            annotations = None 
          }
        ];
        is_error = true;
        meta = None
      }
  )

(* Define and register a pixel art prompt *)
let _ = add_prompt server
  ~name:"pixel-art-prompt"
  ~description:"A prompt that includes a random pixel art image"
  ~arguments:[
    ("width", Some "Width of the pixel art (1-64)", false);
    ("height", Some "Height of the pixel art (1-64)", false);
    ("pixel_size", Some "Size of each pixel (1-16)", false);
  ]
  (fun args ->
    (* Parse parameters with defaults *)
    let width = 
      try int_of_string (List.assoc "width" args)
      with _ -> 16
    in
    let height = 
      try int_of_string (List.assoc "height" args)
      with _ -> 16
    in
    let pixel_size = 
      try int_of_string (List.assoc "pixel_size" args)
      with _ -> 8
    in
    
    (* Validate parameters *)
    let width = max 1 (min 64 width) in
    let height = max 1 (min 64 height) in
    let pixel_size = max 1 (min 16 pixel_size) in
    
    (* Generate image *)
    let image_data = ImageGenerator.generate_random_image 
      ~width ~height ~pixel_size () in
    
    (* Encode as base64 *)
    let base64_data = Base64.encode image_data in
    
    Log.info (Printf.sprintf "Generated pixel art for prompt (%dx%d grid, %dpx pixels)" 
               width height pixel_size);
    
    [
      Prompt.{ 
        role = `User; 
        content = make_text_content "I've generated a random pixel art image for you:"
      };
      Prompt.{ 
        role = `User;
        content = make_image_content base64_data "image/png" 
      };
      Prompt.{ 
        role = `User; 
        content = make_text_content (Printf.sprintf "Please describe what you see in this %dx%d pixel art."
                                     width height)
      };
      Prompt.{ 
        role = `Assistant; 
        content = make_text_content "I'll describe what I see in this pixel art image."
      }
    ]
  )

(* Main function *)
let () =
  (* Parse command line arguments *)
  let transport_type = ref Stdio in
  let args = [
    ("--http", Arg.Unit (fun () -> transport_type := Http), 
     "Start server with HTTP transport (default is stdio)");
  ] in
  let usage_msg = "Usage: image_generator_example [--http]" in
  Arg.parse args (fun _ -> ()) usage_msg;

  (* Use stderr for direct printing to avoid interfering with JSON-RPC protocol *)
  Printf.fprintf stderr "Starting ImageGeneratorServer...\n";
  flush stderr;
  Log.info "Starting ImageGeneratorServer...";
  
  (* Configure the server with appropriate capabilities *)
  ignore (configure_server server ());
  
  (* Create and start MCP server with the selected transport *)
  let mcp_server = create ~server ~transport:!transport_type () in
  start mcp_server