import json
import os
import random
import re
import shutil
import subprocess
import sys
import threading
import time

import httpx
from prompt_toolkit import PromptSession, prompt as pt_prompt
from prompt_toolkit.application import get_app
from prompt_toolkit.completion import Completer, Completion
from prompt_toolkit.document import Document
from prompt_toolkit.formatted_text import HTML, FormattedText
from prompt_toolkit import print_formatted_text as pt_print
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.styles import Style
from rich.console import Console
from rich.text import Text
from rich.theme import Theme

BASE_URL = os.getenv("KEDY_BACKEND_URL", "http://localhost:8000")
VERSION = "26.3.1"
SESSION_PATH = os.path.join(os.path.expanduser("~"), ".kedy", "session.json")

SALMON = "#FA8072"

theme = Theme({
    "salmon": SALMON,
    "info": SALMON,
    "success": "bold green",
    "warning": "bold yellow",
    "error": "bold red",
    "dim": "dim white",
})

console = Console(theme=theme)

PT_STYLE = Style.from_dict({
    "prompt":                                "#FA8072",
    "placeholder":                           "italic #555555",
    "bottom-toolbar":                        "noinherit bg:default",
    "bottom-toolbar.text":                   "noinherit bg:default",
    "completion-menu":                       "noinherit bg:default",
    "completion-menu.completion":            "noinherit bg:default #aaaaaa",
    "completion-menu.completion.current":    "noinherit bg:default bold #FA8072",
    "completion-menu.progressbar":           "noinherit bg:default",
    "completion-menu.progressbar.used":      "noinherit bg:default",
    "scrollbar":                             "noinherit bg:default",
    "scrollbar.background":                  "noinherit bg:default",
    "scrollbar.button":                      "noinherit bg:default",
})

LOGIN_STYLE = Style.from_dict({
    "prompt":                                "#FA8072",
    "placeholder":                           "italic #555555",
})

HINTS = [
    "/generate turkish surnames with city and age",
    "type / to see all commands",
    "/generate e-commerce product reviews",
    "/generate hospital patient records",
]

COMMANDS = {
    "/help": "show available commands",
    "/generate": "generate a dataset (e.g. /generate turkish surnames)",
    "/train": "train a model on project datasets",
    "/models": "list trained models in current project",
    "/deploy": "deploy a model and get an API key",
    "/info": "show project status, datasets, and models",
    "/servers": "manage servers",
    "/projects": "manage projects",
    "/logout": "sign out and exit",
    "/exit": "quit",
}

# Active project state (set when user enters a project)
active_project: dict | None = None
active_project_index: int = -1


# ──────────────────────────────────────────────
#  Session persistence
# ──────────────────────────────────────────────

def load_session() -> dict | None:
    try:
        with open(SESSION_PATH) as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return None


def save_session(data: dict):
    os.makedirs(os.path.dirname(SESSION_PATH), exist_ok=True)
    with open(SESSION_PATH, "w") as f:
        json.dump(data, f, indent=2)


def clear_session():
    try:
        os.remove(SESSION_PATH)
    except FileNotFoundError:
        pass


# ──────────────────────────────────────────────
#  Completers
# ──────────────────────────────────────────────

class OptionCompleter(Completer):
    def __init__(self, options):
        self.options = options

    def get_completions(self, document: Document, complete_event):
        text = document.text_before_cursor.lower()
        for i, opt in enumerate(self.options):
            label = f"{i + 1}. {opt}"
            if not text or opt.lower().startswith(text) or str(i + 1).startswith(text):
                yield Completion(opt, start_position=-len(document.text_before_cursor), display=label)


class SlashCompleter(Completer):
    def get_completions(self, document: Document, complete_event):
        text = document.text_before_cursor
        if not text.startswith("/"):
            return
        for cmd, desc in COMMANDS.items():
            if cmd.startswith(text):
                yield Completion(cmd, start_position=-len(text), display=cmd + "  " + desc)


# ──────────────────────────────────────────────
#  UI helpers
# ──────────────────────────────────────────────

def rule():
    w = shutil.get_terminal_size().columns
    pt_print(FormattedText([("#333333", "─" * w)]))


def thin_rule():
    w = shutil.get_terminal_size().columns
    console.print(f"  {'─' * (w - 4)}", style="#333333")


def print_banner(session_data: dict):
    home = os.path.expanduser("~")
    cwd = os.getcwd().replace(home, "~")
    team = session_data.get("team", {})
    team_name = team.get("name", "Team")
    servers = team.get("servers", [])
    server_count = len(servers)

    console.print()
    cat_lines = [r" /\_/\ ", r"( o.o )", r" > ^ < "]
    third_line = [
        (f"{server_count} server{'s' if server_count != 1 else ''}", SALMON),
        ("  ", ""),
    ]
    if active_project:
        third_line.append((f"{active_project['name']}", f"bold {SALMON}"))
        third_line.append(("  ", ""))
        third_line.append((active_project.get("server_name", ""), "#777777"))
    else:
        third_line.append((cwd, "#777777"))

    right_lines = [
        [("Kedy ", "bold white"), (f"v{VERSION}", "#777777")],
        [(f"{team_name} ", "bold white"), ("· ", SALMON), ("Kedy Pro", f"bold {SALMON}")],
        third_line,
    ]

    for i, cat_line in enumerate(cat_lines):
        line = Text()
        line.append(f"  {cat_line}", style=SALMON)
        line.append("   ")
        for text, style in right_lines[i]:
            line.append(text, style=style)
        console.print(line)
    console.print()


def print_response(text):
    clean = re.sub(r"\[TOOL:\s*[^\]]+\]", "", text).strip()
    if clean:
        console.print()
        console.print(f"  {clean}", style="dim")
        console.print()


def fmt_elapsed(seconds: float) -> str:
    if seconds < 60:
        return f"{seconds:.1f}s"
    m, s = divmod(int(seconds), 60)
    return f"{m}m {s}s"


def print_rows_summary(rows, mode="add"):
    if not rows:
        return
    n = len(rows)
    if mode == "add":
        keys = ", ".join(rows[0].keys()) if isinstance(rows[0], dict) else ""
        console.print(f"  [bold green]+[/bold green] {n} rows generated", end="")
        if keys:
            console.print(f"  [dim]({keys})[/dim]")
        else:
            console.print()
    else:
        console.print(f"  [bold red]-[/bold red] {n} rows removed")


def print_validation(validation: dict, total_rows: int):
    if not validation:
        return
    issues = validation.get("issues", [])
    if not issues:
        console.print(f"  [bold green]\u2713[/bold green] validated")
    else:
        console.print(f"  [bold red]-[/bold red] {len(issues)} issues found")


def create_client() -> httpx.Client:
    return httpx.Client(base_url=BASE_URL, timeout=120)


# ──────────────────────────────────────────────
#  Welcome & Login
# ──────────────────────────────────────────────

