import Scanner.Protocol
import Std.Data.TreeMap
import Std.Data.TreeSet

/-!
# Processing of Protocol Descriptions

This is mainly used to resolve interface/enum references within protocols, converting a `Raw.ProtocolDesc` into a
`Client.ProtocolDesc`. Note that this requires a global view of all protocols since interfaces can mention interfaces
from other protocols. Thus, this needs to be done after parsing all protocols.
-/

namespace WaylandScanner

private structure InterfaceInfo where
  protocol : Identifier
  enums : Array Identifier

private abbrev InterfaceMap := Std.TreeMap String InterfaceInfo

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

private structure TrackingState where
  referencedInterfaces : Std.TreeSet Client.InterfaceRef

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

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 (← read).interfaces.get? name with
    | .some info => pure info
    | .none => .error s!"Unable to find referenced interface {name}"

private def resolveInterface (name : String) : TrackingProcessor Client.InterfaceRef := do
  let info ← interfaceInfo name
  let ref := {
    protocol := info.protocol
    interface := name
  }
  let old ← get
  set { old with referencedInterfaces := old.referencedInterfaces.insert ref }
  return ref

private def resolveEnum (name : String) : Processor Client.EnumRef := do
  let (interface, enum) ←
    match (name.split ".").toArray with
      | #[enum]            => pure ((← read).currentInterface.interface, 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 {
    protocol := info.protocol
    interface
    enum
  }

private def processEnum (raw : Raw.EnumDesc) : Processor Client.EnumDesc := do
  return {
    name := if raw.name != "event" then raw.name else "event'" -- We define an `Iface.Event` type, so avoid any possible conflicts.
    description := raw.description
    since := raw.since
    bitfield := raw.bitfield
    entries := ← raw.entries.mapM processEnumEntry
  }

private def processArgument (raw : Raw.ArgumentDesc) : TrackingProcessor Client.ArgumentDesc := do
  return {
    name := raw.name
    type :=
      ← match raw.type with
        | .int    => Client.ArgumentType.int <$> raw.enum.mapM (liftM ∘ resolveEnum)
        | .uint   => Client.ArgumentType.uint <$> raw.enum.mapM (liftM ∘ 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) : TrackingProcessor 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
    rawArguments := raw.arguments
    arguments := ← otherArgs.mapM processArgument
    returnType := returnType
  }

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

private def processInterface (raw : Raw.InterfaceDesc) : Processor Client.InterfaceDesc := do
  let state : TrackingState := {
    referencedInterfaces := { (← read).currentInterface }
  }
  let (requests, state) ← StateT.run (raw.requests.mapM processRequest) state
  let (events, state) ← StateT.run (raw.events.mapM processEvent) state
  return {
    name := raw.name
    version := raw.version
    description := raw.description
    enums := ← raw.enums.mapM processEnum
    requests := requests
    events := events
    referencedInterfaces := state.referencedInterfaces
  }

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 := { protocol := raw.name, interface := iface.name }
      }
  }

/-! Compute the transitive closure w.r.t. `referencedInterfaces`. -/
private partial def transitiveClosure (protocols : Array Client.ProtocolDesc) : Array Client.ProtocolDesc :=
  let successors : Std.TreeMap Client.InterfaceRef (List Client.InterfaceRef) := .ofArray <| protocols.flatMap
    fun protocol ↦ protocol.interfaces.map
      fun iface ↦ ((⟨protocol.name, iface.name⟩ : Client.InterfaceRef), iface.referencedInterfaces.toList)
  let rec dfs : Std.TreeSet Client.InterfaceRef → List Client.InterfaceRef → Std.TreeSet Client.InterfaceRef
    | visited, [] => visited
    | visited, (x::xs) =>
      if visited.contains x
      then dfs visited xs
      else dfs (visited.insert x) (successors.get! x ++ xs)
  let interfaceClosure (iface : Client.InterfaceDesc) : Client.InterfaceDesc :=
    { iface with referencedInterfaces := dfs .empty iface.referencedInterfaces.toList }
  let protocolClosure (protocol : Client.ProtocolDesc) : Client.ProtocolDesc :=
    { protocol with interfaces := protocol.interfaces.map interfaceClosure }
  protocols.map protocolClosure

-- TODO: Correctness proof for `transitiveClosure`

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

end WaylandScanner