import Scanner.Protocol
import Std.Data.TreeMap

namespace WaylandScanner

private structure InterfaceInfo where
  moduleName : Identifier
  enums : Array Identifier

private abbrev InterfaceMap := Std.TreeMap String InterfaceInfo

private structure Context where
  interfaces : InterfaceMap
  currentInterface : Client.InterfaceRef
  referencedInterfaces : Array Client.InterfaceRef

private abbrev Processor : Type → Type := StateT Context (Except String)

private def processEnumEntry (raw : Raw.EnumEntryDesc) : Processor Client.EnumEntryDesc := do
  return {
    name := raw.name
    value := raw.value
    description := raw.description
    since := raw.since
    deprecatedSince := raw.deprecatedSince
  }

private def interfaceInfo (name : String) : Processor InterfaceInfo := do
  match (← get).interfaces.get? name with
  | .some info => pure info
  | .none => .error s!"Unable to find referenced interface {name}"

private def resolveInterface (name : String) : Processor Client.InterfaceRef := do
  let info ← interfaceInfo name
  let ref := {
    moduleName := info.moduleName
    typeName := name
  }
  let old ← get
  set { old with referencedInterfaces := old.referencedInterfaces ++ #[ref] }
  return ref

private def resolveEnum (name : String) : Processor Client.EnumRef := do
  let (interface, enum) ←
    match (name.split ".").toArray with
    | #[enum]            => pure ((← get).currentInterface.typeName, enum.toString)
    | #[interface, enum] => pure (interface.toString, enum.toString)
    | _                  => throw s!"Malformed enum reference {name}"
  let info ← interfaceInfo interface
  if not (info.enums.contains enum) then
    throw s!"Interface {repr interface} does not contain enum {repr enum}"
  return {
    moduleName := info.moduleName
    typeName := interface
    enumName := enum
  }

private def processEnum (raw : Raw.EnumDesc) : Processor Client.EnumDesc := do
  return {
    name := raw.name
    description := raw.description
    since := raw.since
    bitfield := raw.bitfield
    entries := ← raw.entries.mapM processEnumEntry
  }

private def processArgument (raw : Raw.ArgumentDesc) : Processor Client.ArgumentDesc := do
  return {
    name := raw.name
    type :=
      ← match raw.type with
      | .int    => Client.ArgumentType.int <$> raw.enum.mapM resolveEnum
      | .uInt   => Client.ArgumentType.uInt <$> raw.enum.mapM resolveEnum
      | .fixed  => Client.ArgumentType.fixed |> pure
      | .string => Client.ArgumentType.string (.ofAllowNull raw.allowNull) |> pure
      | .array  => Client.ArgumentType.array |> pure
      | .fd     => Client.ArgumentType.fd |> pure
      | .newID  => Client.ArgumentType.object .notNull <$> raw.interface.mapM resolveInterface
      | .object => Client.ArgumentType.object (.ofAllowNull raw.allowNull) <$> raw.interface.mapM resolveInterface
    description := raw.description
  }

private def processRequest (raw : Raw.MessageDesc) : Processor Client.RequestDesc := do
  let newIDArgs := raw.arguments.filter fun arg ↦ arg.type == .newID
  let otherArgs := raw.arguments.filter fun arg ↦ arg.type != .newID
  let returnType ←
    match newIDArgs with
    | #[]    => Client.ReturnType.unit |> pure
    | #[arg] => Client.ReturnType.object <$> arg.interface.mapM resolveInterface
    | _      => throw s!"Expected at most one new_id argument for request in {raw.name}"
  return {
    name := raw.name
    opcode := raw.opcode
    description := raw.description
    isDestructor := raw.type == .destructor
    since := raw.since
    deprecatedSince := raw.deprecatedSince
    arguments := ← otherArgs.mapM processArgument
    returnType := returnType
  }

private def processEvent (raw : Raw.MessageDesc) : Processor Client.EventDesc := do
  return {
    name := raw.name
    opcode := raw.opcode
    description := raw.description
    isDestructor := raw.type == .destructor
    since := raw.since
    deprecatedSince := raw.deprecatedSince
    arguments := ← raw.arguments.mapM processArgument
  }

private def processInterface (raw : Raw.InterfaceDesc) : Processor Client.InterfaceDesc := do
  let requests ← raw.requests.mapM processRequest
  let events ← raw.events.mapM processEvent
  let state ← get
  let referencedInterfaces := state.referencedInterfaces ++ #[state.currentInterface]
  return {
    name := raw.name
    version := raw.version
    description := raw.description
    enums := ← raw.enums.mapM processEnum
    requests := requests
    events := events
    referencedInterfaces := referencedInterfaces
      |>.toList
      |>.eraseDups
      |>.toArray
  }

private def processProtocol (interfaces : InterfaceMap) (raw : Raw.ProtocolDesc) : Except String Client.ProtocolDesc := do
  return {
    name := raw.name
    copyright := raw.copyright
    description := raw.description
    interfaces := ← raw.interfaces.mapM
      fun iface ↦ (processInterface iface).run' {
        interfaces := interfaces
        currentInterface := { moduleName := raw.name, typeName := iface.name }
        referencedInterfaces := #[]
      }
  }

def processProtocols (raw : Array Raw.ProtocolDesc) : Except String (Array Client.ProtocolDesc) :=
  let interfaces : InterfaceMap := .ofArray <| raw.flatMap
    fun protocol ↦ protocol.interfaces.map
      fun iface ↦ (iface.name, (⟨protocol.name, iface.enums.map Raw.EnumDesc.name⟩ : InterfaceInfo))
  raw.mapM (processProtocol interfaces)

end WaylandScanner