/// The macro expansions of the last check of each file, and the expansion of
/// one call, printed: one step, or all the way down.
///
/// The compiler hands each expansion to `Macro.observe` while it checks. A
/// call is found again by its range, which is its call site; a template's own
/// nodes carry that range too, so a call whose expansion is another macro
/// call expands again at the same site, and the form expanded is what tells
/// those steps apart.
module BjolangLsp.Expansion

open System
open System.Text
open Bjolang
open Bjolang.Lexer
open Bjolang.Ast

type Record =
    { Macro: string
      Site: Range
      Input: SExpr
      Output: SExpr }

// ---------------------------------------------------------------------------
// Printing
// ---------------------------------------------------------------------------

let private punctuation =
    Map.ofList
        [ Colon, ":"
          Dot, "."
          Comma, ","
          CommaAt, ",@"
          Quote, "'"
          SynQuote, "#'"
          Spread, "..."
          Hash, "#" ]

let private escaped (s: string) =
    let b = StringBuilder("\"")

    for c in s do
        match c with
        | '"' -> b.Append "\\\"" |> ignore
        | '\\' -> b.Append "\\\\" |> ignore
        | '\n' -> b.Append "\\n" |> ignore
        | '\t' -> b.Append "\\t" |> ignore
        | c -> b.Append c |> ignore

    b.Append('"').ToString()

let private character (c: int) =
    match c with
    | 32 -> @"#\space"
    | 10 -> @"#\newline"
    | 9 -> @"#\tab"
    | c -> @"#\" + Char.ConvertFromUtf32 c

/// An atom as source. `name` spells a symbol, so that a hygiene renaming can
/// be shown as the template wrote it. A keyword at the head of a list is a
/// clause's, `:for`; anywhere else an argument's, `#:when`.
let private atom (name: string -> string) (atHead: bool) (t: Token) =
    match t with
    | Symbol s
    | ResolvedSymbol s -> name s
    | Keyword k -> if atHead then ":" + k else "#:" + k
    | QuotedSymbol s -> "'" + s
    | TypeVar v -> "%" + v.TrimStart '%'
    | NumberLit n -> n
    | StringLit s -> escaped s
    | CharLit c -> character c
    | BoolLit b -> if b then "#t" else "#f"
    | other -> Map.tryFind other punctuation |> Option.defaultValue "?"

/// A form on one line.
let rec private flat (name: string -> string) (atHead: bool) (s: SExpr) : string =
    match s with
    | SAtom t -> atom name atHead t.Token
    | SList(items, _) -> "(" + (items |> List.mapi (fun i x -> flat name (i = 0) x) |> String.concat " ") + ")"

/// The heads whose first argument stays on their line and whose body is
/// indented by two: a binding form or a definition, as Emacs indents them.
let private bodyForms =
    set
        [ "let"; "let*"; "letrec"; "let/mono"; "when"; "unless"; "fun"; "bjoroutine"; "defun"; "defbjo"
          "def"; "match"; "case"; "with-open"; "with-return"; "loop"; "seql"; "if"; "begin"; "parameterize"
          "try"; "def/macro" ]

let private width = 78

/// A form, indented as a reader of Lisp expects: on one line if it fits, else
/// its head and first argument on the first line and the rest beneath, two in
/// for a body form and under the first argument for a call.
let rec private pretty (name: string -> string) (indent: int) (atHead: bool) (s: SExpr) : string =
    let one = flat name atHead s

    match s with
    | SList(head :: args, _) when indent + one.Length > width && not args.IsEmpty ->
        let headText = flat name true head

        // A template's `match` arrives as `match__12`, which is still a match.
        let isBody =
            match head with
            | SAtom { Token = Symbol h } -> bodyForms.Contains(RegularExpressions.Regex.Replace(h, @"__\d+$", ""))
            | _ -> false

        let first = pretty name (indent + headText.Length + 2) false args.Head

        let restIndent =
            match head with
            | SAtom _ when not isBody -> indent + headText.Length + 2
            | _ -> indent + 2

        let rest =
            args.Tail
            |> List.map (fun a -> "\n" + String(' ', restIndent) + pretty name restIndent false a)
            |> String.concat ""

        // A list headed by a list, such as a `let`'s bindings or a `match`
        // arm, keeps its items in a column.
        match head with
        | SList _ ->
            let inner = indent + 1

            "("
            + pretty name inner true head
            + (args |> List.map (fun a -> "\n" + String(' ', inner) + pretty name inner false a) |> String.concat "")
            + ")"
        | _ -> "(" + headText + " " + first + rest + ")"
    | _ -> one

