import Lean

import NKL

open Lean Meta

abbrev Handle := IO.FS.Handle

private def dq (str : String) := s!"\"{str}\""

private def cname : Name -> MetaM String

| .str _ s => return s

| n => throwError s!"Invalid Constructor Name {n}"

private def printTuple

(h : Handle) (isStruct : Bool)

(name : String) (fields : List String) : MetaM Unit :=

do

let fields := List.map dq fields

let fn := if isStruct then "struct" else "cons"

h.putStrLn s!"{name.capitalize} = {fn}({dq name}, {fields})"

private def genStructure (h : Handle) (si : StructureInfo) : MetaM Unit := do

let name <- cname si.structName

let ns <- List.mapM cname si.fieldNames.toList

printTuple h True name ns

private def genInductive (h : Handle) (tc : Name) : MetaM Unit := do

let mut names : Array String := #[]

let tci <- getConstInfoInduct tc

for c in tci.ctors do

let ci <- getConstInfoCtor c

let name <- cname ci.name

names := names.push name

forallTelescopeReducing ci.type fun xs _ => do

let mut ns := []

for i in [:ci.numFields] do

let ld <- xs[ci.numParams + i]!.fvarId!.getDecl

ns := .cons ld.userName.toString ns

printTuple h False name ns.reverse

let rhs := String.intercalate " | " (.map .capitalize names.toList)

h.putStrLn s!"{<- cname tci.name} = {rhs}"

private def genPython (h : Handle) (name : Name) : MetaM Unit := do

h.putStrLn ""

match getStructureInfo? (<- getEnv) name with

| some si => genStructure h si

| none => genInductive h name

private def header :=

run_meta

let h <- IO.FS.Handle.mk "interop/nkl/lean_types.py" IO.FS.Mode.write

h.putStr header

flip List.forM (genPython h)

[ `NKL.Const

, `NKL.Expr

, `NKL.Stmt

, `NKL.Fun

]