import "cffc.hff";
import "common.hff";
import "vec.hff";
import "map.hff";

static outfp *FILE = {};

struct Value {
   ty *const Type,
   u enum union {
      IImm i64,
      FImm f64,
      BImm bool,
      Null,
      ZeroIni,
      StrLit [#]const u8,
      StrConstRef int,
      Tmp int,
      LocalRef *Var,
      GlobalRef *Var,
      Fn *Fn,
   }
}

static ty_i1 *Type = as(*void)&outfp; // dummy

fn gen(fmt *const u8, ...) void;
fn gen(fmt *const u8, ...) void {
   let ap va_list #?;
   ap->start(fmt);

   fn pritype(ty *const Type) void {
      if ty == ty_i1 {
         gen("i1");
         return;
      }
      switch ty.u {
      case Void; gen("void");
      case Bool; gen("i%z", ty.size*8);
      case Int; gen("i%z", ty.size*8);
      case Flo; gen(ty.size == 4 ? "float" : "double");
      case Ptr p; gen("%t*", p->is(:Void) ? ty_i8 : p);
      case Slice p; gen("{ %t*, %t }", p, ty_usize);
      case Arr arr; gen("[%z x %t]", arr.length, arr.child);
      case Agg agg; gen("%%%s.%d", agg.name ?? as(*const u8)"_", agg.id);
      case Enum enu; pritype(enu.intty);
      case Fn f;
         gen("%t(", f.ret);
         foreach(ty, i, f.params) {
            gen("%t", ty);
            if i < f.params.#len - 1 {
               gen(", ");
            } else if f.variadic {
               gen(", ...");
            }
         }
         gen(") ");
      }
   }