def show_welcome():
    console.print()
    console.print()

    w = shutil.get_terminal_size().columns
    pad = max((w - 36) // 2, 2)
    sp = " " * pad

    cat_lines = [r" /\_/\ ", r"( o.o )", r" > ^ < "]
    for cat_line in cat_lines:
        line = Text()
        line.append(f"{sp}{cat_line}", style=SALMON)
        console.print(line)

    console.print()
    console.print(f"{sp}[bold white]Welcome to Kedy[/bold white]  [{SALMON}]v{VERSION}[/{SALMON}]")
    console.print(f"{sp}[#999999]Sign in to continue[/#999999]")
    console.print()
    thin_rule()
    console.print()


def do_login(client: httpx.Client) -> dict | None:
    """Attempt cached session first, then interactive login."""

    # Try cached session
    cached = load_session()
    if cached and cached.get("email") and cached.get("key"):
        try:
            r = client.post("/auth", json={"email": cached["email"], "key": cached["key"]})
            r.raise_for_status()
            data = r.json()
            if data.get("ok"):
                user = data["user"]
                team = data.get("team", {})
                console.print()
                console.print(f"  [{SALMON}]Signed in as[/{SALMON}] [bold white]{user['name']}[/bold white]")
                console.print()
                return {
                    "user": user,
                    "team": team,
                }
        except Exception:
            pass
        # Cached session invalid, clear it
        clear_session()

    # Interactive login
    show_welcome()

    try:
        email = pt_prompt(
            HTML("<prompt>  email  </prompt>"),
            style=LOGIN_STYLE,
        ).strip()
        if not email:
            return None

        key = pt_prompt(
            HTML("<prompt>  key    </prompt>"),
            style=LOGIN_STYLE,
            is_password=True,
        ).strip()
        if not key:
            return None
    except (KeyboardInterrupt, EOFError):
        return None

    try:
        r = client.post("/auth", json={"email": email, "key": key})
        r.raise_for_status()
        data = r.json()
    except Exception as e:
        console.print(f"\n  [bold red]-[/bold red] Connection error: {e}")
        return None

    if not data.get("ok"):
        console.print(f"\n  [bold red]-[/bold red] {data.get('error', 'Authentication failed')}")
        return None

    user = data["user"]
    team = data.get("team", {})

    # Persist session
    save_session({"email": email, "key": key})

    console.print()
    console.print(f"  [bold green]+[/bold green] Signed in as [bold white]{user['name']}[/bold white]  [{SALMON}]{user['email']}[/{SALMON}]")
    console.print(f"  [#999999]{user['user_type']}  ·  {team.get('name', user['team'])}[/#999999]")
    console.print()

    return {
        "user": user,
        "team": team,
    }


# ──────────────────────────────────────────────
#  Admin Setup (first-time only)
# ──────────────────────────────────────────────

def run_admin_setup(client: httpx.Client, team_key: str):
    console.print()
    thin_rule()
    console.print()

    cat_lines = [r" /\_/\ ", r"( o.o )", r" > ^ < "]
    setup_right = [
        [("Kedy ", "bold white"), ("Setup", f"bold {SALMON}")],
        [(f"v{VERSION}", "#777777")],
        [("Configure your workspace", "#999999")],
    ]
    for i, cat_line in enumerate(cat_lines):
        line = Text()
        line.append(f"  {cat_line}", style=SALMON)
        line.append("   ")
        for text, style in setup_right[i]:
            line.append(text, style=style)
        console.print(line)

    console.print()
    thin_rule()
    console.print()
    console.print(f"  [bold white]This setup will walk you through two steps to get your[/bold white]")
    console.print(f"  [bold white]team up and running.[/bold white]")
    console.print()
    console.print(f"  [{SALMON}]1.[/{SALMON}]  [bold white]Team[/bold white]     Add developers who will use this workspace.")
    console.print(f"  [#999999]              Each member gets an email, name, and access key.[/#999999]")
    console.print()
    console.print(f"  [{SALMON}]2.[/{SALMON}]  [bold white]Servers[/bold white]  Register the machines where Kedy agents will run.")
    console.print(f"  [#999999]              Agents handle data generation across your infra.[/#999999]")
    console.print()
    console.print(f"  [#777777]You can update these later. Press enter to skip any section.[/#777777]")
    console.print()
    thin_rule()

    # ── Team members ──
    console.print()
    console.print(f"  [bold white]Team Members[/bold white]  [{SALMON}]max 10[/{SALMON}]")
    console.print()

    members = []
    for i in range(10):
        try:
            email = pt_prompt(
                HTML(f"<prompt>  [{i + 1}] email  </prompt>"),
                style=LOGIN_STYLE,
            ).strip()
            if not email:
                break

            name = pt_prompt(
                HTML(f"<prompt>  [{i + 1}] name   </prompt>"),
                style=LOGIN_STYLE,
            ).strip()

            key = pt_prompt(
                HTML(f"<prompt>  [{i + 1}] key    </prompt>"),
                style=LOGIN_STYLE,
                is_password=True,
            ).strip()
            if not key:
                console.print("  [#444444]skipped[/#444444]")
                continue

            members.append({"email": email, "name": name or email.split("@")[0], "key": key})
            console.print(f"  [bold green]+[/bold green] [bold white]{name or email}[/bold white]")
            console.print()
        except (KeyboardInterrupt, EOFError):
            break

    if members:
        console.print(f"  [{SALMON}]{len(members)} member(s) added[/{SALMON}]")
    console.print()

    # ── Servers ──
    thin_rule()
    console.print()
    console.print("  [bold white]Servers[/bold white]")
    console.print()

    try:
        count_str = pt_prompt(
            HTML("<prompt>  how many?  </prompt>"),
            style=LOGIN_STYLE,
        ).strip()
        server_count = int(count_str) if count_str.isdigit() and int(count_str) > 0 else 0
    except (KeyboardInterrupt, EOFError):
        server_count = 0

    servers = []
    for i in range(server_count):
        console.print()
        console.print(f"  [bold white]Server {i + 1}[/bold white]")
        try:
            ip = pt_prompt(HTML("<prompt>  ip        </prompt>"), style=LOGIN_STYLE).strip()
            if not ip:
                break
            username = pt_prompt(HTML("<prompt>  username  </prompt>"), style=LOGIN_STYLE).strip()
            if not username:
                break
            password = pt_prompt(HTML("<prompt>  password  </prompt>"), style=LOGIN_STYLE, is_password=True).strip()
            if not password:
                break

            servers.append({"ip": ip, "username": username, "password": password})
            console.print(f"  [bold green]+[/bold green] [bold white]{username}@{ip}[/bold white]")
        except (KeyboardInterrupt, EOFError):
            break

    # ── Submit ──
    console.print()
    with console.status("  Saving...", spinner="circle", spinner_style=SALMON):
        try:
            r = client.post("/setup", json={
                "team": team_key,
                "team_members": members,
                "servers": servers,
            })
            r.raise_for_status()
            data = r.json()
        except Exception as e:
            console.print(f"  [bold red]-[/bold red] Setup failed: {e}")
            return

    console.print(f"  [bold green]+[/bold green] [bold white]Setup complete[/bold white]")
    console.print(f"  [{SALMON}]{data.get('users', 0)} users · {data.get('servers', 0)} servers[/{SALMON}]")
    console.print()


# ──────────────────────────────────────────────
#  /help
# ──────────────────────────────────────────────

def handle_help():
    console.print()
    console.print(f"  [bold white]Commands[/bold white]")
    console.print()
    cmds = [
        ("/help", "Show this help screen"),
        ("/generate", "Generate a dataset (e.g. /generate turkish surnames)"),
        ("/train", "Train a model on project datasets"),
        ("/models", "List trained models in current project"),
        ("/deploy", "Deploy a model and get an API key"),
        ("/info", "Project status, datasets, and models"),
        ("/servers", "List, add, edit, or remove servers (max 2)"),
        ("/projects", "Create, delete, and manage projects"),
        ("/logout", "Sign out and return to login"),
        ("/exit", "Quit Kedy"),
    ]
    for cmd, desc in cmds:
        console.print(f"  [{SALMON}]{cmd:<12}[/{SALMON}] [white]{desc}[/white]")
    console.print()
    thin_rule()
    console.print()
    console.print(f"  [bold white]Usage[/bold white]")
    console.print()
    console.print(f"  [white]Type a natural language description to generate a dataset.[/white]")
    console.print(f"  [white]Kedy will search for references, build a plan, ask clarifying[/white]")
    console.print(f"  [white]questions, then generate and save the data as JSON.[/white]")
    console.print()
    console.print(f"  [#999999]Examples:[/#999999]")
    console.print(f"  [{SALMON}]>[/{SALMON}]  [white]50 rows of Turkish e-commerce transactions[/white]")
    console.print(f"  [{SALMON}]>[/{SALMON}]  [white]cybersecurity incident logs with severity levels[/white]")
    console.print(f"  [{SALMON}]>[/{SALMON}]  [white]financial portfolio data with risk scores[/white]")
    console.print()


# ──────────────────────────────────────────────
#  /servers
# ──────────────────────────────────────────────

def _numbered_menu(options: list[str]) -> int:
    """Show numbered options, return selected index or -1."""
    completer = OptionCompleter(options)
    try:
        def _open():
            get_app().current_buffer.start_completion()

        session = PromptSession(
            HTML("<prompt>  ~  </prompt>"),
            style=PT_STYLE,
            completer=completer,
            complete_while_typing=True,
        )
        answer = session.prompt(pre_run=_open).strip()
    except (KeyboardInterrupt, EOFError):
        return -1

    if answer.isdigit() and 1 <= int(answer) <= len(options):
        return int(answer) - 1
    for i, opt in enumerate(options):
        if opt.lower() == answer.lower():
            return i
    return -1


def handle_servers(client: httpx.Client, team_key: str):
    def fetch():
        try:
            r = client.post("/servers", json={"team": team_key, "action": "list"})
            r.raise_for_status()
            d = r.json()
            return d.get("servers", []) if d.get("ok") else []
        except Exception:
            return []

    servers = fetch()
    console.print()
    console.print(f"  [bold white]Servers[/bold white]  [{SALMON}]{len(servers)}/2[/{SALMON}]")
    if servers:
        for i, s in enumerate(servers):
            console.print(f"    [{SALMON}]{i+1}.[/{SALMON}] [bold white]{s.get('name', s['ip'])}[/bold white]  [#777777]{s.get('username','?')}@{s['ip']}[/#777777]")
    else:
        console.print(f"    [#777777]No servers yet[/#777777]")
    console.print()

    options = []
    if len(servers) < 2:
        options.append("Add")
    if servers:
        options.extend(["Edit", "Remove"])
    options.append("Back")
    inline = "    ".join(f"[{SALMON}]{i+1}[/{SALMON}] [white]{o}[/white]" for i, o in enumerate(options))
    console.print(f"  {inline}")
    console.print()

    choice = _numbered_menu(options)
    if choice < 0 or options[choice] == "Back":
        return
    selected = options[choice]

    if selected == "Add":
        try:
            name = pt_prompt(HTML("<prompt>  name  </prompt>"), style=LOGIN_STYLE).strip()
            if not name: return
            ip = pt_prompt(HTML("<prompt>  ip  </prompt>"), style=LOGIN_STYLE).strip()
            if not ip: return
            user = pt_prompt(HTML("<prompt>  user  </prompt>"), style=LOGIN_STYLE).strip()
            if not user: return
            pw = pt_prompt(HTML("<prompt>  pass  </prompt>"), style=LOGIN_STYLE, is_password=True).strip()
            if not pw: return
        except (KeyboardInterrupt, EOFError):
            return
        try:
            r = client.post("/servers", json={"team": team_key, "action": "add", "server": {"name": name, "ip": ip, "username": user, "password": pw}})
            r.raise_for_status()
            d = r.json()
            console.print(f"  [bold green]+[/bold green] {name}" if d.get("ok") else f"  [bold red]-[/bold red] {d.get('error')}")
        except Exception as e:
            console.print(f"  [bold red]-[/bold red] {e}")

    elif selected == "Edit":
        names = [s.get("name", s["ip"]) for s in servers]
        idx = _numbered_menu(names)
        if idx < 0: return
        s = servers[idx]
        console.print(f"  [#777777]enter to keep current[/#777777]")
        try:
            name = pt_prompt(HTML("<prompt>  name  </prompt>"), style=LOGIN_STYLE, default=s.get("name", "")).strip()
            ip = pt_prompt(HTML("<prompt>  ip  </prompt>"), style=LOGIN_STYLE, default=s.get("ip", "")).strip()
            user = pt_prompt(HTML("<prompt>  user  </prompt>"), style=LOGIN_STYLE, default=s.get("username", "")).strip()
            pw = pt_prompt(HTML("<prompt>  pass  </prompt>"), style=LOGIN_STYLE, is_password=True).strip()
        except (KeyboardInterrupt, EOFError):
            return
        updates = {k: v for k, v in [("name", name), ("ip", ip), ("username", user), ("password", pw)] if v}
        if not updates:
            return
        try:
            r = client.post("/servers", json={"team": team_key, "action": "edit", "index": idx, "server": updates})
            r.raise_for_status()
            d = r.json()
            console.print(f"  [bold green]+[/bold green] Updated" if d.get("ok") else f"  [bold red]-[/bold red] {d.get('error')}")
        except Exception as e:
            console.print(f"  [bold red]-[/bold red] {e}")

    elif selected == "Remove":
        names = [s.get("name", s["ip"]) for s in servers]
        idx = _numbered_menu(names)
        if idx < 0: return
        try:
            r = client.post("/servers", json={"team": team_key, "action": "remove", "index": idx})
            r.raise_for_status()
            d = r.json()
            console.print(f"  [{SALMON}]-[/{SALMON}] {names[idx]} removed" if d.get("ok") else f"  [bold red]-[/bold red] {d.get('error')}")
        except Exception as e:
            console.print(f"  [bold red]-[/bold red] {e}")


# ──────────────────────────────────────────────
#  /projects
# ──────────────────────────────────────────────

def handle_projects(client: httpx.Client, team_key: str, session_data: dict):
    global active_project, active_project_index

    def fetch():
        try:
            r = client.post("/projects", json={"team": team_key, "action": "list"})
            r.raise_for_status()
            d = r.json()
            return d.get("projects", []) if d.get("ok") else []
        except Exception:
            return []

    projects = fetch()
    console.print()
    console.print(f"  [bold white]Projects[/bold white]  [{SALMON}]{len(projects)}[/{SALMON}]")
    if projects:
        for i, p in enumerate(projects):
            dot = f"[bold green]●[/bold green]" if active_project and active_project.get("name") == p["name"] else " "
            console.print(f"   {dot} [{SALMON}]{i+1}.[/{SALMON}] [bold white]{p['name']}[/bold white]  [#777777]{p.get('server_name', '?')}[/#777777]")
    else:
        console.print(f"    [#777777]No projects yet[/#777777]")
    console.print()

    options = ["New"]
    if projects:
        options.extend(["Enter", "Datasets", "Delete"])
    if active_project:
        options.append("Leave")
    options.append("Back")
    inline = "    ".join(f"[{SALMON}]{i+1}[/{SALMON}] [white]{o}[/white]" for i, o in enumerate(options))
    console.print(f"  {inline}")
    console.print()

    choice = _numbered_menu(options)
    if choice < 0 or options[choice] == "Back":
        return
    selected = options[choice]

    if selected == "New":
        servers = session_data.get("team", {}).get("servers", [])
        if not servers:
            console.print(f"  [bold red]-[/bold red] No servers. Use /servers first")
            return
        try:
            name = pt_prompt(HTML("<prompt>  name  </prompt>"), style=LOGIN_STYLE).strip()
            if not name: return
        except (KeyboardInterrupt, EOFError):
            return
        srv_names = [s.get("name", s["ip"]) for s in servers]
        idx = _numbered_menu(srv_names)
        if idx < 0: return
        with console.status("  Creating...", spinner="circle", spinner_style=SALMON):
            try:
                r = client.post("/projects", json={"team": team_key, "action": "create", "project": {"name": name, "server_index": idx}})
                r.raise_for_status()
                d = r.json()
            except Exception as e:
                console.print(f"  [bold red]-[/bold red] {e}"); return
        console.print(f"  [bold green]+[/bold green] {name}" if d.get("ok") else f"  [bold red]-[/bold red] {d.get('error')}")

    elif selected == "Enter":
        pnames = [p["name"] for p in projects]
        idx = _numbered_menu(pnames)
        if idx < 0: return
        active_project = projects[idx]
        active_project_index = idx
        console.print(f"  [bold green]+[/bold green] [bold white]{active_project['name']}[/bold white]  [#777777]{active_project.get('server_name','')}:{active_project.get('data_path','')}[/#777777]")

    elif selected == "Datasets":
        pnames = [p["name"] for p in projects]
        idx = _numbered_menu(pnames)
        if idx < 0: return
        with console.status("  Loading...", spinner="circle", spinner_style=SALMON):
            try:
                r = client.post("/projects", json={"team": team_key, "action": "datasets", "index": idx})
                r.raise_for_status()
                d = r.json()
            except Exception as e:
                console.print(f"  [bold red]-[/bold red] {e}"); return
        if not d.get("ok"):
            console.print(f"  [bold red]-[/bold red] {d.get('error')}"); return
        datasets = d.get("datasets", [])
        console.print(f"  [bold white]{projects[idx]['name']}[/bold white]  [{SALMON}]{len(datasets)} file(s)[/{SALMON}]")
        if not datasets:
            console.print(f"  [#777777]Empty. Enter project and generate data.[/#777777]"); return
        for i, ds in enumerate(datasets):
            console.print(f"    [{SALMON}]{i+1}.[/{SALMON}] [white]{ds}[/white]")
        console.print()
        vi = _numbered_menu(datasets)
        if vi < 0: return
        with console.status("  Loading...", spinner="circle", spinner_style=SALMON):
            try:
                r = client.post("/projects", json={"team": team_key, "action": "view", "index": idx, "filename": datasets[vi]})
                r.raise_for_status()
                vd = r.json()
            except Exception as e:
                console.print(f"  [bold red]-[/bold red] {e}"); return
        if not vd.get("ok"):
            console.print(f"  [bold red]-[/bold red] {vd.get('error')}"); return
        rows = vd.get("data", [])
        console.print(f"  [bold white]{datasets[vi]}[/bold white]  [{SALMON}]{len(rows)} rows[/{SALMON}]")
        for row in (rows[:5] if isinstance(rows, list) else []):
            if isinstance(row, dict):
                parts = [f"[white]{k}[/white]=[#777777]{v}[/#777777]" for k, v in list(row.items())[:4]]
                console.print(f"    {', '.join(parts)}")
        if len(rows) > 5:
            console.print(f"    [#777777]+{len(rows)-5} more[/#777777]")

    elif selected == "Delete":
        pnames = [p["name"] for p in projects]
        idx = _numbered_menu(pnames)
        if idx < 0: return
        name = projects[idx]["name"]
        try:
            r = client.post("/projects", json={"team": team_key, "action": "delete", "index": idx})
            r.raise_for_status()
            d = r.json()
            if d.get("ok"):
                console.print(f"  [{SALMON}]-[/{SALMON}] {name} deleted")
                if active_project and active_project.get("name") == name:
                    active_project = None
                    active_project_index = -1
            else:
                console.print(f"  [bold red]-[/bold red] {d.get('error')}")
        except Exception as e:
            console.print(f"  [bold red]-[/bold red] {e}")

    elif selected == "Leave":
        n = active_project["name"]
        active_project = None
        active_project_index = -1
        console.print(f"  [{SALMON}]-[/{SALMON}] Left {n}")


# ──────────────────────────────────────────────
#  Plan display and questions
# ──────────────────────────────────────────────

def _clean_plan_summary(text: str) -> str:
    text = re.sub(r"```[\s\S]*?```", "", text)
    text = re.sub(r"\{[\s\S]*\}", "", text)
    text = re.sub(r"\[[\s\S]*\]", "", text)
    lines = text.split("\n")
    cleaned = []
    for line in lines:
        stripped = line.strip()
        if stripped.startswith(("* `", "- `", "\u2022 `", "* **", "- **")):
            continue
        if stripped.lower().startswith("note:"):
            continue
        if stripped.lower().startswith(("the dataset includes", "the dataset contains")):
            continue
        if re.search(r"(?i)(here is|here's|generated).*(json|csv|dataset|data)", stripped):
            continue
        if stripped.startswith(('"', "'", "{", "}", "[", "]")):
            continue
        cleaned.append(line)
    text = "\n".join(cleaned).strip()
    text = re.sub(r"\n{3,}", "\n\n", text)
    if len(text) > 200:
        cut = text[:200].rfind(". ")
        if cut > 50:
            text = text[:cut + 1]
        else:
            text = text[:200] + "\u2026"
    return text


def show_plan(plan_data: dict):
    summary = plan_data.get("summary", "")
    steps = plan_data.get("steps", [])

    console.print()
    if summary:
        clean = _clean_plan_summary(summary)
        if clean:
            console.print(f"  [bold]Plan[/bold]  [dim]{clean}[/dim]")
        else:
            console.print(f"  [bold]Plan[/bold]")
    for i, step in enumerate(steps, 1):
        console.print(f"  [salmon]{i}.[/salmon] {step}")
    console.print()


def ask_plan_questions(questions: list) -> dict:
    answers = {}

    for q_data in questions:
        question = q_data.get("question", "")
        options = q_data.get("options", [])
        if not question or not options:
            continue

        console.print(f"  [bold]{question}[/bold]")
        completer = OptionCompleter(options)

        try:
            def _open():
                get_app().current_buffer.start_completion()

            session = PromptSession(
                HTML("<prompt>  &gt; </prompt>"),
                style=PT_STYLE,
                completer=completer,
                complete_while_typing=True,
            )
            answer = session.prompt(pre_run=_open).strip()
        except (KeyboardInterrupt, EOFError):
            console.print("  [dim]skipped[/dim]")
            continue

        if answer.isdigit() and 1 <= int(answer) <= len(options):
            answer = options[int(answer) - 1]

        answers[question] = answer
        console.print(f"  [dim]-> {answer}[/dim]")
        console.print()

    return answers


def confirm_action(description: str) -> bool:
    console.print(f"  [bold]{description}[/bold]")
    completer = OptionCompleter(["Yes", "No"])

    try:
        def _open():
            get_app().current_buffer.start_completion()

        session = PromptSession(
            HTML("<prompt>  allow? &gt; </prompt>"),
            style=PT_STYLE,
            completer=completer,
            complete_while_typing=True,
        )
        answer = session.prompt(pre_run=_open).strip().lower()
    except (KeyboardInterrupt, EOFError):
        return False

    return answer in ("yes", "y", "1")


# ──────────────────────────────────────────────
#  Chat / data generation
# ──────────────────────────────────────────────

SPIN_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]


