import system.io

universes u v

namespace hide

inductive format.color
| red | green | orange | blue | pink | cyan | grey

inductive format : Type
| nil          : format
| text         : string → format
| choice       : format → format → format
| compose      : format → format → format
| flat_compose : format → format → format
| line         : format
| nest         : nat → format → format
| color_begin  : color → format
| color_end    : format

def format.of_string : string → format := format.text
def format.of_nat (n : nat) : format := format.text n^.to_string

class has_to_format (α : Type u) :=
(to_format : α → format)

instance : has_to_format format :=
⟨id⟩

def to_fmt {α : Type u} [has_to_format α] : α → format :=
has_to_format.to_format

instance has_to_format_string : has_to_format string :=
⟨λ s, format.of_string s⟩

instance has_to_format_nat : has_to_format nat :=
⟨λ n, format.of_nat n⟩

instance has_to_format_char : has_to_format char :=
⟨λ c : char, format.of_string [c]⟩

namespace format

def color.code : color → string
| color.red    := "1"
| color.green  := "2"
| color.orange := "3"
| color.blue   := "4"
| color.pink   := "5"
| color.cyan   := "6"
| color.grey   := "7"

instance : has_coe string format := ⟨text⟩

def space : format := " "

instance : inhabited format :=
⟨format.space⟩

instance : has_append format :=
⟨format.compose⟩

def concat (x y : format) : format :=
x ++ " " ++ y
infixr ^ := concat

def highlight (c : color) (f : format) : format :=
color_begin c ++ f ++ color_end

def above (x y : format) : format :=
x ++ line ++ y

instance nat_to_format : has_coe nat format :=
⟨format.of_nat⟩

instance string_to_format : has_coe string format :=
⟨format.of_string⟩

def is_nil : format → bool
| nil := tt
| _   := ff

def flatten_core : format → format × bool
| (nest i x)    := ((flatten_core x).1, tt)
| (compose x y) := let (fx, hx) := flatten_core x, (fy, hy) := flatten_core y in
                (flat_compose fx fy, hx || hy)
| (choice x y)  := ((flatten_core x).1, tt)
| line          := (" ", tt)
| x             := (x, ff)

def flatten (f : format) : format := (flatten_core f).1

def group (f : format) : format :=
let (flat_f, diff) := flatten_core f in
if diff then choice flat_f f else flat_f

def bracket (l : string) (x : format) (r : string) : format :=
group (nest (utf8_length l) (l ++ x ++ r))

def paren (f : format) : format :=
bracket "(" f ")"

def cbrace (f : format) : format :=
bracket "{" f "}"

def sbracket (f : format) : format :=
bracket "[" f "]"

def dcbrace (f : format) : format :=
bracket "⦃" f "⦄"

def wrap (x y : format) : format :=
x ++ (choice " " line) ++ y

