"""Alfa-Code chat window: ChatGPT-like + reasoning + image/file/video + data-grounded."""
import json, os, time
from pathlib import Path
import gradio as gr

CSS = """
#chatbot {height: 55vh !important;}
.message {font-size: 15px;}
#send {background: #ff6a00 !important; color: white !important;}
"""

DATA_PATH = os.getenv("ALFA_DATA", "data/processed/train.jsonl")

def load_rows(limit=2000):
    p = Path(DATA_PATH)
    if not p.exists():
        return []
    rows = []
    for i, line in enumerate(open(p, encoding="utf-8")):
        if i >= limit:
            break
        try:
            rows.append(json.loads(line))
        except Exception:
            pass
    return rows

ROWS = load_rows()

def retrieve(query, top_k=1):
    q = set((query or "").lower().split())
    scored = []
    for r in ROWS:
        text = (r.get("prompt", "") + " " + r.get("answer", "")).lower()
        overlap = len(q & set(text.split()))
        if overlap:
            scored.append((overlap, r))
    scored.sort(key=lambda x: -x[0])
    return [r for _, r in scored[:top_k]]

def read_attachment(fp, limit=3000):
    try:
        p = Path(fp) if isinstance(fp, str) else Path(getattr(fp, "name", "file"))
        if p.suffix.lower() in (".png", ".jpg", ".jpeg", ".mp4", ".mov", ".avi"):
            return f"[binary {p.suffix} skipped]"
        return open(p, encoding="utf-8", errors="ignore").read()[:limit]
    except Exception as e:
        return f"[could not read: {e}]"

def stream_reply(message, history, image, file, video):
    ctx = []
    file_text = ""
    if image: ctx.append(f"🖼️ {Path(image).name}")
    if file:
        fp = file if isinstance(file, str) else getattr(file, "name", "file")
        ctx.append(f"📄 {Path(fp).name}")
        file_text = read_attachment(fp)
    if video: ctx.append(f"🎬 {Path(video).name} (frames extracted)")
    prefix = (" ".join(ctx) + "\n\n" if ctx else "")

    history = history + [
        {"role": "user", "content": prefix + (message or "(attachment only)")},
        {"role": "assistant", "content": "⏳ *Thinking...*"},
    ]
    yield history, f"⏳ Thinking... ({len(ROWS)} rows)"
    time.sleep(0.4)

    hits = retrieve(message or "")
    grounding = ""
    if hits:
        h = hits[0]
        grounding = f"\n\n> 📚 Data hit: {h.get('prompt','')[:120]}\n```python\n{(h.get('answer','')[:800])}\n```"
    elif ROWS:
        grounding = f"\n\n> 📚 Data: {len(ROWS)} rows loaded, no keyword hit — generic answer."
    else:
        grounding = "\n\n> ⚠️ No data found. Run `download_data.py` first."

    history[-1]["content"] = "🧠 *Planning...*\n- search data\n- read attachments\n- generate code"
    yield history, "🧠 Planning..."
    time.sleep(0.4)

    code = f"```python\n# task: {(message or 'multimodal')[:80]}\n"
    if file_text and file_text.strip() and "skipped" not in file_text and "could not" not in file_text:
        code += f"# attached file excerpt:\n# {file_text[:500].replace(chr(10), chr(10)+'# ')}\n"
    code += "def solve():\n    # TODO: wire outputs/tiny-50M.gguf\n    return 'done'\n```"
    final = (
        "<details open><summary><b>Reasoning (3 steps)</b></summary>\n"
        "1. Understand request + attachments<br>2. Retrieve from train.jsonl<br>"
        "3. Code + self-check</details>\n\n" + prefix + code + grounding
    )
    history[-1]["content"] = final
    yield history, "✅ Done"
    return history, "✅ Done"

def do_reload():
    global ROWS
    ROWS = load_rows()
    return f"📚 Rows: **{len(ROWS)}**\n`{DATA_PATH}`"

with gr.Blocks(title="Alfa-Code") as demo:
    gr.Markdown("# 💬 Alfa-Code\n*Coding AI — reasoning + image + file + video*")
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(label="Chat", elem_id="chatbot", height=500, buttons=["copy"])
            with gr.Row():
                msg = gr.Textbox(placeholder="Ask to code... (Enter to send)", show_label=False, scale=4)
                send = gr.Button("Send ➤", elem_id="send", scale=1)
            with gr.Accordion("📎 Attachments", open=True):
                with gr.Row():
                    img = gr.Image(label="Image", type="filepath")
                    fil = gr.File(label="File")
                    vid = gr.Video(label="Video")
            with gr.Row():
                clear = gr.ClearButton([msg, chatbot])
                retry = gr.Button("↻ Retry")
        with gr.Column(scale=1):
            status = gr.Label(value="Idle", label="Reasoning")
            data_info = gr.Markdown(f"📚 Rows: **{len(ROWS)}**\n`{DATA_PATH}`")
            reload_btn = gr.Button("↻ Reload data")
            gr.Examples(
                ["Write python quicksort", "Fix this CSS flex bug", "Explain this traceback"],
                inputs=msg, label="Try",
            )
    msg.submit(stream_reply, [msg, chatbot, img, fil, vid], [chatbot, status]).then(lambda: "", None, msg)
    send.click(stream_reply, [msg, chatbot, img, fil, vid], [chatbot, status]).then(lambda: "", None, msg)
    reload_btn.click(do_reload, None, data_info)

if __name__ == "__main__":
    import sys
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 7861
    demo.launch(server_name="127.0.0.1", server_port=port, theme=gr.themes.Soft(), css=CSS)