   for let c = *fmt; c != 0; c = *++fmt {
      if c != '%' {
         fputc(c, outfp);
         continue;
      }
      switch c = *++fmt {
      case '%';
         fprintf(outfp, "%%");
      case 'c';
         fprintf(outfp, "%c", ap->arg(int));
      case 's';
         fprintf(outfp, "%s", ap->arg(*const u8));
      case 'S';
         let str = ap->arg([#]const u8);
         fputc('"', outfp);
         foreach(c, _, str) {
            extern fn isprint(int) int;
            if isprint(c) == 0 {
               fprintf(outfp, "\\%.2X", c);
            } else {
               fputc(c, outfp);
            }
         }
         fprintf(outfp, "\\00\"");
      case 'd';
         fprintf(outfp, "%d", ap->arg(int));
      case 'I';
         fprintf(outfp, "%lld", as(c_llong)ap->arg(i64));
      case 'f';
         fprintf(outfp, "%.17f", ap->arg(f64));
      case 'z';
         fprintf(outfp, "%zu", ap->arg(usize));
      case 'v';
         switch ap->arg(Value).u {
         case IImm i; gen("%I", i);
         case FImm f; gen("%f", f);
         case BImm b; gen("%s", b ? "true" : "false");
         case Null; gen("null");
         case ZeroIni; gen("zeroinitializer");
         case Tmp t; fprintf(outfp, "%%t%d", t);
         case LocalRef var; gen("%%%s.%d", container_of(var, Decl, u.Let).name, var.id);
         case Fn f;
            let decl = container_of(f, Decl, u.Fn);
            if !decl.externp {
               gen("@%s.%d", decl.name, f.id);
            } else {
               gen("@%s", decl.name);
            }
         case GlobalRef v;
            let decl = container_of(v, Decl, u.Static);
            if !decl.externp {
               gen("@%s.%d", decl.name, v.id);
            } else {
               gen("@%s", decl.name);
            }
         case StrConstRef id;
            gen("@.str.%d", id);
            
         }
      case 't';
         pritype(ap->arg(*Type));

      case else
         assert(#f, "bad fmt '%c'", c);
      }
   }
   ap->end();
}

static tmpid int = 1;
fn mktmp(ty *const Type) Value {
   return Value{ty, :Tmp(tmpid++)};
}

fn genexpr(f *Fn, ex *Expr) Value;
static arena Arena = {};
static alloc Allocator = {};
static strs Vec<[#]const u8> = {};

fn genaddr(f *Fn, e *Expr) Value {
   switch e.u {
   case StrLit s;
      let p *u8 = alloc->alloc(s.#len + 1, 1);
      memcpy(p, s.#ptr, s.#len + 1);
      let s = p[0::s.#len];
      strs->push(s);
      return {mkptrtype(e.ty), :StrConstRef(strs.len - 1)};
   case Symbol decl;
      switch decl.u {
      case Let *var; return {var.ty, :LocalRef(var)};
      case Static *var; return {var.ty, :GlobalRef(var)};
      }
   }
   assert(#f, "genaddr");
}

fn convert(f *Fn, to *const Type, ex *Expr) Value {
   let from = ex.ty;

   defmacro cvt(inst) [
      (do let t = mktmp(to),
              val = genexpr(f, ex);
          gen("\t%v = %s %t %v to %t\n", t, inst, from, val, to);
          t;
      )
   ]
   switch {
   case from == to;
      return genexpr(f, ex);

   case to->is(:Int) and from->is(:Int) and to.size == from.size and to.u.Int.sgn != from.u.Int.sgn;
      return genexpr(f, ex);

   case to->is(:Int) and from->is(:Int) and to.size == from.size;
      return cvt("bitcast");

   case to->is(:Int) and from->is(:Int) and to.size < from.size;
      return cvt("trunc");

   case to->is(:Int) and from->is(:Int) and to.size > from.size;
      return cvt(from.u.Int.sgn ? "sext" : "zext");

   case to->is(:Int) and from->is(:Int) and to.size < from.size;
      return cvt("trunc");

   case to->is(:Int) and from->is(:Int) and to.size > from.size;
      return cvt(from.u.Int.sgn ? "sext" : "zext");

   case to->is(:Flo) and from->is(:Flo) and to.size == from.size;
      return genexpr(f, ex);

   case to->is(:Flo) and from->is(:Flo) and to.size > from.size;
      return cvt("fpext");

   case to->is(:Flo) and from->is(:Flo) and to.size < from.size;
      return cvt("fptrunc");

   case to->is(:Ptr) and from->is(:Ptr);
      let t = mktmp(to),
          val = genexpr(f, ex);
      gen("\t%v = bitcast %t %v to %t\n", t, from, val, to);
      return t;
   
   case to->is(:Ptr) and from->is(:Arr);
      let addr = genaddr(f, ex);
      let t = mktmp(ex.ty);
      gen("\t%v = getelementptr %t, %t %v, %t 0, %t 0\n", t, addr.ty.u.Ptr, addr.ty, addr,
                                                          ty_usize, ty_usize);
      return t;   
   case else
      assert(#f, "convert");
   }
}

fn genblock(f *Fn, block [#]Stmt) void;

fn genexpr(f *Fn, ex *Expr) Value {
   fold(ex);
   switch ex.u {
   case IntLit i;  return {ex.ty, :IImm(i.i)};
   case FloLit f;  return {ex.ty, :FImm(f)};
   case BoolLit b; return {ex.ty, :BImm(b)};
   case NullLit;   return {ex.ty, :Null};
   case ZeroIni;  return {ex.ty, :ZeroIni};
   case StrLit s; return {ex.ty, :StrLit(s)};
   case BinOp b;
      defmacro genbinop(inst) [ {
         let lhs = convert(f, ex.ty, b.lhs),
             rhs = convert(f, ex.ty, b.rhs),
             t = mktmp(ex.ty);
         gen("\t%v = %s %t %v, %v\n", t, inst, ex.ty, lhs, rhs);
         return t;
      } ]
      let ty2 = typeof2(b.lhs.ty, b.rhs.ty);
      defmacro gencmp(op) [ {
         let lhs = convert(f, ty2, b.lhs),
             rhs = convert(f, ty2, b.rhs),
             t0 = mktmp(ex.ty),
             t1 = mktmp(ex.ty);
         gen("\t%v = %s %s %t %v, %v\n", t0, ty2->is(:Flo) ? "fcmp" : "icmp", op,
                                  ty2, lhs, rhs);
         gen("\t%v = zext i1 %v to %t\n", t1, t0, ex.ty);
         return t1;
      } ]
      switch b.op {
      case  '+';
      // TODO ptr arithmetic
         genbinop(!ex.ty->is(:Flo) ? "add" : "fadd");
      case  '-';
         genbinop(!ex.ty->is(:Flo) ? "sub" : "fsub");
      case  '*'; genbinop(!ex.ty->is(:Flo) ? "mul" : "fmul");
      case  '/'; genbinop(ex.ty->is(:Flo) ? "fdiv" : ex.ty.u.Int.sgn ? "sdiv" : "udiv");
      case  '%'; genbinop(ex.ty.u.Int.sgn ? "srem" : "urem");
      case  '&'; genbinop("and");
      case  '|'; genbinop("or");
      case  '^'; genbinop("or");
      case '<<'; genbinop("shl");
      case '>>'; genbinop(ex.ty.u.Int.sgn ? "ashr" : "lshr");
      case '=='; gencmp("eq");
      case  '<'; gencmp(ty2->is(:Flo) ? "olt" : ty2.u.Int.sgn ? "slt" : "ult");
      case '<='; gencmp(ty2->is(:Flo) ? "ole" : ty2.u.Int.sgn ? "sle" : "ule");
      case  '>'; gencmp(ty2->is(:Flo) ? "ogt" : ty2.u.Int.sgn ? "sgt" : "ugt");
      case '>='; gencmp(ty2->is(:Flo) ? "oge" : ty2.u.Int.sgn ? "sge" : "uge");
      case else
         assert(#f, "binop? %d", b.op);
      }
   case UnOp un;
      switch un.op {
      case :preinc, :predec;
         let one = Value{ex.ty, :IImm(un.op == :preinc ? 1 : -1)};
         let addr = genaddr(f, un.ex);
         let var = mktmp(ex.ty);
         let val = mktmp(ex.ty);
         gen("\t%v = load %t, %t* %v\n", var, ex.ty, ex.ty, addr);
         gen("\t%v = add %t %v, %v\n", val, ex.ty, var, one);
         gen("\tstore %t %v, %t* %v\n", ex.ty, val, ex.ty, addr);
         return val;
      case else
         assert(#f, "unop?");
      }

   case Symbol decl;
      switch decl.u {
      case Let *var;
         let t = mktmp(ex.ty);
         gen("\t%v = load %t, %t* %%%s.%d\n", t, ex.ty, ex.ty, decl.name, var.id);
         return t;
      case Static *var;
         let t = mktmp(ex.ty);
         if decl.externp {
            gen("\t%v = load %t, %t* @%s\n", t, ex.ty, ex.ty, decl.name);
         } else {
            gen("\t%v = load %t, %t* @%s.%d\n", t, ex.ty, ex.ty, decl.name, var.id);
         }
         return t;
      case Fn *f;
         return {ex.ty, :Fn(f)};
      case else
         assert(#f, "decl");
      }
   case Call call;
      let lhs = genexpr(f, call.lhs);
      let fnty = &(lhs.ty->is(:Ptr) ? lhs.ty.u.Ptr : lhs.ty).u.Fn;
      let args *Value = xcalloc(call.args.#len, sizeof Value);
      defer free(args);
      foreach_ptr(arg, i, call.args) {
         let ty = i < fnty.params.#len ? fnty.params[i] : typeof2(arg.ty, arg.ty);
         args[i] = convert(f, ty, arg);
      }
      let t = mktmp(ex.ty);
      gen("\t%v = call %t %v(", t, call.lhs.ty, lhs);
      foreach_ptr(arg, i, call.args) {
         let ty = i < fnty.params.#len ? fnty.params[i] : typeof2(arg.ty, arg.ty);
         gen("%t %v", ty, args[i]);
         if i < call.args.#len - 1 {
            gen(", ");
         }
      }
      gen(")\n");
      return t;
      
   case else
      assert(#f, "expr? %d", ex.u.#tag);
   }
}

fn llvmbool(f *Fn, ex *Expr) Value {
   let v = genexpr(f, ex);
   let t = mktmp(ty_i1);
   if v.ty->is(:Bool) or v.ty->is(:Int) {
      gen("\t%v = icmp ne %t %v, 0\n", t, v.ty, v);
   } else {
      gen("\t%v = icmp ne %t %v, null\n", t, v.ty, v);
   }
   return t;
}

fn nop() void {
   gen("\t%%t%d = bitcast i8 0 to i8 ; NOP\n", tmpid++);
}

fn genstmt(f *Fn, st *Stmt) void {
   switch st.u {
   case Expr *e;
      genexpr(f, e);

   case Decl decl;
      switch decl.u {
      case Let var;
         let nam = decl.name,
             ty = var.ty;
         gen("\t%%%s.%d = alloca %t\n", nam, var.id, ty);
         if !var.ini->empty() {
            gen("\tstore %t %v, %t* %%%s.%d\n",
                ty, convert(f, ty, &var.ini.Some), ty, nam, var.id);
         }
      }
   case If *cnd;
      static ifid int = {};
      let id = ifid++;
      gen("\tbr i1 %v, label %%IfT%d, label %%IfF%d\n",
          llvmbool(f, &cnd.test), id, id);
      gen("IfT%d:", id);
      nop();
      genblock(f, cnd.t);
      if cnd.f.#ptr {
         gen("\tbr label %%IfSk%d\n", id);
      }
      gen("IfF%d:", id);
      nop();
      genblock(f, cnd.f);
      if cnd.f.#ptr {
         gen("IfSk%d:", id);
         nop();
      }
   case While loop;
      gen("Cont%d:", loop.id);
      nop();
      gen("\tbr i1 %v, label %%Next%d, label %%Brk%d\n",
          llvmbool(f, &loop.test), loop.id, loop.id);
      gen("Next%d:", loop.id);
      nop();
      genblock(f, loop.body);
      gen("\tbr label %%Cont%d\n", loop.id);
      gen("Brk%d:",loop.id);
      nop();
   case For loop;
      genblock(f, loop.ini);
      gen("\tbr label %%Cont%d\n", loop.id);
      gen("Cont%d:", loop.id);
      nop();
      gen("\tbr i1 %v, label %%Next%d, label %%Brk%d\n",
          llvmbool(f, &loop.test), loop.id, loop.id);
      gen("Next%d:", loop.id);
      nop();
      genblock(f, loop.body);
      if !loop.next->empty() {
         genexpr(f, &loop.next.Some);
      }
      gen("\tbr label %%Cont%d\n", loop.id);
      gen("Brk%d:",loop.id);
      nop();
   case Return *e;
      switch *e {
      case None;
         gen("\tret void\n");
      case Some *e;
         let retty = f.ty.u.Fn.ret;
         gen("\tret %t %v\n", retty, convert(f, retty, e));
      }
   case else
      assert(#f, "stmt? %d", st.u.#tag);
   }
}

fn genblock(f *Fn, block [#]Stmt) void {
   foreach_ptr (st, _, block) {
      genstmt(f, st);
   }
}

struct FuncKey {
   name *const u8,
   externp bool,
}
struct FuncEntry {
   t enum {
      Idk, Def, Decl
   },
   ty *const Type,
   id int,
}
static funcs Map<FuncKey, FuncEntry, struct {
   fn hash(e FuncKey) u32 {
      let h = fnv1a_s(FNV1A_INI, e.name);
      h = fnv1a_i(h, as(int)e.externp);
      return h;
   }
   fn eq(a FuncKey, b FuncKey) bool {
      return streq(a.name, b.name) and a.externp == b.externp;
   }
}> = {};

extern fn llvm_addfn(decl *Decl) void {
   assert(decl.toplevel, "llvm-addfn");
   let slot = funcs->get_slot({decl.name, decl.externp});
   if slot.t != :Def {
      slot.t = :Decl;
      slot.ty = decl.u.Fn.ty;
      slot.id = decl.u.Fn.id;
   }
}

extern fn llvm_genfn(externp bool, name *const u8, f *Fn) void {
   funcs->put({name, externp}, {:Def, f.ty, f.id});
   tmpid = 0;
   gen("define%s %t ", externp ? "" : " internal", f.ty.u.Fn.ret);
   if !externp {
      gen("@%s.%d(", name, f.id);
   } else {
      gen("@%s(", name);
   }
   foreach(ty, i, f.ty.u.Fn.params) {
      gen("%t", ty);
      if f.paramnames[i] {
         gen(" %%%s", f.paramnames[i]);
      }
      if i < f.paramnames.#len - 1 {
         gen(", ");
      }
   }
   gen(") { \n");
   let id = 0;
   foreach(nam, i, f.paramnames) {
      if nam {
         let ty = f.ty.u.Fn.params[i];
         gen("%%%s.%d = alloca %t   ", nam, id, ty);
         gen("store %t %%%s, %t* %%%s.%d\n", ty, nam, ty, nam, id);
         ++id;
      }
   }
   genblock(f, f.body.Some);
   if f.ty.u.Fn.ret->is(:Void) {
      gen("\tret void\n");
   } else {
      let retty = f.ty.u.Fn.ret;
      gen("\tret %t undef\n", retty);
   }
   gen("}\n");
}

extern fn llvm_addtype(ty *const Type) void {
   if ty.konst { return; }
   switch ty.u {
   case Agg agg;
      assert(agg.kind == :Struct, "addtype");
      if agg.fwd {
         gen("%%%s.%d = type opaque\n", agg.name, agg.id);
      } else {
      }
      
   }
}

extern fn llvm_fini() void {
   vec_each(s, i, strs) {
      gen("@.str.%z = internal constant [%z x i8] c%S;\n", i, s.#len + 1, s);
   }
   map_each(f, k, funcs) {
      if f.t == :Decl {
         let ty = f.ty;
         gen("declare %t ", ty.u.Fn.ret);
         if !k.externp {
            gen("@%s.%d(", k.name, f.id);
         } else {
            gen("@%s(", k.name);
         }
         foreach(tyy, i, ty.u.Fn.params) {
            gen("%t", tyy);
            if i < ty.u.Fn.params.#len - 1 {
               gen(", ");
            } else if ty.u.Fn.variadic {
               gen(", ...");
            }
         }
         gen(");\n");
      }
   }
   strs->clear();
   funcs->clear();
   arena->destroy();
}

extern fn llvm_init(out *FILE) void {
   alloc = {&arena, &Arena:allocf, #null};
   outfp = out;
}