class Spinner:
    """Animated inline spinner for agent steps."""

    def __init__(self, msg: str):
        self.msg = msg
        self._stop = threading.Event()
        self._thread = None

    def start(self):
        self._stop.clear()
        self._thread = threading.Thread(target=self._animate, daemon=True)
        self._thread.start()

    def _animate(self):
        i = 0
        while not self._stop.is_set():
            frame = SPIN_FRAMES[i % len(SPIN_FRAMES)]
            sys.stdout.write(f"\r  \033[38;2;250;128;114m{frame}\033[0m \033[38;5;245m{self.msg}\033[0m\033[K")
            sys.stdout.flush()
            i += 1
            self._stop.wait(0.08)

    def stop(self, result: str, ok: bool = True):
        self._stop.set()
        if self._thread:
            self._thread.join()
        dot = f"\033[38;2;250;128;114m●\033[0m" if ok else "\033[38;5;240m○\033[0m"
        color = "\033[97m" if ok else "\033[38;5;245m"
        sys.stdout.write(f"\r  {dot} {color}{result}\033[0m\033[K\n")
        sys.stdout.flush()


# Background agent state
agent_status: dict = {
    "running": False,
    "project": "",
    "step": "",
    "step_num": 0,
    "total_steps": 4,
    "start_time": 0.0,
    "step_start_time": 0.0,
    "result": None,  # {"ok": bool, "message": str} when done
}
agent_lock = threading.Lock()

