type position = {
  line : int;
  col : int;
  offset : int;
}
[@@deriving show, eq]

type span = {
  start : position;
  end_ : position;
  file : string;
}
[@@deriving show, eq]

type severity =
  | Error
  | Warning
  | Info
  | Hint
[@@deriving show, eq]

type error_kind =
  | LexError
  | ParseError
  | TypeError
  | IRError
  | CodeGenError
  | InternalError
[@@deriving show, eq]

type diagnostic = {
  severity : severity;
  kind : error_kind;
  message : string;
  span : span option;
  notes : (string * span option) list;
  fix : string option;
}
[@@deriving show]

type context = {
  mutable diagnostics : diagnostic list;
  mutable error_count : int;
  mutable warning_count : int;
  source_files : (string, string) Hashtbl.t;
}

let create_context () : context = {
  diagnostics = [];
  error_count = 0;
  warning_count = 0;
  source_files = Hashtbl.create 16;
}

let register_source ctx filename content =
  Hashtbl.replace ctx.source_files filename content

let add_diagnostic ctx diag =
  ctx.diagnostics <- diag :: ctx.diagnostics;
  match diag.severity with
  | Error -> ctx.error_count <- ctx.error_count + 1
  | Warning -> ctx.warning_count <- ctx.warning_count + 1
  | _ -> ()

let error ?span ?notes ?fix kind message =
  {
    severity = Error;
    kind;
    message;
    span;
    notes = Option.value notes ~default:[];
    fix;
  }

let warning ?span ?notes ?fix kind message =
  {
    severity = Warning;
    kind;
    message;
    span;
    notes = Option.value notes ~default:[];
    fix;
  }

let report_error ctx ?span ?notes ?fix kind message =
  add_diagnostic ctx (error ?span ?notes ?fix kind message)

let report_warning ctx ?span ?notes ?fix kind message =
  add_diagnostic ctx (warning ?span ?notes ?fix kind message)

let has_errors ctx = ctx.error_count > 0

let get_source_lines filename start_line end_line =
  try
    let ic = open_in filename in
    let rec read_lines acc line_num =
      if line_num > end_line then (
        close_in ic;
        List.rev acc
      ) else
        match input_line ic with
        | line ->
            let acc' =
              if line_num >= start_line then (line_num, line) :: acc
              else acc
            in
            read_lines acc' (line_num + 1)
        | exception End_of_file ->
            close_in ic;
            List.rev acc
    in
    read_lines [] 1
  with Sys_error _ -> []

let get_cached_source_lines ctx filename start_line end_line =
  match Hashtbl.find_opt ctx.source_files filename with
  | Some content ->
      let lines = String.split_on_char '\n' content in
      let indexed =
        List.mapi (fun i line -> (i + 1, line)) lines
        |> List.filter (fun (num, _) -> num >= start_line && num <= end_line)
      in
      indexed
  | None -> get_source_lines filename start_line end_line

let use_color () =
  match Sys.getenv_opt "NO_COLOR" with
  | Some _ -> false
  | None -> (
      match Sys.getenv_opt "TERM" with
      | Some term when term <> "dumb" -> Unix.isatty Unix.stdout
      | _ -> false)

let color_reset = "\027[0m"
let color_red = "\027[31m"
let color_yellow = "\027[33m"
let color_blue = "\027[34m"
let color_cyan = "\027[36m"
let color_bold = "\027[1m"

let with_color color s =
  if use_color () then color ^ s ^ color_reset else s

let format_severity = function
  | Error -> with_color (color_bold ^ color_red) "error"
  | Warning -> with_color (color_bold ^ color_yellow) "warning"
  | Info -> with_color (color_bold ^ color_blue) "info"
  | Hint -> with_color (color_bold ^ color_cyan) "hint"

let format_kind = function
  | LexError -> "lex"
  | ParseError -> "parse"
  | TypeError -> "type"
  | IRError -> "ir"
  | CodeGenError -> "codegen"
  | InternalError -> "internal"

let format_position pos =
  Printf.sprintf "%d:%d" pos.line (pos.col + 1)

let format_span span =
  Printf.sprintf "%s:%s" span.file (format_position span.start)

