import Std.Data.TreeSet
import Scanner.Protocol

/-!
# Generating Code from Protocol Descriptions

This turns a `Client.ProtocolDesc` into the various Lean source files containing bindings for said protocol.
-/

namespace WaylandScanner

open Client

-- TODO: Maybe use Lean's `Syntax` API here instead of a custom formatter?

private structure GeneratorState where
  imports : Std.TreeSet String := .empty
  code : Array String := #[]

private structure GeneratorContext where
  rootModule : String
  currentModule : String

private abbrev Generator := StateT GeneratorState (ReaderM GeneratorContext)

private def emitImport (mod : String) : Generator Unit := do
  let current := (← read).currentModule
  if current != mod then
    let oldState ← get
    set { oldState with imports := oldState.imports.insert mod }
  return ()

private def emitLocalImport (protocolName : Identifier) (suffix : String := "") : Generator Unit := do
  let root := (← read).rootModule
  emitImport s!"{root}.{protocolName.pascalCase}{suffix}"

private def emitLine (line : String) : Generator Unit := do
  let oldState ← get
  set { oldState with code := oldState.code ++ #[line] }

private def emitNamespace {α : Type} (space : String) (gen : Generator α) : Generator α := do
  emitLine s!"namespace {space}"
  emitLine s!""
  let res ← gen
  emitLine s!""
  emitLine s!"end {space}"
  return res

private def Generator.run (gen : Generator Unit) : ReaderM GeneratorContext String := do
  let (_, res) ← StateT.run gen {}
  let imports := res.imports
    |>.toArray
    |>.map fun m ↦ s!"import {m}"
  let lines := imports ++ #[""] ++ res.code
  return String.intercalate "\n" lines.toList

private def Generator.runProtocol (rootModule : String) (protocol : ProtocolDesc) (suffix : String := "") (gen : Generator Unit) : String :=
  gen.run.run { rootModule, currentModule := s!"{rootModule}.{protocol.name.pascalCase}{suffix}" }

private def Generator.runIndex (rootModule : String) (gen : Generator Unit) : String :=
  gen.run.run { rootModule, currentModule := rootModule }

private def resolveInterface (iface : InterfaceRef) (suffix : String := ".Header") : Generator String := do
  emitLocalImport iface.moduleName suffix
  return iface.typeName.pascalCase

private def resolveEnum (enum : EnumRef) : Generator String := do
  let iface ← resolveInterface enum.toInterfaceRef
  return s!"{iface}.{enum.enumName.pascalCase}"

section Header

private def emitNameDefinitions (protocol : ProtocolDesc) : Generator Unit := do
  for iface in protocol.interfaces do
    emitLine s!"abbrev {iface.name.pascalCase}.name : InterfaceName := {iface.name.quoted}"

private def emitObjectDefinitions (protocol : ProtocolDesc) : Generator Unit := do
  for iface in protocol.interfaces do
    emitLine s!"def {iface.name.pascalCase} : Type := Object {iface.name.pascalCase}.name"

private def emitEnumDefinitions (protocol : ProtocolDesc) : Generator Unit := do
  for iface in protocol.interfaces do
    for enum in iface.enums do
      emitLine s!""
      if enum.bitfield then
        emitLine s!"structure {iface.name.pascalCase}.{enum.name.pascalCase} where"
        for entry in enum.entries do
          emitLine s!"  {entry.name.camelCase} : Bool"
      else
        emitLine s!"inductive {iface.name.pascalCase}.{enum.name.pascalCase} where"
        for entry in enum.entries do
          emitLine s!"  | {entry.name.camelCase}"

def generateHeader (rootModule : String) (protocol : ProtocolDesc) : String :=
  Generator.runProtocol rootModule protocol ".Header" do
    emitNamespace rootModule do
      emitImport "Wayland.Display"
      emitNameDefinitions protocol
      emitLine ""
      emitObjectDefinitions protocol
      emitEnumDefinitions protocol

end Header

section Interface

private def emitEventDefinitions (protocol : ProtocolDesc) : Generator Unit := do
  for iface in protocol.interfaces do
    emitLine s!"inductive {iface.name.pascalCase}.Event where"
    for event in iface.events do
      emitLine s!"  | {event.name.camelCase}"
      for arg in event.arguments do
        let argName := arg.name.camelCase
        match arg.type with
          | .int .none                      => emitLine s!"    ({argName} : Int32)"
          | .int (.some enum)               => emitLine s!"    ({argName} : {← resolveEnum enum})"
          | .uInt .none                     => emitLine s!"    ({argName} : UInt32)"
          | .uInt (.some enum)              => emitLine s!"    ({argName} : {← resolveEnum enum})"
          | .fixed                          => emitLine s!"    ({argName} : Float)"
          | .string .nullable               => emitLine s!"    ({argName} : Option String)"
          | .string .notNull                => emitLine s!"    ({argName} : String)"
          | .array                          => emitLine s!"    ({argName} : ByteArray)"
          | .fd                             => emitLine s!"    ({argName} : Int32)"
          | .object .nullable .none         => emitLine s!"    ({argName} : Option AnyObject)"
          | .object .nullable (.some iface) => emitLine s!"    ({argName} : Option {← resolveInterface iface})"
          | .object .notNull  .none         => emitLine s!"    ({argName} : AnyObject)"
          | .object .notNull  (.some iface) => emitLine s!"    ({argName} : {← resolveInterface iface})"
    emitLine s!""

private def emitInterfaceInstances (protocol : ProtocolDesc) : Generator Unit := do
  let emitMessageDesc (msg : MessageDesc) : Generator Unit := do
    emitLine   "    {"
    emitLine s!"      name := {msg.name.quoted}"
    let nullability (arg : Raw.ArgumentDesc) : String :=
      if arg.allowNull then ".nullable" else ".notNull"
    let interfaceName (arg : Raw.ArgumentDesc) : String :=
      match arg.interface with
        | .some iface => s!"(.some \"{iface}\")"
        | .none       =>   ".none"
    let argumentDesc (arg : Raw.ArgumentDesc) : String :=
      match arg.type with
        | .int    =>   ".int"
        | .uInt   =>   ".uint"
        | .fixed  =>   ".fixed"
        | .string => s!".string {nullability arg}"
        | .object => s!".object {nullability arg} {interfaceName arg}"
        | .newID  => s!".newID {interfaceName arg}"
        | .array  =>   ".array"
        | .fd     =>   ".fd"
    let args := msg.rawArguments
      |>.map argumentDesc
      |>.toList
      |> String.intercalate ", "
    emitLine s!"      args := #[{args}]"
    emitLine s!"      since := {msg.since}"
    emitLine s!"      isDestructor := {msg.isDestructor}"
    emitLine   "    },"
  for iface in protocol.interfaces do
    emitLine s!""
    emitLine s!"instance : Interface {iface.name.pascalCase}.name where"
    emitLine s!"  event := {iface.name.pascalCase}.Event"
    emitLine s!"  version := {iface.version}"
    emitLine s!"  requests := #["
    for req in iface.requests do
      emitMessageDesc req.toMessageDesc
    emitLine s!"  ]"
    emitLine s!"  events := #["
    for ev in iface.events do
      emitMessageDesc ev.toMessageDesc
    emitLine s!"  ]"
    emitLine s!"  demarshalEvent := sorry"
    -- TODO

def generateInterface (rootModule : String) (protocol : ProtocolDesc) : String :=
  Generator.runProtocol rootModule protocol ".Interface" do
    emitNamespace rootModule do
      emitImport "Wayland.Display"
      emitLocalImport protocol.name ".Header"
      emitEventDefinitions protocol
      emitInterfaceInstances protocol

end Interface

section Resolved

private def emitResolvedInterfaces (protocol : ProtocolDesc) : Generator Unit := do
  for iface in protocol.interfaces do
    emitLine s!"instance : ResolvedInterface {iface.name.pascalCase}.name where"
    emitLine s!"  referencedInterfaces := #["
    for ref in iface.referencedInterfaces do
      let name ← resolveInterface ref (suffix := ".Interface")
      emitLine s!"    ⟨{name}.name, inferInstance⟩,"
    emitLine s!"  ]"
    emitLine s!""

def generateResolved (rootModule : String) (protocol : ProtocolDesc) : String :=
  Generator.runProtocol rootModule protocol ".Resolved" do
    emitNamespace rootModule do
      emitImport "Wayland.Display"
      emitLocalImport protocol.name ".Interface"
      emitResolvedInterfaces protocol

end Resolved

section Requests

private def emitRequests (protocol : ProtocolDesc) : Generator Unit := do
  for iface in protocol.interfaces do
    for req in iface.requests do
      emitLine s!""
      emitLine s!"def {iface.name.pascalCase}.{req.name.camelCase}"
      emitLine s!"    (ω : {iface.name.pascalCase})"
      -- TODO: Special wl_callback handling
      for arg in req.arguments do
        let argName := arg.name.camelCase
        match arg.type with
          | .int .none                      => emitLine s!"    ({argName} : Int32)"
          | .int (.some enum)               => emitLine s!"    ({argName} : {← resolveEnum enum})"
          | .uInt .none                     => emitLine s!"    ({argName} : UInt32)"
          | .uInt (.some enum)              => emitLine s!"    ({argName} : {← resolveEnum enum})"
          | .fixed                          => emitLine s!"    ({argName} : Float)"
          | .string .nullable               => emitLine s!"    ({argName} : Option String)"
          | .string .notNull                => emitLine s!"    ({argName} : String)"
          | .array                          => emitLine s!"    ({argName} : ByteArray)"
          | .fd                             => emitLine s!"    ({argName} : Int32)"
          | .object .nullable .none         => emitLine s!"    ({argName} : Option AnyObject)"
          | .object .nullable (.some iface) => emitLine s!"    ({argName} : Option {← resolveInterface iface (suffix := ".Resolved")})"
          | .object .notNull  .none         => emitLine s!"    ({argName} : AnyObject)"
          | .object .notNull  (.some iface) => emitLine s!"    ({argName} : {← resolveInterface iface (suffix := ".Resolved")})"
      match req.returnType with
        | .unit                 => emitLine s!"    : DisplayM σ Unit :="
        | .object .none         => emitLine s!"    : DisplayM σ AnyObject :="
        | .object (.some iface) => emitLine s!"    : DisplayM σ {← resolveInterface iface (suffix := ".Resolved")} :="
      emitLine s!"  sorry"
      -- TODO

def generateRequests (rootModule : String) (protocol : ProtocolDesc) : String :=
  Generator.runProtocol rootModule protocol ".Requests" do
    emitNamespace rootModule do
      emitImport "Wayland.Display"
      emitLocalImport protocol.name ".Resolved"
      emitLine "variable {σ : Type}"
      emitRequests protocol

end Requests

section Umbrella

def generateUmbrella (rootModule : String) (protocol : ProtocolDesc) : String :=
  Generator.runProtocol rootModule protocol "" do
    emitLocalImport protocol.name ".Header"
    emitLocalImport protocol.name ".Interface"
    emitLocalImport protocol.name ".Resolved"
    emitLocalImport protocol.name ".Requests"

end Umbrella

section Index

def generateIndex (rootModule : String) (protocols : Array ProtocolDesc) : String :=
  Generator.runIndex rootModule do
    for protocol in protocols do
      emitLocalImport protocol.name

end Index

end WaylandScanner