# Average durations per step (seconds) — used for ETA estimation
_STEP_AVG_DURATIONS = {
    1: 5,    # Search
    2: 45,   # Generate
    3: 10,   # Validate
    4: 5,    # Save
}


# Background training state
train_status: dict = {
    "running": False,
    "project": "",
    "epoch": 0,
    "total_epochs": 0,
    "loss": 0.0,
    "start_time": 0.0,
    "result": None,  # {"ok": bool, "message": str} when done
}
train_lock = threading.Lock()


def _agent_worker(client: httpx.Client, user_input: str, enriched_message: str,
                  plan_data: dict | None, team_key: str,
                  project: dict, project_index: int):
    """Background worker that runs search → generate → validate → save."""
    global agent_status

    def set_step(step: str, step_num: int = 0):
        with agent_lock:
            agent_status["step"] = step
            if step_num:
                agent_status["step_num"] = step_num
                agent_status["step_start_time"] = time.time()

    def finish(ok: bool, message: str):
        with agent_lock:
            agent_status["running"] = False
            agent_status["step"] = ""
            agent_status["result"] = {"ok": ok, "message": message}
        # Terminal bell
        sys.stdout.write("\a")
        sys.stdout.flush()
        # macOS notification
        try:
            title = "Kedy — Data Agent"
            body = message if ok else f"Failed: {message}"
            subprocess.Popen(
                ["osascript", "-e", f'display notification "{body}" with title "{title}"'],
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            )
        except Exception:
            pass

    try:
        # Step 1: Search
        set_step("Searching for references", 1)
        try:
            sr = client.post("/search", json={"topic": user_input})
            sr.raise_for_status()
        except Exception:
            pass

        # Step 2: Generate
        set_step("Generating dataset", 2)
        chat_payload = {"message": enriched_message}
        if plan_data:
            chat_payload["plan"] = plan_data

        try:
            r = client.post("/chat", json=chat_payload)
            r.raise_for_status()
            data = r.json()
        except Exception as e:
            finish(False, f"Generation failed: {e}")
            return

        response_text = data.get("response", "")
        extracted_rows = data.get("extracted_rows", [])
        validation = data.get("validation")

        # Retry on error
        if not extracted_rows and re.search(r"(?i)(error|bad request|internal server error)", response_text):
            set_step("Retrying generation", 2)
            try:
                r = client.post("/chat", json=chat_payload)
                r.raise_for_status()
                data = r.json()
                extracted_rows = data.get("extracted_rows", [])
                validation = data.get("validation")
            except Exception as e:
                finish(False, f"Generation failed: {e}")
                return

        if not extracted_rows:
            finish(False, "No structured data could be extracted")
            return

        # Build filename
        ctx = data.get("context", {})
        title = ctx.get("Title") or ctx.get("title") or ""
        if not title or title.lower() in ("unknown", "dataset", "untitled"):
            slug = re.sub(r"[^\w\s-]", "", user_input.lower())
            slug = re.sub(r"\s+", "_", slug).strip("_")
            title = slug[:40] or "dataset"
        filename = f"{title.replace(' ', '_')[:40]}.json"

        # Step 3: Validate
        set_step("Validating data", 3)
        issues = validation.get("issues", []) if validation else []

        # Step 4: Save
        set_step("Saving to server", 4)
        if project and project_index >= 0:
            try:
                r = client.post("/projects", json={
                    "team": team_key,
                    "action": "save_remote",
                    "index": project_index,
                    "filename": filename,
                    "data": extracted_rows,
                })
                r.raise_for_status()
                save_data = r.json()
            except Exception as e:
                finish(False, f"Save failed: {e}")
                return

            if save_data.get("ok"):
                path = save_data.get("path", "")
                tokens = data.get("tokens", 0)
                parts = [f"{len(extracted_rows)} rows"]
                if issues:
                    parts.append(f"{len(issues)} issues")
                if tokens:
                    parts.append(f"{tokens:,} tokens")
                finish(True, f"{filename} → {project.get('server_name', '')}:{path}  ({', '.join(parts)})")
            else:
                finish(False, f"Save failed: {save_data.get('error', 'Unknown')}")
        else:
            save_path = os.path.join(os.getcwd(), filename)
            try:
                with open(save_path, "w") as f:
                    json.dump(extracted_rows, f, indent=2)
                tokens = data.get("tokens", 0)
                parts = [f"{len(extracted_rows)} rows"]
                if tokens:
                    parts.append(f"{tokens:,} tokens")
                finish(True, f"Saved → {save_path}  ({', '.join(parts)})")
            except Exception as e:
                finish(False, f"Save failed: {e}")

    except Exception as e:
        finish(False, f"Unexpected error: {e}")