def space_upto_line_break (avail : nat) : format → nat → option (nat × bool)
| nil                len := some (len, ff)
| (color_begin c)    len := some (len, ff)
| color_end          len := some (len, ff)
| (nest i x)         len := space_upto_line_break x len
| line               len := some (len, tt)
| (text t)           len := some (len + t^.length, ff)
| (choice x y)       len := space_upto_line_break y len
| (compose x y)      len :=
  do a ← space_upto_line_break x 0,
     let len' := len + a.1 in
     if len' > avail then none else
     if a.2 then some (len', a.2) else
     space_upto_line_break y len'
| (flat_compose x y) len :=
  do a ← space_upto_line_break x 0,
     let len' := len + a.1 in
     if len' > avail then none else
     if a.2 then some (len', a.2) else
     space_upto_line_break y len'

def space_upto_line_break_list_exceeded : list format → nat → bool
| []        avail := ff
| (f :: fs) avail :=
  match space_upto_line_break avail f 0 with
  | none              := tt
  | some (len, found) :=
    if len > avail then tt else
    if found then ff else
    space_upto_line_break_list_exceeded fs (avail - len)
  end

def pretty_core {α : Type u} (app : α → string → α) (w : nat) (colors : bool) :
  format → list format → α → nat → nat → (α × nat)
| nil                todo s pos indent := (s, pos)
| (color_begin c)    todo s pos indent := (if colors then app s ("\x1B[" ++ c^.code ++ "m") else s, pos)
| color_end          todo s pos indent := (if colors then app s "\x1B[0m" else s, pos)
| (compose x y)      todo s pos indent :=
  let (s', pos') := pretty_core x (y :: todo) s pos indent in pretty_core y todo s' pos' indent
| (flat_compose x y) todo s pos indent :=
  let (s', pos') := pretty_core x (y :: todo) s pos indent in pretty_core y todo s' pos' indent
| (nest i x)         todo s pos indent := pretty_core x todo s pos (indent + i)
| line               todo s pos indent :=
  (nat.rec_on indent (app s "\n") (λ_ s', app s " "), indent)
| (text t)           todo s pos indent := (app s t, pos + t^.length)
| (choice x y)       todo s pos indent :=
  if space_upto_line_break_list_exceeded (x::todo) (w - pos)
  then pretty_core y todo s pos indent
  else pretty_core x todo s pos indent

def pretty {α} (f : format) (app : α → string → α) (w : nat) (colors : bool) (out : α) : α :=
(pretty_core app w colors f [] out 0 0).1

meta def get_option : tactic unit :=
tactic.get_nat_option `pp.width 120 >>= (tactic.to_expr ∘ quote) >>= tactic.exact
meta def get_pp_width : tactic unit :=
tactic.get_nat_option `pp.width 120 >>= (tactic.to_expr ∘ quote) >>= tactic.exact
meta def get_pp_colors : tactic unit :=
tactic.get_bool_option `pp.colors ff >>= (tactic.to_expr ∘ quote) >>= tactic.exact

def print [io.interface] (f : format) (w : auto_param nat ``get_pp_width)
  (colors : auto_param bool ``get_pp_colors) : io unit :=
pretty f (λm s, m >> io.put_str s) w colors (return ())

def to_string (f : format) (w : auto_param nat ``get_pp_width)
  (colors : auto_param bool ``get_pp_colors) : string :=
pretty f string.concat w colors ""

instance : has_to_string format :=
⟨λf, format.to_string f⟩

def trace_fmt {α : Type u} (f : format) (fn : unit → α) : α :=
trace (to_string f) (fn ())

def indent (f : format) (n : nat) : format :=
nest n (line ++ f)

def when {α : Type u} [has_to_format α] : bool → α → format
| tt a := to_fmt a
| ff a := nil

def join (xs : list format) : format :=
xs^.foldl compose (of_string "")
end format

open hide.format list

instance has_to_format_bool : has_to_format bool :=
⟨λ b, if b then of_string "tt" else of_string "ff"⟩

instance has_to_format_decidable {p : Prop} : has_to_format (decidable p) :=
⟨λ b : decidable p, @ite p b _ (of_string "tt") (of_string "ff")⟩

def list.to_format {α : Type u} [has_to_format α] : list α → format
| [] := to_fmt "[]"
| xs := to_fmt "[" ++ group (nest 1 $ join $ list.intersperse ("," ++ line) $ xs^.map to_fmt) ++ to_fmt "]"

instance has_to_format_list {α : Type u} [has_to_format α] : has_to_format (list α) :=
⟨list.to_format⟩

attribute [instance] string.has_to_format

instance has_to_format_name : has_to_format name :=
⟨λ n, to_fmt (name.to_string n)⟩

instance has_to_format_unit : has_to_format unit :=
⟨λ u, to_fmt "()"⟩

instance has_to_format_option {α : Type u} [has_to_format α] : has_to_format (option α) :=
⟨λ o, option.cases_on o
  (to_fmt "none")
  (λ a, to_fmt "(some " ++ nest 6 (to_fmt a) ++ to_fmt ")")⟩

instance sum_has_to_format {α : Type u} {β : Type v} [has_to_format α] [has_to_format β] : has_to_format (sum α β) :=
⟨λ s, sum.cases_on s
  (λ a, to_fmt "(inl " ++ nest 5 (to_fmt a) ++ to_fmt ")")
  (λ b, to_fmt "(inr " ++ nest 5 (to_fmt b) ++ to_fmt ")")⟩

open prod

instance has_to_format_prod {α : Type u} {β : Type v} [has_to_format α] [has_to_format β] : has_to_format (prod α β) :=
⟨λ ⟨a, b⟩, group (nest 1 (to_fmt "(" ++ to_fmt a ++ to_fmt "," ++ line ++ to_fmt b ++ to_fmt ")"))⟩

open sigma

instance has_to_format_sigma {α : Type u} {β : α → Type v} [has_to_format α] [s : ∀ x, has_to_format (β x)]
                                          : has_to_format (sigma β) :=
⟨λ ⟨a, b⟩, group (nest 1 (to_fmt "⟨"  ++ to_fmt a ++ to_fmt "," ++ line ++ to_fmt b ++ to_fmt "⟩"))⟩

open subtype

instance has_to_format_subtype {α : Type u} {p : α → Prop} [has_to_format α] : has_to_format (subtype p) :=
⟨λ s, to_fmt (val s)⟩

end hide