import os, shutil, subprocess, threading, uuid, time, zipfile, gzip, glob
import gradio as gr
from transformers import pipeline

LOG_FILE = "train.log"
MODEL_DIR = "trained_model"
ZIP_FILE = "trained_model.zip"
ZIP_PART = ZIP_FILE + ".part"

def _human(n):
    u=["B","KB","MB","GB"]; i=0; x=float(n)
    while x>=1024 and i<len(u)-1: x/=1024; i+=1
    return f"{x:.1f} {u[i]}"

def _read(path, fb="Waiting..."):
    try:
        with open(path,"r",encoding="utf-8",errors="ignore") as f: return f.read()
    except: return fb

def _zip_dir_atomic(src, out_path, tmp_path):
    if os.path.exists(tmp_path): os.remove(tmp_path)
    with zipfile.ZipFile(tmp_path,"w",zipfile.ZIP_DEFLATED) as z:
        for root,_,files in os.walk(src):
            for fn in files:
                fp = os.path.join(root, fn)
                z.write(fp, arcname=os.path.relpath(fp, src))
    if os.path.exists(out_path): os.remove(out_path)
    os.replace(tmp_path, out_path)

def upload_file(f):
    if f is None: return "❌ No file.", ""
    os.makedirs("uploads", exist_ok=True)
    dst = os.path.join("uploads", f"dataset_{uuid.uuid4().hex}.jsonl")
    shutil.copy(f.name, dst)
    return f"✅ Uploaded → {dst}", dst

def _train_single(dataset, log, quick, steps, subset):
    args = ["python","train.py","--dataset",dataset,"--output",MODEL_DIR]
    if quick:
        args += ["--quick","1","--max_steps",str(steps or 8),"--subset",str(subset or 32)]
    p = subprocess.Popen(args, stdout=log, stderr=subprocess.STDOUT)
    p.wait()
    log.write(f"\n    ↳ train.py exited {p.returncode} for {os.path.basename(dataset)}\n")
    return p.returncode == 0

def _worker(dataset_path, shards_folder, quick, steps, subset):
    with open(LOG_FILE,"w") as log: log.write("🔥 Starting training (C# AI)…\n")
    ok=True
    with open(LOG_FILE,"a") as log:
        if shards_folder:
            log.write(f"📂 Folder mode: {shards_folder}\n")
            paths = sorted(glob.glob(os.path.join(shards_folder,"*.jsonl"))) + \
                    sorted(glob.glob(os.path.join(shards_folder,"*.jsonl.gz")))
            paths = [p for p in paths if "manifest" not in os.path.basename(p).lower()]
            if not paths:
                log.write("❌ No shards (*.jsonl / *.jsonl.gz).\n"); ok=False
            else:
                tmp="tmp_train.jsonl"
                for i,pth in enumerate(paths,1):
                    log.write(f"\n[{i}/{len(paths)}] Shard: {os.path.basename(pth)}\n")
                    if pth.endswith(".gz"):
                        try:
                            with gzip.open(pth,"rt",encoding="utf-8") as rf, open(tmp,"w",encoding="utf-8") as wf:
                                for line in rf: wf.write(line)
                            shard = tmp
                        except Exception as e:
                            log.write(f"❌ GZ read failed: {e}\n"); ok=False; break
                    else:
                        shard = pth
                    if not _train_single(shard, log, quick, steps, subset): ok=False; break
                if os.path.exists(tmp):
                    try: os.remove(tmp)
                    except: pass
        else:
            if not dataset_path or not os.path.exists(dataset_path):
                log.write("❌ Upload a valid dataset.\n"); ok=False
            else:
                ok = _train_single(dataset_path, log, quick, steps, subset)

        if ok and os.path.isdir(MODEL_DIR):
            try:
                _zip_dir_atomic(MODEL_DIR, ZIP_FILE, ZIP_PART)
                sz = _human(os.path.getsize(ZIP_FILE))
                log.write(f"\n✅ Model zipped → {ZIP_FILE} ({sz})\n")
            except Exception as e:
                log.write(f"\n❌ Zip failed: {e}\n")
        else:
            log.write("\n❌ Training failed; no zip.\n")

def start_training(dataset_path, shards_folder, quick, steps, subset):
    for p in (ZIP_FILE, ZIP_PART):
        if os.path.exists(p):
            try: os.remove(p)
            except: pass
    threading.Thread(target=_worker, args=(dataset_path, shards_folder, quick, steps, subset), daemon=True).start()
    return "🚀 Training started. Use Refresh buttons."

def read_logs(): return _read(LOG_FILE, "Waiting for logs...")