def handle_chat(client: httpx.Client, user_input: str, auto_accept: bool, team_key: str = "core"):
    global agent_status

    # Check if agent is already running
    with agent_lock:
        if agent_status["running"]:
            console.print(f"  [{SALMON}]![/{SALMON}] [white]Agent is already working. Wait for it to finish.[/white]")
            return

    proj_name = active_project["name"] if active_project else ""

    # Step 1: Plan (synchronous — needs user input)
    sp = Spinner("Building plan")
    sp.start()
    plan_data = None
    answers = {}
    context = {}
    try:
        r = client.post("/plan", json={"message": user_input})
        r.raise_for_status()
        data = r.json()
        plan_data = data.get("plan")
        context = data.get("context", {})
        category = context.get("Category", "")
        cat_label = f" · {category}" if category and category != "Unknown" else ""
        sp.stop(f"Plan ready{cat_label}", ok=True)
    except Exception as e:
        sp.stop("Planning failed", ok=False)
        console.print(f"    [bold red]{e}[/bold red]")
        return

    # Show plan & collect answers
    if plan_data:
        show_plan(plan_data)

        questions = plan_data.get("questions", [])
        if questions:
            answers = ask_plan_questions(questions)

        if not auto_accept:
            if not confirm_action("Proceed?"):
                console.print("  [dim]cancelled[/dim]")
                return

        # Clear plan/questions from screen
        steps = plan_data.get("steps", [])
        summary = plan_data.get("summary", "")
        lines_to_clear = 2
        if summary:
            lines_to_clear += 1
        lines_to_clear += len(steps) + 1
        for q_data in questions:
            if q_data.get("question") and q_data.get("options"):
                lines_to_clear += 3
        if not auto_accept:
            lines_to_clear += 2
        # Also clear the "Plan ready" spinner line
        lines_to_clear += 1

        for _ in range(lines_to_clear):
            sys.stdout.write("\033[A\033[2K")
        sys.stdout.flush()

    enriched_message = user_input
    if answers:
        prefs = "\n".join(f"- {q}: {a}" for q, a in answers.items())
        enriched_message = f"{user_input}\n\nUser preferences:\n{prefs}"

    # Launch background worker
    with agent_lock:
        agent_status["running"] = True
        agent_status["project"] = proj_name
        agent_status["step"] = "Starting"
        agent_status["step_num"] = 0
        agent_status["start_time"] = time.time()
        agent_status["step_start_time"] = time.time()
        agent_status["result"] = None

    worker = threading.Thread(
        target=_agent_worker,
        args=(client, user_input, enriched_message, plan_data, team_key,
              active_project, active_project_index),
        daemon=True,
    )
    worker.start()

    console.print(f"  [{SALMON}]●[/{SALMON}] [white]Data Agent is working on[/white] [bold white]{proj_name}[/bold white]")
    console.print(f"    [#777777]You can keep using commands. Status is shown in the toolbar.[/#777777]")
    console.print()


# ──────────────────────────────────────────────
#  /train
# ──────────────────────────────────────────────

def _train_poller(client: httpx.Client, team_key: str, project_index: int, proj_name: str, total_epochs: int):
    """Background thread that polls training status from the server."""
    global train_status
    start = time.time()

    while True:
        time.sleep(3)
        with train_lock:
            if not train_status["running"]:
                return

        try:
            r = client.post("/train", json={
                "team": team_key,
                "action": "status",
                "project_index": project_index,
            })
            r.raise_for_status()
            data = r.json()
        except Exception:
            continue

        epoch = data.get("epoch", 0)
        status = data.get("status", "running")
        loss = data.get("loss", 0.0)

        with train_lock:
            train_status["epoch"] = epoch
            train_status["loss"] = loss

        if status == "done":
            with train_lock:
                train_status["running"] = False
                train_status["result"] = {"ok": True, "message": f"Training complete — {total_epochs} epochs, loss {loss:.4f}"}
            sys.stdout.write("\a")
            sys.stdout.flush()
            try:
                subprocess.Popen(
                    ["osascript", "-e", f'display notification "Training complete on {proj_name}" with title "Kedy — Training"'],
                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                )
            except Exception:
                pass
            return

        if status == "error":
            err = data.get("error", "Unknown error")
            with train_lock:
                train_status["running"] = False
                train_status["result"] = {"ok": False, "message": f"Training failed: {err}"}
            sys.stdout.write("\a")
            sys.stdout.flush()
            try:
                subprocess.Popen(
                    ["osascript", "-e", f'display notification "Training failed on {proj_name}" with title "Kedy — Training"'],
                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                )
            except Exception:
                pass
            return

        if status == "stopped":
            with train_lock:
                train_status["running"] = False
                train_status["result"] = {"ok": False, "message": "Training stopped"}
            return


def _stop_training(client: httpx.Client, team_key: str):
    """Stop a running training session."""
    global train_status

    proj_name = train_status.get("project", "")
    sp = Spinner("Stopping training")
    sp.start()
    try:
        r = client.post("/train", json={
            "team": team_key,
            "action": "stop",
            "project_index": active_project_index,
        })
        r.raise_for_status()
        d = r.json()
    except Exception as e:
        sp.stop(f"Error: {e}", ok=False)
        return

    with train_lock:
        train_status["running"] = False
        train_status["result"] = None
    sp.stop(d.get("message", "Stopped"), ok=True)


def handle_train(client: httpx.Client, team_key: str):
    global train_status

    if not active_project:
        console.print(f"  [{SALMON}]![/{SALMON}] [white]Select a project first with[/white] [{SALMON}]/projects[/{SALMON}]")
        return

    # If training is already running, offer stop
    with train_lock:
        is_running = train_status["running"]
        running_proj = train_status.get("project", "")
        running_epoch = train_status.get("epoch", 0)
        running_total = train_status.get("total_epochs", 0)
        running_loss = train_status.get("loss", 0.0)

    if is_running:
        console.print()
        console.print(f"  [bold white]Training in progress[/bold white]  [{SALMON}]{running_proj}[/{SALMON}]")
        console.print(f"    [white]Epoch {running_epoch}/{running_total}  loss {running_loss:.4f}[/white]")
        console.print()
        options = ["Stop", "Back"]
        inline = "    ".join(f"[{SALMON}]{i+1}[/{SALMON}] [white]{o}[/white]" for i, o in enumerate(options))
        console.print(f"  {inline}")
        console.print()
        choice = _numbered_menu(options)
        if choice == 0:
            _stop_training(client, team_key)
        return

    proj_name = active_project["name"]

    # Fetch datasets
    sp = Spinner("Loading datasets")
    sp.start()
    try:
        r = client.post("/projects", json={"team": team_key, "action": "datasets", "index": active_project_index})
        r.raise_for_status()
        d = r.json()
    except Exception as e:
        sp.stop("Failed to load datasets", ok=False)
        console.print(f"    [bold red]{e}[/bold red]")
        return

    if not d.get("ok"):
        sp.stop(d.get("error", "Error"), ok=False)
        return

    datasets = [f for f in d.get("datasets", []) if f.endswith(".json")]
    if not datasets:
        sp.stop("No datasets found", ok=False)
        console.print(f"    [#777777]Generate data first with /generate[/#777777]")
        return

    sp.stop(f"{len(datasets)} dataset(s) found", ok=True)
    console.print()

    # Show datasets for selection
    for i, ds in enumerate(datasets):
        console.print(f"    [{SALMON}]{i+1}.[/{SALMON}] [white]{ds}[/white]")
    console.print()
    console.print(f"    [#777777]Enter numbers separated by comma, or 'all'[/#777777]")

    try:
        answer = pt_prompt(HTML("<prompt>  datasets ~  </prompt>"), style=PT_STYLE).strip()
    except (KeyboardInterrupt, EOFError):
        return

    if not answer:
        return

    if answer.lower() == "all":
        selected = datasets[:]
    else:
        selected = []
        for part in answer.split(","):
            part = part.strip()
            if part.isdigit() and 1 <= int(part) <= len(datasets):
                selected.append(datasets[int(part) - 1])
        if not selected:
            console.print(f"  [{SALMON}]![/{SALMON}] [white]No valid selection[/white]")
            return

    console.print()
    console.print(f"  [bold white]Selected:[/bold white] [white]{', '.join(selected)}[/white]")

    # Ask for epochs
    try:
        epochs_str = pt_prompt(HTML("<prompt>  epochs (10) ~  </prompt>"), style=PT_STYLE).strip()
    except (KeyboardInterrupt, EOFError):
        return
    epochs = int(epochs_str) if epochs_str.isdigit() and int(epochs_str) > 0 else 10

    # Start training
    sp = Spinner("Starting training")
    sp.start()
    try:
        r = client.post("/train", json={
            "team": team_key,
            "action": "start",
            "project_index": active_project_index,
            "datasets": selected,
            "epochs": epochs,
        }, timeout=60)
        r.raise_for_status()
        d = r.json()
    except Exception as e:
        sp.stop("Failed to start training", ok=False)
        console.print(f"    [bold red]{e}[/bold red]")
        return

    if not d.get("ok"):
        sp.stop(d.get("error", "Error"), ok=False)
        return

    data_type = d.get("data_type", "text")
    model_name = d.get("model_name", "")
    sp.stop(d.get("message", "Training started"), ok=True)

    # Initialize train status and start poller
    with train_lock:
        train_status["running"] = True
        train_status["project"] = proj_name
        train_status["epoch"] = 0
        train_status["total_epochs"] = epochs
        train_status["loss"] = 0.0
        train_status["start_time"] = time.time()
        train_status["result"] = None

    poller = threading.Thread(
        target=_train_poller,
        args=(client, team_key, active_project_index, proj_name, epochs),
        daemon=True,
    )
    poller.start()

    console.print(f"    [#777777]Model: {model_name} ({data_type}). Progress shown in toolbar.[/#777777]")
    console.print()