let format_source_context ctx span =
  let start_line = max 1 (span.start.line - 3) in
  let end_line = span.end_.line + 3 in
  let lines = get_cached_source_lines ctx span.file start_line end_line in
  
  let max_line_num = List.fold_left (fun acc (num, _) -> max acc num) 0 lines in
  let line_num_width = String.length (string_of_int max_line_num) in
  
  let format_line (line_num, line_text) =
    let line_num_str = Printf.sprintf "%*d" line_num_width line_num in
    let is_error_line = line_num >= span.start.line && line_num <= span.end_.line in
    
    if is_error_line then
      let prefix = with_color color_blue (line_num_str ^ " | ") in
      let underline =
        if line_num = span.start.line && line_num = span.end_.line then
          let spaces = String.make (line_num_width + 3 + span.start.col) ' ' in
          let carets = String.make (max 1 (span.end_.col - span.start.col)) '^' in
          spaces ^ with_color color_red carets
        else if line_num = span.start.line then
          let spaces = String.make (line_num_width + 3 + span.start.col) ' ' in
          let carets = String.make (String.length line_text - span.start.col) '^' in
          spaces ^ with_color color_red carets
        else if line_num = span.end_.line then
          let spaces = String.make (line_num_width + 3) ' ' in
          let carets = String.make span.end_.col '^' in
          spaces ^ with_color color_red carets
        else
          let spaces = String.make (line_num_width + 3) ' ' in
          let carets = String.make (String.length line_text) '^' in
          spaces ^ with_color color_red carets
      in
      prefix ^ line_text ^ "\n" ^ underline
    else
      let prefix = with_color color_blue (line_num_str ^ " | ") in
      prefix ^ line_text
  in
  
  String.concat "\n" (List.map format_line lines)

let format_diagnostic ctx diag =
  let buf = Buffer.create 256 in
  
  let header =
    match diag.span with
    | Some span ->
        Printf.sprintf "%s: [%s] %s"
          (format_severity diag.severity)
          (format_kind diag.kind)
          (format_span span)
    | None ->
        Printf.sprintf "%s: [%s]"
          (format_severity diag.severity)
          (format_kind diag.kind)
  in
  Buffer.add_string buf (with_color color_bold header);
  Buffer.add_char buf '\n';
  
  Buffer.add_string buf ("  " ^ diag.message);
  Buffer.add_char buf '\n';
  
  (match diag.span with
  | Some span ->
      Buffer.add_char buf '\n';
      Buffer.add_string buf (format_source_context ctx span);
      Buffer.add_char buf '\n'
  | None -> ());
  
  List.iter (fun (note, span_opt) ->
    Buffer.add_string buf ("\n  " ^ with_color color_cyan "note:" ^ " " ^ note);
    Buffer.add_char buf '\n';
    match span_opt with
    | Some span ->
        Buffer.add_string buf (format_source_context ctx span);
        Buffer.add_char buf '\n'
    | None -> ()
  ) diag.notes;
  
  (match diag.fix with
  | Some fix ->
      Buffer.add_string buf ("\n  " ^ with_color color_cyan "help:" ^ " " ^ fix);
      Buffer.add_char buf '\n'
  | None -> ());
  
  Buffer.contents buf

let format_all_diagnostics ctx =
  let diags = List.rev ctx.diagnostics in
  String.concat "\n\n" (List.map (format_diagnostic ctx) diags)

let print_diagnostics ctx =
  if ctx.diagnostics <> [] then begin
    prerr_endline (format_all_diagnostics ctx);
    prerr_endline "";
    
    let summary =
      match (ctx.error_count, ctx.warning_count) with
      | (0, 0) -> ""
      | (e, 0) ->
          with_color color_red (Printf.sprintf "%d error%s" e (if e = 1 then "" else "s"))
      | (0, w) ->
          with_color color_yellow (Printf.sprintf "%d warning%s" w (if w = 1 then "" else "s"))
      | (e, w) ->
          with_color color_red (Printf.sprintf "%d error%s" e (if e = 1 then "" else "s")) ^
          " and " ^
          with_color color_yellow (Printf.sprintf "%d warning%s" w (if w = 1 then "" else "s"))
    in
    if summary <> "" then
      prerr_endline (with_color color_bold ("Compilation finished with " ^ summary))
  end

let clear_diagnostics ctx =
  ctx.diagnostics <- [];
  ctx.error_count <- 0;
  ctx.warning_count <- 0

let make_position ~line ~col ~offset = { line; col; offset }

let make_span ~start ~end_ ~file = { start; end_; file }

let merge_spans s1 s2 =
  if s1.file <> s2.file then s1
  else { file = s1.file; start = s1.start; end_ = s2.end_ }