import sys, os, json, torch, requests, time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from fsi_echo import FSIEchoModel, CodeTokenizer, ClosedLoopDebugger
import gradio as gr

@torch.no_grad()
def load():
    MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'prod2_final.pt')
    if not os.path.exists(MODEL_PATH):
        url = "https://huggingface.co/FerrellSyntheticIntelligence/FSI_ECHO/resolve/main/prod2_final.pt"
        r = requests.get(url, stream=True, timeout=300)
        r.raise_for_status()
        with open(MODEL_PATH, 'wb') as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
    ckpt = torch.load(MODEL_PATH, map_location='cpu', weights_only=True)
    m = FSIEchoModel()
    m.load_state_dict(ckpt['model'])
    m.eval()
    return m, CodeTokenizer()

model, tok = load()
debugger = ClosedLoopDebugger(model, tok)

def generate(prompt, temp, max_t):
    t0 = time.time()
    r = model.generate(tok, prompt, max_tokens=int(max_t), temperature=float(temp), top_k=5, top_p=0.95)
    return r['generated'], f"{r['tokens']}t | conf:{r['confidence']:.2f} | {time.time()-t0:.1f}s"

def debug(code):
    r = debugger.debug(code)
    return r.get('code', 'Could not fix')

with gr.Blocks(title="FSI_ECHO", theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("# 🧬 FSI_ECHO — 2.6M Param Code AI\n*Morphing Code Swarm*")
    with gr.Tab("Generate"):
        p = gr.Textbox(label="Prompt", value="def is_even")
        t = gr.Slider(0.1, 1.0, 0.1, label="Temperature")
        m = gr.Slider(10, 200, 50, label="Max Tokens")
        b = gr.Button("Generate")
        o = gr.Textbox(label="Output", lines=8)
        s = gr.Textbox(label="Stats")
        b.click(generate, [p, t, m], [o, s])
    with gr.Tab("Debug"):
        c = gr.Textbox(label="Buggy Code", value="def add(a, b):\n    a + b", lines=6)
        b2 = gr.Button("Debug")
        o2 = gr.Textbox(label="Fixed Code", lines=8)
        b2.click(debug, [c], [o2])

if __name__ == '__main__':
    demo.launch(server_name='0.0.0.0', server_port=7860)