# ──────────────────────────────────────────────
#  /models
# ──────────────────────────────────────────────

def handle_models(client: httpx.Client, team_key: str):
    if not active_project:
        console.print(f"  [{SALMON}]![/{SALMON}] [white]Select a project first with[/white] [{SALMON}]/projects[/{SALMON}]")
        return

    sp = Spinner("Loading models")
    sp.start()
    try:
        r = client.post("/models", json={"team": team_key, "project_index": active_project_index})
        r.raise_for_status()
        d = r.json()
    except Exception as e:
        sp.stop(f"Error: {e}", ok=False)
        return

    if not d.get("ok"):
        sp.stop(d.get("error", "Error"), ok=False)
        return

    models = d.get("models", [])
    train_log = d.get("train_status")
    sp.stop(f"{len(models)} model(s)", ok=True)
    console.print()

    if not models:
        console.print(f"    [#777777]No models yet. Use /train to train one.[/#777777]")
    else:
        for i, m in enumerate(models):
            mtype = m.get("model_type", "?")
            mname = m.get("model_name", "?")
            badge = f"[#00CC00]●[/#00CC00]" if mtype == "text" else f"[#4488FF]●[/#4488FF]"
            console.print(f"    {badge} [{SALMON}]{i+1}.[/{SALMON}] [bold white]{mname}[/bold white]  [#777777]{mtype}[/#777777]")
            if m.get("features"):
                console.print(f"       [#777777]features: {', '.join(m['features'])}[/#777777]")
            if m.get("target"):
                console.print(f"       [#777777]target: {m['target']}[/#777777]")
            if m.get("rows"):
                console.print(f"       [#777777]{m['rows']} rows, {m.get('epochs', '?')} epochs[/#777777]")

    # Show training status if active
    if train_log and train_log.get("status") == "running":
        ep = train_log.get("epoch", 0)
        total = train_log.get("total_epochs", 0)
        loss = train_log.get("loss", 0.0)
        console.print()
        console.print(f"    [{SALMON}]⠿[/{SALMON}] [white]Training: epoch {ep}/{total}  loss {loss:.4f}[/white]")

    console.print()


# ──────────────────────────────────────────────
#  /info
# ──────────────────────────────────────────────

def handle_info(client: httpx.Client, team_key: str):
    if not active_project:
        console.print(f"  [{SALMON}]![/{SALMON}] [white]Select a project first with[/white] [{SALMON}]/projects[/{SALMON}]")
        return

    sp = Spinner("Loading project info")
    sp.start()
    try:
        r = client.post("/info", json={"team": team_key, "project_index": active_project_index})
        r.raise_for_status()
        d = r.json()
    except Exception as e:
        sp.stop(f"Error: {e}", ok=False)
        return

    if not d.get("ok"):
        sp.stop(d.get("error", "Error"), ok=False)
        return

    sp.stop(active_project["name"], ok=True)

    project = d.get("project", {})
    datasets = d.get("datasets", [])
    models = d.get("models", [])
    training = d.get("training")

    console.print()
    console.print(f"  [bold white]{project.get('name', '')}[/bold white]  [#777777]{project.get('server_name', '')}:{project.get('data_path', '')}[/#777777]")
    console.print()

    # Active tasks
    console.print(f"  [bold white]Active[/bold white]")
    has_active = False
    with agent_lock:
        if agent_status["running"]:
            step = agent_status.get("step", "Working")
            step_num = agent_status.get("step_num", 0)
            total_s = agent_status.get("total_steps", 4)
            console.print(f"    [{SALMON}]●[/{SALMON}] [white]Data Agent: {step} [{step_num}/{total_s}][/white]")
            has_active = True
    if training:
        ep = training.get("epoch", 0)
        total_ep = training.get("total_epochs", 0)
        loss = training.get("loss", 0.0)
        mname = training.get("model_name", "")
        console.print(f"    [{SALMON}]●[/{SALMON}] [white]Training: {mname} epoch {ep}/{total_ep}  loss {loss:.4f}[/white]")
        has_active = True
    if not has_active:
        console.print(f"    [#777777]No active tasks[/#777777]")
    console.print()

    # Datasets
    console.print(f"  [bold white]Datasets[/bold white]  [{SALMON}]{len(datasets)}[/{SALMON}]")
    if datasets:
        for i, ds in enumerate(datasets):
            console.print(f"    [{SALMON}]{i+1}.[/{SALMON}] [white]{ds}[/white]")
    else:
        console.print(f"    [#777777]No datasets. Use /generate to create one.[/#777777]")
    console.print()

    # Models
    console.print(f"  [bold white]Models[/bold white]  [{SALMON}]{len(models)}[/{SALMON}]")
    if models:
        for i, m in enumerate(models):
            mtype = m.get("model_type", "?")
            mname = m.get("model_name", "?")
            badge = f"[#00CC00]●[/#00CC00]" if mtype == "text" else f"[#4488FF]●[/#4488FF]"
            extra = ""
            if m.get("rows"):
                extra = f"  [#777777]{m['rows']} rows, {m.get('epochs', '?')} epochs[/#777777]"
            console.print(f"    {badge} [{SALMON}]{i+1}.[/{SALMON}] [bold white]{mname}[/bold white]  [#777777]{mtype}[/#777777]{extra}")
    else:
        console.print(f"    [#777777]No models. Use /train to train one.[/#777777]")
    console.print()


# ──────────────────────────────────────────────
#  /deploy
# ──────────────────────────────────────────────

def _prediction_example_payload(model_type: str, schema: dict | None) -> str:
    schema = schema or {}
    example_request = schema.get("example_request")
    if isinstance(example_request, dict) and example_request:
        return json.dumps(example_request, ensure_ascii=False)
    if model_type == "numeric":
        return ""
    return json.dumps({"input": "your prompt text here"}, ensure_ascii=False)


def _indented_json_lines(obj: dict, indent: int = 0) -> list[str]:
    text = json.dumps(obj, ensure_ascii=False, indent=2)
    prefix = " " * indent
    return [prefix + line for line in text.splitlines()]


def _is_id_like_key(name: str) -> bool:
    key = (name or "").strip().lower()
    return key in {"id", "_id", "uuid"} or key.endswith("_id") or key.endswith("id")


def _display_example_request(schema: dict | None) -> dict:
    schema = schema or {}
    example_request = schema.get("example_request")
    if not isinstance(example_request, dict):
        return {}

    req = json.loads(json.dumps(example_request))
    data = req.get("data")
    if isinstance(data, dict):
        filtered = {k: v for k, v in data.items() if not _is_id_like_key(k)}
        if filtered:
            req["data"] = filtered
    return req


