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.protocol suffix
  return iface.interface.pascalCase

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

private def typeToLean : ArgumentType → Generator String
  | .int .none                      => do return "Int32"
  | .int (.some enum)               => resolveEnum enum
  | .uint .none                     => do return "UInt32"
  | .uint (.some enum)              => resolveEnum enum
  | .fixed                          => do return "Float32"
  | .string .nullable               => do return "Option String"
  | .string .notNull                => do return "String"
  | .array                          => do return "ByteArray"
  | .fd                             => do return "Fd"
  | .object .nullable .none         => do return "Option AnyObject"
  | .object .nullable (.some iface) => do return s!"Option {← resolveInterface iface}"
  | .object .notNull  .none         => do return s!"AnyObject"
  | .object .notNull  (.some iface) => resolveInterface iface
  | .newID .none                    => do return s!"AnyObject"
  | .newID (.some iface)            => resolveInterface iface

private def marshalType : ArgumentType → String
  | .int _      => ".int"
  | .uint _     => ".uint"
  | .fixed      => ".fixed"
  | .string _   => ".string"
  | .object _ _ => ".object"
  | .newID _    => ".object"
  | .array      => ".array"
  | .fd         => ".fd"

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!"abbrev {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}"

private def emitEnumMarshallable (protocol : ProtocolDesc) : Generator Unit := do
  for iface in protocol.interfaces do
    for enum in iface.enums do
      emitLine s!""
      emitNamespace s!"{iface.name.pascalCase}.{enum.name.pascalCase}" do
        emitLine s!"instance : Marshallable {iface.name.pascalCase}.{enum.name.pascalCase} where"
        if enum.bitfield then
          emitLine s!"  marshal"
          emitLine s!"    | .uint, μ => ArgumentMarshaller.marshalUInt <|"
          emitLine s!"        0"
          for entry in enum.entries do
            emitLine s!"        ||| (if μ.{entry.name.camelCase} then {entry.value} else 0)"
          emitLine s!"    | _, _ => throw \"Marshalling unexpected type, expected uint.\""
          emitLine s!"  demarshal"
          emitLine s!"    | .uint => do"
          emitLine s!"        let μ ← ArgumentDemarshaller.demarshalUInt"
          emitLine   "        return {"
          for entry in enum.entries do
            emitLine s!"          {entry.name.camelCase} := (μ &&& {entry.value}) == {entry.value}"
          emitLine   "        }"
          emitLine s!"    | _ => throw \"Demarshalling unexpected type, expected uint.\""
        else
          emitLine s!"  marshal"
          for entry in enum.entries do
            emitLine s!"    | .int,  .{entry.name.camelCase} => ArgumentMarshaller.marshalInt {entry.value}"
            emitLine s!"    | .uint, .{entry.name.camelCase} => ArgumentMarshaller.marshalUInt {entry.value}"
          emitLine s!"    | _, _ => throw \"Marshalling unexpected type, expected int or uint.\""
          emitLine s!"  demarshal"
          let emitMatch (demarshal : String) : Generator Unit := do
            emitLine s!"      match (← ArgumentDemarshaller.{demarshal}) with"
            for entry in enum.entries do
              emitLine s!"        | {entry.value} => pure .{entry.name.camelCase}"
            emitLine s!"        | _ => throw \"Invalid enum value.\""
          emitLine s!"    | .int => do"
          emitMatch "demarshalInt"
          emitLine s!"    | .uint => do"
          emitMatch "demarshalUInt"
          emitLine s!"    | _ => throw \"Demarshalling unexpected type, expected int or uint.\""

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
      emitEnumMarshallable 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
        emitLine s!"    ({arg.name.camelCase} : {← typeToLean arg.type})"
    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 : Nullability → String
      | .nullable => ".nullable"
      | .notNull  => ".notNull"
    let interfaceName : Option InterfaceRef → String
      | .some iface => s!"(.some {iface.interface.quoted})"
      | .none       =>   ".none"
    let argumentDesc (arg : ArgumentDesc) : String :=
      match arg.type with
        | .int _               =>   ".int"
        | .uint _              =>   ".uint"
        | .fixed               =>   ".fixed"
        | .string null         => s!".string {nullability null}"
        | .object null iface   => s!".object {nullability null} {interfaceName iface}"
        | .newID (.some iface) => s!".newID (.some {iface.interface.quoted})"
        | .newID .none         =>   ".string .notNull, .uint, .newID .none"
        | .array               =>   ".array"
        | .fd                  =>   ".fd"
    let args := msg.arguments
      |>.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 instInterface{iface.name.pascalCase} : 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"
    for ev in iface.events do
      emitLine s!"    | {ev.opcode} => do"
      emitLine s!"      return .{ev.name.camelCase}"
      for arg in ev.arguments do
        emitLine s!"        (← Marshallable.demarshal {marshalType arg.type})"
    emitLine "    | _ => throw s!\"Received invalid opcode\""

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 instResolvedInterface{iface.name.pascalCase} : 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
      let baseName := s!"{iface.name.pascalCase}.{req.name.camelCase}"
      let rawName :=
        if req.returnType.isWlCallback
        then baseName ++ ".raw"
        else baseName
      emitLine s!""
      emitLine s!"def {rawName}"
      emitLine s!"    (ω : {iface.name.pascalCase})"
      for arg in req.inputArguments do
        emitLine s!"    ({arg.name.camelCase} : {← typeToLean arg.type})"
      match req.returnType with
        | .unit                 => emitLine s!"    : DisplayM σ Unit :="
        | .object .none         => emitLine s!"    : DisplayM σ (Object ι) :="
        | .object (.some iface) => emitLine s!"    : DisplayM σ {← resolveInterface iface (suffix := ".Resolved")} :="
      match req.returnType with
        | .unit     => emitLine s!"  ω.send do"
        | .object _ => emitLine s!"  ω.sendConstructor do"
      for arg in req.arguments do
        match arg.type with
          | .newID (.some _) =>
            emitLine s!"    ArgumentMarshaller.marshalAnyObject .none"
          | .newID .none =>
            emitLine s!"    ArgumentMarshaller.marshalString (.some ι)"
            emitLine s!"    ArgumentMarshaller.marshalUInt (inferInstanceAs (Interface ι)).version"
            emitLine s!"    ArgumentMarshaller.marshalAnyObject .none"
          | _ =>
            emitLine s!"    Marshallable.marshal {marshalType arg.type} {arg.name.camelCase}"
      emitLine s!"    return {req.opcode}"
      -- Special case wl_callback handling
      if req.returnType.isWlCallback then
        emitLine s!""
        emitLine s!"def {baseName}"
        emitLine s!"    (ω : {iface.name.pascalCase})"
        for arg in req.inputArguments do
          emitLine s!"    ({arg.name.camelCase} : {← typeToLean arg.type})"
        emitLine s!"    (γ : DisplayM σ Unit)"
        emitLine s!"    : DisplayM σ Unit := do"
        let argNames := #["ω"] ++ (req.inputArguments.map fun arg ↦ arg.name.camelCase)
          |>.toList
          |> String.intercalate " "
        emitLine s!"  let ρ ← {rawName} {argNames}"
        emitLine s!"  ρ.listen fun"
        emitLine s!"    | .done _ => do ρ.keep; γ"

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} {ι : InterfaceName} [iface : ResolvedInterface ι]"
      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