#!/usr/bin/env python3
"""Fix compilation errors and test again."""
import subprocess, os

TOKEN = "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"
subprocess.run(["git", "clone", f"https://{TOKEN}@github.com/ticketguy/Lila.git", "/app/lila"], check=True)
os.chdir("/app/lila")
subprocess.run(["git", "config", "user.name", "0xticketguy"], check=True)
subprocess.run(["git", "config", "user.email", "0xticketguy@harboria.dev"], check=True)

# Fix 1: model.c needs _GNU_SOURCE for madvise
with open("engine/runtime/model.c", "r") as f:
    content = f.read()
if "#define _GNU_SOURCE" not in content:
    content = "#define _GNU_SOURCE\n" + content
    with open("engine/runtime/model.c", "w") as f:
        f.write(content)

# Fix 2: attention.c calls dequant_matvec but it's static in inference.c
# Make dequant_matvec non-static and add proper declaration
with open("engine/runtime/inference.c", "r") as f:
    content = f.read()
content = content.replace("static void dequant_matvec(", "void dequant_matvec(")
content = content.replace("static void matvec(", "void matvec(")
with open("engine/runtime/inference.c", "w") as f:
    f.write(content)

# Fix 3: Remove the extern declaration at bottom of attention.c (conflicts)
with open("engine/runtime/attention.c", "r") as f:
    content = f.read()
content = content.replace("extern void dequant_matvec(float *out, const LilaQuantWeight *w, const float *vec);", "")
# Add forward declaration at top instead
if "void dequant_matvec(" not in content.split("void lila_attention")[0]:
    content = content.replace('#include "model.h"', '#include "model.h"\n\n/* From inference.c */\nextern void dequant_matvec(float *out, const LilaQuantWeight *w, const float *vec);')
with open("engine/runtime/attention.c", "w") as f:
    f.write(content)

# Fix 4: transformer.c also needs the extern declaration fixed
with open("engine/runtime/transformer.c", "r") as f:
    content = f.read()
# Remove bad extern and fix
content = content.replace("extern void lila_rmsnorm_avx2(float *out, const float *x, const float *weight, int size, float eps);", "")
content = content.replace("extern void lila_attention(float *output, const float *input, LilaLayer *layer,\n                           LilaKVCache *cache, int layer_idx, int position);", "")
content = content.replace("extern void dequant_matvec(float *out, const LilaQuantWeight *w, const float *vec);", "")
# Add clean declarations at top
header_block = """/* External declarations */
extern void lila_rmsnorm_avx2(float *out, const float *x, const float *weight, int size, float eps);
extern void lila_attention(float *output, const float *input, LilaLayer *layer,
                           LilaKVCache *cache, int layer_idx, int position);
extern void dequant_matvec(float *out, const LilaQuantWeight *w, const float *vec);
"""
if "External declarations" not in content:
    content = content.replace('#include "model.h"\n#include <math.h>', f'#include "model.h"\n#include <math.h>\n\n{header_block}')
with open("engine/runtime/transformer.c", "w") as f:
    f.write(content)

# Fix 5: Add CFLAGS -D_GNU_SOURCE to Makefile
with open("engine/Makefile", "r") as f:
    content = f.read()
content = content.replace("CFLAGS := -O3", "CFLAGS := -D_GNU_SOURCE -O3")
with open("engine/Makefile", "w") as f:
    f.write(content)

# Commit
subprocess.run(["git", "add", "-A"], check=True)
subprocess.run(["git", "commit", "-m", "Fix engine compilation errors\n\n- Add _GNU_SOURCE for madvise/MADV_SEQUENTIAL\n- Fix dequant_matvec linkage (non-static + extern declarations)\n- Fix transformer.c extern declarations\n- Add -D_GNU_SOURCE to CFLAGS"], check=True)
subprocess.run(["git", "push", "origin", "main"], check=True)
print("✅ Fixes pushed")

# Now try to compile
os.chdir("/app/lila/engine")
subprocess.run(["apt-get", "update", "-qq"], check=False, capture_output=True)
subprocess.run(["apt-get", "install", "-y", "-qq", "nasm"], check=False, capture_output=True)

print("\n=== Building engine ===")
result = subprocess.run(["make"], capture_output=True, text=True)
print(result.stdout[-500:] if result.stdout else "no stdout")
if result.returncode != 0:
    print("ERRORS:", result.stderr[-1000:])
    
    # Try with stubs
    print("\n=== Trying C-only with stubs ===")
    with open("/tmp/stubs.c", "w") as f:
        f.write('#include <stdint.h>\n#include <math.h>\n'
                'void lila_matvec_avx2(float *o, const float *m, const float *v, int r, int c) { for(int i=0;i<r;i++){float s=0;for(int j=0;j<c;j++)s+=m[i*c+j]*v[j];o[i]=s;}}\n'
                'void lila_rmsnorm_avx2(float *o, const float *x, const float *w, int sz, float e) { float ss=0;for(int i=0;i<sz;i++)ss+=x[i]*x[i];ss=1.0f/sqrtf(ss/sz+e);for(int i=0;i<sz;i++)o[i]=x[i]*ss*w[i];}\n'
                'void lila_dequant_int4_avx2(float *o, const uint8_t *i, const float *c, const float *s, int n, int g) {}\n'
                'void lila_print_cpu_features(void) { printf("CPU: x86_64 (stubs)\\n"); }\n')
    
    c_files = ["runtime/model.c", "runtime/inference.c", "runtime/attention.c",
               "runtime/transformer.c", "runtime/tokenizer.c", "runtime/dispatch.c",
               "interface/cli.c", "/tmp/stubs.c"]
    cmd = ["gcc", "-D_GNU_SOURCE", "-O2", "-Wall", "-std=c11", "-o", "lila-engine"] + c_files + ["-lm", "-lpthread", "-I", "runtime/"]
    r2 = subprocess.run(cmd, capture_output=True, text=True)
    if r2.returncode == 0:
        print("✅ ENGINE COMPILED (C with scalar stubs)")
        r3 = subprocess.run(["./lila-engine", "--test"], capture_output=True, text=True)
        print(f"Test: {r3.stdout.strip()}")
    else:
        print(f"❌ Still failing: {r2.stderr[:500]}")
else:
    print("✅ ENGINE COMPILED WITH ASSEMBLY KERNELS")
    r3 = subprocess.run(["./lila-engine", "--test"], capture_output=True, text=True)
    print(f"Test: {r3.stdout.strip()}")