def _split_primary_and_extra_params(example_request: dict) -> tuple[dict, list[tuple[str, object]]]:
    req = json.loads(json.dumps(example_request or {}))
    data = req.get("data")
    if not isinstance(data, dict) or not data:
        return req, []

    items = list(data.items())
    primary_key, primary_value = items[0]
    req["data"] = {primary_key: primary_value}
    return req, items[1:]


def _print_prediction_usage(base: str, api_key: str, model_type: str, schema: dict | None = None):
    payload = _prediction_example_payload(model_type, schema)
    example_request = _display_example_request(schema)
    primary_request, extra_params = _split_primary_and_extra_params(example_request)

    console.print(f"  [bold white]Usage[/bold white]")
    console.print()

    if model_type == "numeric" and schema:
        target = schema.get("target", "")
        features = schema.get("features", [])
        feature_types = schema.get("feature_types", {})
        feature_ranges = schema.get("feature_ranges", {})
        sample_input = schema.get("sample_input", {})
        available_targets = schema.get("targets", [])

        if available_targets:
            console.print(f"  [#999999]Available targets:[/#999999] [white]{', '.join(available_targets)}[/white]")
        if features:
            console.print(f"  [#999999]Input:[/#999999] [white]{', '.join(features)}[/white]")
        if target:
            console.print(f"  [#999999]Output:[/#999999] [white]raw numeric prediction for {target}[/white]")
        if feature_ranges:
            console.print(f"  [#999999]Ranges:[/#999999]")
            for key in features:
                bounds = feature_ranges.get(key)
                kind = feature_types.get(key, "numeric")
                if isinstance(bounds, dict):
                    min_v = bounds.get('min')
                    max_v = bounds.get('max')
                    sample_v = sample_input.get(key, min_v)
                    console.print(
                        f"    [white]{key}[/white]  [#777777]{kind} · start {min_v} · end {max_v} · example {sample_v}[/#777777]"
                    )
            console.print()

    if not payload:
        console.print(f"  [bold yellow]Schema unavailable[/bold yellow]")
        console.print(f"  [#777777]No prediction example was printed because this deployed model did not provide a valid input schema.[/#777777]")
        console.print(f"  [#777777]Use GET {base}/predict/schema with the API key, or retrain/redeploy the model.[/#777777]")
        return

    console.print(f"  [#999999]curl:[/#999999]")
    console.print(f'  [white]curl -X POST {base}/predict \\[/white]')
    console.print(f'  [white]  -H "Authorization: Bearer {api_key}" \\[/white]')
    console.print(f'  [white]  -H "Content-Type: application/json" \\[/white]')
    console.print(f"  [white]  -d '{payload}'[/white]")
    console.print()
    console.print(f"  [#999999]Python:[/#999999]")
    console.print(f"  [white]import requests[/white]")
    console.print(f"  [white]r = requests.post([/white]")
    console.print(f'  [white]    "{base}/predict",[/white]')
    console.print(f'  [white]    headers={{"Authorization": "Bearer {api_key}"}},[/white]')
    if model_type == "numeric" and primary_request:
        console.print(f"  [white]    json={json.dumps(primary_request, ensure_ascii=False)}[/white]")
    else:
        console.print(f"  [white]    json={payload}[/white]")
    console.print(f"  [white])[/white]")
    console.print(f"  [white]print(r.json())[/white]")
    console.print()
    console.print(f"  [#999999]JavaScript:[/#999999]")
    console.print(f'  [white]const res = await fetch("{base}/predict", {{[/white]')
    console.print(f'  [white]  method: "POST",[/white]')
    console.print(f'  [white]  headers: {{[/white]')
    console.print(f'  [white]    "Authorization": "Bearer {api_key}",[/white]')
    console.print(f'  [white]    "Content-Type": "application/json"[/white]')
    console.print(f"  [white]  }},[/white]")
    if model_type == "numeric" and primary_request:
        console.print(f'  [white]  body: JSON.stringify({json.dumps(primary_request, ensure_ascii=False)})[/white]')
        for key, value in extra_params:
            console.print(f'  [white]  //"data": {{"{key}": {json.dumps(value, ensure_ascii=False)}}}[/white]')
        console.print(f"  [white]  // you can add multiple data parameters in one request[/white]")
    else:
        console.print(f"  [white]  body: JSON.stringify({payload})[/white]")
    console.print(f"  [white]}});[/white]")
    console.print(f"  [white]const data = await res.json();[/white]")
    console.print(f"  [white]console.log(data)[/white]")


def _fetch_prediction_schema(client: httpx.Client, api_key: str) -> dict:
    try:
        r = client.get(
            "/predict/schema",
            headers={"Authorization": f"Bearer {api_key}"},
        )
        r.raise_for_status()
        data = r.json()
        if data.get("ok") and isinstance(data.get("schema"), dict):
            return data["schema"]
    except Exception:
        pass
    return {}


def handle_deploy(client: httpx.Client, team_key: str):
    if not active_project:
        console.print(f"  [{SALMON}]![/{SALMON}] [white]Select a project first with[/white] [{SALMON}]/projects[/{SALMON}]")
        return

    console.print()
    console.print(f"  [bold white]Deploy[/bold white]  [{SALMON}]{active_project['name']}[/{SALMON}]")
    console.print()

    options = ["Deploy model", "List deployments", "Undeploy", "Back"]
    inline = "    ".join(f"[{SALMON}]{i+1}[/{SALMON}] [white]{o}[/white]" for i, o in enumerate(options))
    console.print(f"  {inline}")
    console.print()

    choice = _numbered_menu(options)
    if choice < 0 or options[choice] == "Back":
        return

    selected = options[choice]

    if selected == "Deploy model":
        sp = Spinner("Deploying model")
        sp.start()
        try:
            r = client.post("/deploy", json={
                "team": team_key,
                "action": "deploy",
                "project_index": active_project_index,
            })
            r.raise_for_status()
            d = r.json()
        except Exception as e:
            sp.stop(f"Error: {e}", ok=False)
            return

        if not d.get("ok"):
            sp.stop(d.get("error", "Error"), ok=False)
            return

        api_key = d["api_key"]
        model_type = d.get("model_type", "")
        model_name = d.get("model_name", "")
        schema = d.get("schema", {})
        if not schema:
            schema = _fetch_prediction_schema(client, api_key)
        sp.stop(f"Deployed {model_name} ({model_type})", ok=True)

        console.print()
        console.print(f"  [bold white]API Key[/bold white]")
        console.print(f"    [{SALMON}]{api_key}[/{SALMON}]")
        console.print()
        thin_rule()
        console.print()

        base = BASE_URL
        _print_prediction_usage(base, api_key, model_type, schema)

        console.print()

    elif selected == "List deployments":
        sp = Spinner("Loading deployments")
        sp.start()
        try:
            r = client.post("/deploy", json={"team": team_key, "action": "list"})
            r.raise_for_status()
            d = r.json()
        except Exception as e:
            sp.stop(f"Error: {e}", ok=False)
            return

        deploys = d.get("deployments", [])
        sp.stop(f"{len(deploys)} deployment(s)", ok=True)
        console.print()

        if not deploys:
            console.print(f"    [#777777]No active deployments[/#777777]")
        else:
            for dep in deploys:
                badge = f"[#00CC00]●[/#00CC00]" if dep.get("model_type") == "text" else f"[#4488FF]●[/#4488FF]"
                console.print(f"    {badge} [bold white]{dep.get('project', '')}[/bold white]  "
                              f"[white]{dep.get('model_name', '')}[/white]  "
                              f"[#777777]{dep.get('model_type', '')}[/#777777]")
                console.print(f"       [#777777]key: {dep.get('api_key', '')}  deployed: {dep.get('deployed_at', '')}[/#777777]")
        console.print()

    elif selected == "Undeploy":
        sp = Spinner("Undeploying")
        sp.start()
        try:
            r = client.post("/deploy", json={
                "team": team_key,
                "action": "undeploy",
                "project_index": active_project_index,
            })
            r.raise_for_status()
            d = r.json()
        except Exception as e:
            sp.stop(f"Error: {e}", ok=False)
            return

        if d.get("ok"):
            sp.stop(d.get("message", "Done"), ok=True)
        else:
            sp.stop(d.get("error", "Error"), ok=False)
        console.print()


# ──────────────────────────────────────────────
#  Main
# ──────────────────────────────────────────────