def refresh_download():
    if os.path.exists(ZIP_FILE):
        size=_human(os.path.getsize(ZIP_FILE))
        return gr.update(visible=True, value=ZIP_FILE), f"*Ready:* {ZIP_FILE} • {size}"
    return gr.update(visible=False, value=None), "No trained model yet."

def load_test_zip(z):
    if z is None: return "❌ No file.", ""
    import zipfile
    root = os.path.join("models", f"test_{uuid.uuid4().hex}")
    os.makedirs(root, exist_ok=True)
    try:
        with zipfile.ZipFile(z.name,"r") as zz: zz.extractall(root)
        return f"✅ Extracted to {root}", root
    except Exception as e:
        return f"❌ Extract failed: {e}", ""

def clear_test_model(): return "Cleared. Will use trained_model/ if present.", ""

def generate(prompt, model_path):
    if not prompt.strip(): return "Enter a prompt."
    try:
        if model_path and os.path.isdir(model_path): m, src = model_path, "(uploaded)"
        elif os.path.isdir(MODEL_DIR):              m, src = MODEL_DIR, "(trained_model/)"
        else:                                       m, src = "sshleifer/tiny-gpt2", "(tiny fallback)"
        gen = pipeline("text-generation", model=m, tokenizer="sshleifer/tiny-gpt2")
        out = gen(prompt, max_length=120, do_sample=True, temperature=0.7, truncation=True)[0]["generated_text"]
        return f"{out}\n\n— using {src}"
    except Exception as e:
        return f"❌ Error: {e}"

with gr.Blocks(title="C# AI Trainer (Quick Test Mode)") as app:
    gr.Markdown("## ⚡ C# AI Trainer — includes *Quick Test* (few steps on a tiny model).")

    ds_state = gr.State("")
    folder_state = gr.State("")
    test_model_state = gr.State("")

    with gr.Tab("🧠 Train"):
        with gr.Row():
            file_in = gr.File(label="Upload dataset (.jsonl)", file_types=[".jsonl"])
            up_btn = gr.Button("📤 Upload")
        with gr.Row():
            shards_folder = gr.Textbox(value="", label="Folder with shards (optional)")
            use_folder = gr.Button("📂 Use Folder")
        with gr.Row():
            quick = gr.Checkbox(value=True, label="⚡ Quick Test (tiny model, few steps)")
            steps = gr.Number(value=8, label="Max steps (Quick mode)", precision=0)
            subset = gr.Number(value=32, label="Subset rows (Quick mode)", precision=0)
        status = gr.Textbox(label="Status", interactive=False)
        with gr.Row():
            start_btn = gr.Button("🚀 Start Training")
            refresh_logs = gr.Button("🔁 Refresh Logs")
            refresh_dl = gr.Button("📦 Refresh Download Area")
        logs = gr.Textbox(label="📜 Logs", lines=18)
        dl_btn = gr.DownloadButton(label="📥 Download Trained Model (.zip)", visible=False, value=None)
        dl_info = gr.Markdown("No trained model yet.")

        up_btn.click(fn=upload_file, inputs=file_in, outputs=[status, ds_state])
        use_folder.click(fn=lambda p: ("✅ Using folder." if p.strip() else "❌ Provide folder path.", p.strip()),
                         inputs=shards_folder, outputs=[status, folder_state])
        start_btn.click(fn=start_training, inputs=[ds_state, folder_state, quick, steps, subset], outputs=status
        ).then(fn=read_logs, outputs=logs
        ).then(fn=refresh_download, outputs=[dl_btn, dl_info])
        refresh_logs.click(fn=read_logs, outputs=logs)
        refresh_dl.click(fn=refresh_download, outputs=[dl_btn, dl_info])

    with gr.Tab("🚀 Test"):
        with gr.Row():
            zip_in = gr.File(label="Upload model ZIP", file_types=[".zip"])
            load_btn = gr.Button("📦 Load ZIP")
            clear_btn = gr.Button("🧹 Clear")
        test_status = gr.Textbox(label="Test Model Status", interactive=False)
        prompt = gr.Textbox(label="Prompt", placeholder="e.g., Write a C# method that reverses a string.")
        go = gr.Button("🔍 Generate")
        out = gr.Textbox(label="AI Response", lines=12)

        load_btn.click(fn=load_test_zip, inputs=zip_in, outputs=[test_status, test_model_state])
        clear_btn.click(fn=clear_test_model, outputs=[test_status, test_model_state])
        go.click(fn=generate, inputs=[prompt, test_model_state], outputs=out)

app.launch()