open Ir

let read_lines path =
  let ic = open_in path in
  let lines = ref [] in
  (try
     while true do lines := input_line ic :: !lines done
   with End_of_file -> close_in ic);
  List.rev !lines

let parse_func (lines : string list) : func =
  let blocks = ref [] in
  let current_id = ref 0 in
  let current_insts = ref [] in
  let finish_block () =
    if !current_id <> 0 then
      blocks := mk_block !current_id (List.rev !current_insts) :: !blocks;
    current_id := 0;
    current_insts := []
  in
  let parse_instr s =
    (* Tokenize: split on whitespace and standalone commas, but keep a
       parenthesized group (e.g. "(10,2)") as a single token. This lets
       "load 200, 100" and "phi 7 (10,2) (20,3)" both parse cleanly. *)
    let tokens = ref [] in
    let buf = Buffer.create 16 in
    let depth = ref 0 in
    let flush () =
      if Buffer.length buf > 0 then begin
        tokens := String.trim (Buffer.contents buf) :: !tokens;
        Buffer.clear buf
      end
    in
    String.iter (fun c ->
      match c with
      | '(' -> incr depth; Buffer.add_char buf c
      | ')' -> decr depth; Buffer.add_char buf c
      | (' ' | '\t' | '\n' | ',') when !depth = 0 ->
          flush ()
      | _ -> Buffer.add_char buf c
    ) s;
    flush ();
    let parts = List.rev !tokens in
    match parts with
    | "alloca" :: id :: [] -> Alloca (int_of_string id)
    | "load" :: r :: p :: [] -> Load (int_of_string r, int_of_string p)
    | "store" :: v :: p :: [] -> Store (int_of_string v, int_of_string p)
    | "move" :: r :: v :: [] -> Move (int_of_string r, int_of_string v)
    | "binop" :: r :: v1 :: v2 :: [] -> BinOp (int_of_string r, int_of_string v1, int_of_string v2)
    | "phi" :: r :: ops ->
      (* Each op is a parenthesized "(v,b)"; strip parens and split on ','. *)
      let parse_op op =
        (* Drop the surrounding parentheses, then split on ','. *)
        let inner = String.sub op 1 (String.length op - 2) in
        match String.split_on_char ',' inner with
        | [v; b] -> (int_of_string (String.trim v), int_of_string (String.trim b))
        | _ -> failwith ("Bad phi operand: " ^ op)
      in
      Phi (int_of_string r, List.map parse_op ops)
    | "ret" :: [] -> Ret None
    | "ret" :: v :: [] -> Ret (Some (int_of_string v))
    | "br" :: t :: [] -> Br (int_of_string t)
    | "condbr" :: c :: t1 :: t2 :: [] -> CondBr (int_of_string c, int_of_string t1, int_of_string t2)
    | _ -> failwith ("Unknown instr: " ^ s)
  in
  List.iter (fun line ->
    let trimmed = String.trim line in
    if trimmed = "" then ()
    else if String.ends_with ~suffix:":" trimmed then
      let label = String.sub trimmed 0 (String.length trimmed - 1) in
      (* Accept either "1" or "block 1" style headers. *)
      let id_str =
        if String.starts_with ~prefix:"block " label
        then String.sub label 6 (String.length label - 6)
        else label
      in
      finish_block ();
      current_id := int_of_string id_str
    else
      current_insts := parse_instr trimmed :: !current_insts
  ) lines;
  finish_block ();
  let blocks = List.rev !blocks in
  let entry = match blocks with
    | [] -> 0
    | b :: _ -> b.id
  in
  { blocks; entry }

let parse_func_from_file path : func =
  parse_func (read_lines path)