def main():
    # Backend check
    try:
        r = httpx.get(f"{BASE_URL}/health", timeout=5)
        r.raise_for_status()
    except Exception:
        console.print(f"[error]Cannot reach backend at {BASE_URL}. Is it running?[/error]")
        raise SystemExit(1)

    client = create_client()

    # ── Login (cached or interactive) ──
    session_data = do_login(client)
    if session_data is None:
        console.print("\n  Goodbye.\n", style="dim")
        return

    user = session_data["user"]
    user_type = user["user_type"]
    team = session_data.get("team", {})

    # ── Admin first-time setup (per team) ──
    if user_type == "admin" and not team.get("setup_done", False):
        run_admin_setup(client, user.get("team", "core"))
        # Re-fetch team data after setup
        try:
            r = client.post("/auth", json=load_session())
            r.raise_for_status()
            fresh = r.json()
            if fresh.get("ok"):
                session_data["team"] = fresh.get("team", {})
        except Exception:
            pass

    # ── Banner ──
    print_banner(session_data)

    # ── Input loop ──
    auto_accept = [True]
    kb = KeyBindings()

    @kb.add("s-tab")
    def _toggle(event):
        auto_accept[0] = not auto_accept[0]

    def bottom_toolbar():
        state = "on" if auto_accept[0] else "off"
        name = user.get("name", "")
        parts = [
            ("#FA8072", name),
            ("", "  "),
        ]
        if active_project:
            parts.append(("bold #FA8072", active_project["name"]))
            parts.append(("", "  "))

        # Show agent status
        with agent_lock:
            if agent_status["running"]:
                step = agent_status.get("step", "Working")
                step_num = agent_status.get("step_num", 0)
                total = agent_status.get("total_steps", 4)
                start_t = agent_status.get("start_time", 0)

                # Blinking salmon circle + step indicator
                step_label = f"[{step_num}/{total}]" if step_num else ""
                blink = int(time.time() * 2) % 2 == 0
                parts.append(("#FA8072" if blink else "#3a1e1a", "● "))
                parts.append(("#FA8072" if blink else "#7a4038", f"Agent: {step} {step_label}"))
                parts.append(("", "  "))

                # ETA calculation
                if start_t and step_num:
                    step_elapsed = time.time() - agent_status.get("step_start_time", start_t)
                    current_avg = _STEP_AVG_DURATIONS.get(step_num, 10)
                    remaining_in_step = max(current_avg - step_elapsed, 1)
                    remaining_future = sum(_STEP_AVG_DURATIONS.get(i, 0) for i in range(step_num + 1, total + 1))
                    eta_secs = int(remaining_in_step + remaining_future)

                    if eta_secs >= 60:
                        eta_str = f"~{eta_secs // 60}m {eta_secs % 60}s"
                    else:
                        eta_str = f"~{eta_secs}s"
                    parts.append(("#777777", f"ETA {eta_str}"))
                    parts.append(("", "  "))
            elif agent_status.get("result") is not None:
                # Just finished — show green or red circle briefly
                res = agent_status["result"]
                if res["ok"]:
                    parts.append(("#00CC00", "● "))
                    parts.append(("#00CC00", "Agent: Done"))
                else:
                    parts.append(("#FF4444", "● "))
                    parts.append(("#FF4444", "Agent: Failed"))
                parts.append(("", "  "))

        # Show training status
        with train_lock:
            if train_status["running"]:
                epoch = train_status.get("epoch", 0)
                total_ep = train_status.get("total_epochs", 0)
                loss = train_status.get("loss", 0.0)
                t_start = train_status.get("start_time", 0)

                blink_t = int(time.time() * 2) % 2 == 0
                parts.append(("#FA8072" if blink_t else "#3a1e1a", "● "))
                parts.append(("#FA8072" if blink_t else "#7a4038", f"Train: epoch {epoch}/{total_ep}"))
                if loss:
                    parts.append(("#777777", f"  loss {loss:.4f}"))
                parts.append(("", "  "))

                # ETA based on elapsed time per epoch
                if epoch > 0 and t_start:
                    elapsed = time.time() - t_start
                    per_epoch = elapsed / epoch
                    remaining = int(per_epoch * (total_ep - epoch))
                    if remaining >= 60:
                        eta_str = f"~{remaining // 60}m {remaining % 60}s"
                    else:
                        eta_str = f"~{remaining}s"
                    parts.append(("#777777", f"ETA {eta_str}"))
                    parts.append(("", "  "))
            elif train_status.get("result") is not None:
                res_t = train_status["result"]
                if res_t["ok"]:
                    parts.append(("#00CC00", "● "))
                    parts.append(("#00CC00", "Train: Done"))
                else:
                    parts.append(("#FF4444", "● "))
                    parts.append(("#FF4444", "Train: Failed"))
                parts.append(("", "  "))

        # Show auto-accept only when nothing else is active
        with agent_lock:
            agent_active = agent_status["running"] or agent_status.get("result") is not None
        with train_lock:
            train_active = train_status["running"] or train_status.get("result") is not None
        if not agent_active and not train_active:
            parts.extend([
                ("#777777", "auto-accept"),
                ("", " "),
                ("#FA8072" if auto_accept[0] else "#555555", state),
                ("", "  "),
                ("#555555", "shift+tab"),
            ])
        return FormattedText(parts)

    def _check_agent_result():
        """Check if background agent finished and show result."""
        with agent_lock:
            result = agent_status["result"]
            if result is None:
                return
            if agent_status["running"]:
                return

        console.print()
        if result["ok"]:
            console.print(f"  [#00CC00]●[/#00CC00] [bold white]Done[/bold white]  [white]{result['message']}[/white]")
        else:
            console.print(f"  [#FF4444]●[/#FF4444] [bold red]Failed[/bold red]  [white]{result['message']}[/white]")
        console.print()

        # Clear result after displaying so toolbar returns to normal
        with agent_lock:
            agent_status["result"] = None

    def _check_train_result():
        """Check if background training finished and show result."""
        with train_lock:
            result = train_status["result"]
            if result is None:
                return
            if train_status["running"]:
                return

        console.print()
        if result["ok"]:
            console.print(f"  [#00CC00]●[/#00CC00] [bold white]Done[/bold white]  [white]{result['message']}[/white]")
        else:
            console.print(f"  [#FF4444]●[/#FF4444] [bold red]Failed[/bold red]  [white]{result['message']}[/white]")
        console.print()

        with train_lock:
            train_status["result"] = None

    session = PromptSession(
        style=PT_STYLE,
        completer=SlashCompleter(),
        bottom_toolbar=bottom_toolbar,
        key_bindings=kb,
        refresh_interval=0.5,
    )

    while True:
        try:
            # Check if background tasks finished before showing prompt
            _check_agent_result()
            _check_train_result()

            rule()
            hint = random.choice(HINTS)

            user_input = session.prompt(
                HTML("<prompt>  ~  </prompt>"),
                placeholder=HTML(f"<placeholder>{hint}</placeholder>"),
            ).strip()

            if not user_input:
                continue

            if user_input.lower() in ("/exit", "exit"):
                console.print("\n  Goodbye.\n", style="dim")
                break

            if user_input.lower() == "/logout":
                clear_session()
                console.print(f"\n  [{SALMON}]Signed out.[/{SALMON}]\n")
                break

            if user_input.lower() == "/help":
                handle_help()
                continue

            if user_input.lower() == "/servers":
                handle_servers(client, user.get("team", "core"))
                continue

            if user_input.lower() == "/projects":
                handle_projects(client, user.get("team", "core"), session_data)
                continue

            if user_input.lower().startswith("/generate"):
                desc = user_input[len("/generate"):].strip()
                if not desc:
                    console.print(f"  [{SALMON}]![/{SALMON}] [white]Usage:[/white] [{SALMON}]/generate <description>[/{SALMON}]")
                    console.print(f"    [#777777]e.g. /generate turkish surnames with city and age[/#777777]")
                    continue
                if not active_project:
                    console.print(f"  [{SALMON}]![/{SALMON}] [white]Select a project first with[/white] [{SALMON}]/projects[/{SALMON}]")
                    continue
                handle_chat(client, desc, auto_accept[0], team_key=user.get("team", "core"))
                continue

            if user_input.lower() == "/train":
                handle_train(client, team_key=user.get("team", "core"))
                continue

            if user_input.lower() == "/models":
                handle_models(client, team_key=user.get("team", "core"))
                continue

            if user_input.lower() == "/deploy":
                handle_deploy(client, team_key=user.get("team", "core"))
                continue

            if user_input.lower() == "/info":
                handle_info(client, team_key=user.get("team", "core"))
                continue

            # Unknown command or free text — show error
            if user_input.startswith("/"):
                console.print(f"  [{SALMON}]![/{SALMON}] [white]Unknown command.[/white] [#777777]Type /help for available commands.[/#777777]")
            else:
                console.print(f"  [{SALMON}]![/{SALMON}] [white]Use[/white] [{SALMON}]/generate <description>[/{SALMON}] [white]to create a dataset.[/white]")
                console.print(f"    [#777777]Type /help for all commands.[/#777777]")

        except KeyboardInterrupt:
            console.print("\n  Goodbye.\n", style="dim")
            break
        except EOFError:
            console.print("\n  Goodbye.\n", style="dim")
            break
        except Exception as e:
            console.print(f"  {e}", style="error")


if __name__ == "__main__":
    main()