// ---------------------------------------------------------------------------
// What the last checks expanded
// ---------------------------------------------------------------------------

/// Each file's expansions, from its last check, in the order they were made.
let private byFile = Collections.Generic.Dictionary<string, Record list>()

/// Each name hygiene made, and the spelling the template wrote.
let private renamed = Collections.Generic.Dictionary<string, string>()

/// Runs `check` with the compiler reporting its expansions, and keeps those
/// made in `path` as the file's.
let recording (path: string) (check: unit -> 'a) : 'a =
    let seen = ResizeArray<Record>()

    Macro.observe <-
        Some(fun macro site input output names ->
            for KeyValue(fresh, original) in names do
                renamed[fresh] <- original

            seen.Add
                { Macro = macro
                  Site = site
                  Input = input
                  Output = output })

    try
        check ()
    finally
        Macro.observe <- None
        let full = IO.Path.GetFullPath path

        // A form can be expanded more than once in one check, as a
        // declaration is checked again in a copy; the first is kept.
        byFile[full] <-
            seen
            |> Seq.filter (fun r -> r.Site.File <> "" && IO.Path.GetFullPath r.Site.File = full)
            |> Seq.distinctBy (fun r -> r.Site, flat id false r.Input)
            |> List.ofSeq

let private holds (line: int) (column: int) (r: Range) =
    (line > r.Start.Line || (line = r.Start.Line && column >= r.Start.Column))
    && (line < r.End.Line || (line = r.End.Line && column <= r.End.Column))

let private size (r: Range) = (r.End.Line - r.Start.Line) * 100000 + (r.End.Column - r.Start.Column)

/// The expansion of the innermost macro call at `line` (1-based) and
/// `column` (0-based) of `path`, from its last check: the macro's name and
/// the expansion as text. One step is what the macro made; `full` expands
/// every macro call in that too, as the compiler went on to.
///
/// `hygienic` keeps the names hygiene made, `tmp__37`, rather than showing
/// them as the template wrote them, which is what tells a template's name
/// from the caller's of the same spelling.
let at (path: string) (line: int) (column: int) (full: bool) (hygienic: bool) : (string * string) option =
    match byFile.TryGetValue(IO.Path.GetFullPath path) with
    | false, _ -> None
    | true, records ->
        let calls = records |> List.filter (fun r -> holds line column r.Site)

        match calls |> List.sortBy (fun r -> size r.Site) |> List.tryHead with
        | None -> None
        | Some innermost ->
            // The first expansion at a site is of the call as written; any
            // after it are of what that made.
            let call = records |> List.find (fun r -> r.Site = innermost.Site)

            let name (s: string) =
                if hygienic then s
                else
                    match renamed.TryGetValue s with
                    | true, original -> original
                    | _ -> s

            /// Every macro call in `s` replaced by its expansion, again and
            /// again, as far as the check went.
            let rec expand (depth: int) (s: SExpr) : SExpr =
                match s with
                | SList(items, r) when depth < 200 ->
                    let key = flat id false s

                    match records |> List.tryFind (fun x -> x.Site = r && flat id false x.Input = key) with
                    | Some x -> expand (depth + 1) x.Output
                    | None -> SList(items |> List.map (expand depth), r)
                | _ -> s

            let result = if full then expand 0 call.Output else call.Output
            Some(call.Macro, pretty name 0 false result)