"""Interactive chat with a North ML Flare model.

REMOTE (recommended — runs the model on an SSH device, nothing to download):
  python3 chat.py --ssh mac-mini
  python3 chat.py --ssh mac-mini --model fiber-1

LOCAL (downloads HF weights, runs on this machine):
  python3 chat.py --model North-ML1/fiber-pro

Remote needs: ssh access to the host + NORTH_API_KEY or ~/.config/northml/api.key
Local needs:  pip install torch transformers safetensors huggingface_hub
Commands: /model /search /think /new /quit
"""
import argparse, json, os, sys, threading, importlib.util

IM_S = "<" + "|im_start" + "|>"
IM_E = "<" + "|im_end" + "|>"
TH_S = "<" + "think" + ">"
TH_E = "<" + "/think" + ">"
EOT = "<" + "|endoftext" + "|>"

DIM = "\033[2m"; CYAN = "\033[36m"; BOLD = "\033[1m"; RESET = "\033[0m"; GRAY = "\033[90m"
if not sys.stdout.isatty():
    DIM = CYAN = BOLD = RESET = GRAY = ""


# ============================================================ remote (ssh)

REMOTE_PORT = 8095
LOCAL_PORT = 18095


def get_key():
    k = os.environ.get("NORTH_API_KEY", "").strip()
    if k:
        return k
    p = os.path.expanduser("~/.config/northml/api.key")
    if os.path.exists(p):
        return open(p).read().strip()
    sys.exit("No API key. Set NORTH_API_KEY or write it to " + p)


def health():
    import urllib.request
    try:
        return urllib.request.urlopen(
            "http://127.0.0.1:%d/health" % LOCAL_PORT, timeout=3).status == 200
    except Exception:
        return False


def ensure_tunnel(host):
    import subprocess, time
    if health():
        return
    print("%sopening ssh tunnel to %s...%s" % (DIM, host, RESET))
    subprocess.Popen(
        ["ssh", "-N", "-L", "%d:127.0.0.1:%d" % (LOCAL_PORT, REMOTE_PORT),
         "-o", "ExitOnForwardFailure=yes", "-o", "ServerAliveInterval=15", host],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    for _ in range(40):
        if health():
            print("%stunnel up on :%d%s" % (DIM, LOCAL_PORT, RESET))
            return
        time.sleep(0.5)
    sys.exit("Tunnel failed — is `ssh %s` working and the API running?" % host)


def ask(key, model, messages, search, think):
    import urllib.request
    req = urllib.request.Request(
        "http://127.0.0.1:%d/v1/chat/completions" % LOCAL_PORT,
        data=json.dumps({"model": model, "messages": messages, "stream": True,
                         "search": search, "think": think}).encode(),
        headers={"Authorization": "Bearer " + key,
                 "Content-Type": "application/json"})
    return urllib.request.urlopen(req, timeout=180)


def handle_sse(resp):
    answer, thinking = [], False
    for raw in resp:
        line = raw.decode("utf-8", "replace").strip()
        if not line.startswith("data: "):
            continue
        payload = line[6:]
        if payload == "[DONE]":
            break
        try:
            ev = json.loads(payload)
        except json.JSONDecodeError:
            continue
        agent = ev.get("agent")
        if agent:
            t = agent.get("type")
            if t == "status":
                print("\n%s  ▸ %s%s" % (GRAY, agent.get("stage", t), RESET), flush=True)
            elif t == "route":
                print("%s  ▸ route: search=%s think=%s%s" % (
                    GRAY, agent.get("search"), agent.get("think"), RESET), flush=True)
            elif t == "search":
                print("%s  ▸ query: %s  [%s]%s" % (
                    GRAY, agent.get("query", ""), agent.get("engine", ""), RESET), flush=True)
            elif t == "source":
                mark = "✓" if agent.get("kept") else "·"
                print("%s    %s %s  %s%s" % (GRAY, mark, agent.get("title", "?"),
                                             agent.get("url", ""), RESET))
            elif t == "thinking":
                print("%s%s%s" % (DIM, agent.get("text", ""), RESET), end="", flush=True)
                thinking = True
            elif t == "verify":
                print("%s  ▸ verify: %s%s" % (
                    GRAY, str(agent.get("verdict", agent.get("text", "")))[:160], RESET),
                    flush=True)
            elif t == "error":
                print("\033[31m  error: %s%s" % (agent.get("message", agent), RESET))
            continue
        for ch in ev.get("choices") or []:
            d = ch.get("delta", {}).get("content")
            if d:
                if thinking:
                    print()
                    thinking = False
                print(d, end="", flush=True)
                answer.append(d)
    print()
    return "".join(answer)


def remote_loop(host, model):
    import urllib.error
    key = get_key()
    ensure_tunnel(host)
    search, think = "auto", "auto"
    history = []
    print("%sNorth ML chat — model on %s%s" % (BOLD, host, RESET))
    print("model: %s%s%s  ·  /model /search /think /new /quit\n" % (CYAN, model, RESET))
    while True:
        try:
            q = input("%syou > %s" % (BOLD, RESET)).strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not q:
            continue
        if q.startswith("/"):
            parts = q.split()
            cmd = parts[0]
            if cmd in ("/quit", "/q"):
                break
            if cmd == "/new":
                history = []
                print("%s(history cleared)%s" % (DIM, RESET))
                continue
            if cmd == "/model" and len(parts) > 1:
                model = parts[1]
                continue
            if cmd == "/search" and len(parts) > 1:
                search = parts[1]
                continue
            if cmd == "/think" and len(parts) > 1:
                think = parts[1]
                continue
            print("%scommands: /model ink-1|fiber-1|fiber-1-pro  "
                  "/search auto|on|off  /think auto|on|off  /new  /quit%s" % (DIM, RESET))
            continue
        history.append({"role": "user", "content": q})
        try:
            ans = handle_sse(ask(key, model, history, search, think))
            history.append({"role": "assistant", "content": ans})
        except urllib.error.HTTPError as e:
            print("%sHTTP %s: %s%s" % (DIM, e.code, e.read()[:200].decode("utf-8", "replace"), RESET))
            history.pop()
        except Exception as e:
            print("%serror: %s%s" % (DIM, e, RESET))
            history.pop()


# ============================================================== local (hf)

def pick_device():
    import torch
    if torch.cuda.is_available():
        return "cuda", torch.bfloat16
    if torch.backends.mps.is_available():
        return "mps", torch.float16
    return "cpu", torch.float32


def load(repo):
    import torch
    from huggingface_hub import snapshot_download
    from transformers import AutoModelForCausalLM, AutoTokenizer

    local = snapshot_download(repo) if not os.path.isdir(repo) else repo
    device, dtype = pick_device()
    tok = AutoTokenizer.from_pretrained(local)
    model = AutoModelForCausalLM.from_pretrained(local, dtype=dtype).to(device).eval()

    arch_p = os.path.join(local, "starlight_arch.py")
    reg_p = os.path.join(local, "agent_register.pt")
    meta_p = os.path.join(local, "starlight.json")
    if all(os.path.exists(p) for p in (arch_p, reg_p, meta_p)):
        spec = importlib.util.spec_from_file_location("starlight_arch", arch_p)
        arch = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(arch)
        meta = json.load(open(meta_p))
        regs = arch.attach_registers(model, arch.hidden_size_of(model),
                                     slots=int(meta.get("slots", 16)))
        regs.load_state_dict(torch.load(reg_p, map_location="cpu")["registers"])
        regs.to(device=device, dtype=dtype)
        print("registers attached (%s slots)" % meta.get("slots", 16))
    print("loaded on %s (%s)" % (device, dtype))
    return model, tok, device


def render(history, think):
    s = ""
    for role, text in history:
        s += IM_S + role + "\n" + text + IM_E + "\n"
    s += IM_S + "assistant\n"
    if think:
        s += TH_S + "\n"
    return s


def gen_stream(model, tok, device, input_ids, cap, eos_ids):
    import torch
    from transformers import TextIteratorStreamer
    streamer = TextIteratorStreamer(tok, skip_special_tokens=True, skip_prompt=True)
    holder = {}

    def _run():
        holder["out"] = model.generate(
            input_ids=input_ids, max_new_tokens=cap, do_sample=False,
            repetition_penalty=1.05, eos_token_id=eos_ids,
            pad_token_id=tok.eos_token_id, streamer=streamer)

    t = threading.Thread(target=_run)
    t.start()
    text = ""
    for chunk in streamer:
        text += chunk
        sys.stdout.write(chunk)
        sys.stdout.flush()
    t.join()
    return text, holder["out"]


def chat_turn(model, tok, device, history, think):
    import torch
    im_end = tok.convert_tokens_to_ids(IM_E)
    eot = tok.convert_tokens_to_ids(EOT)
    th_e = tok.convert_tokens_to_ids(TH_E)
    eos = list({i for i in (tok.eos_token_id, im_end, eot) if i is not None})

    prompt = render(history, think)
    ids = tok(prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device)

    if think:
        sys.stdout.write("\n[think] ")
        th_text, out1 = gen_stream(model, tok, device, ids, 220, eos + [th_e])
        if th_e not in out1[0].tolist()[ids.shape[1]:]:
            extra = tok(TH_E + "\n", add_special_tokens=False,
                        return_tensors="pt").input_ids.to(device)
            ids2 = torch.cat([out1, extra], dim=1)
        else:
            ids2 = out1
        sys.stdout.write("\n[answer] ")
        ans_text, _ = gen_stream(model, tok, device, ids2, 500, eos)
        sys.stdout.write("\n")
        return ans_text
    ans_text, _ = gen_stream(model, tok, device, ids, 500, eos)
    sys.stdout.write("\n")
    return ans_text


def local_loop(repo, think):
    print("downloading %s ..." % repo)
    model, tok, device = load(repo)
    history = []
    print("\nReady. Ask anything.   /think on|off   /new   /quit\n")
    while True:
        try:
            q = input("you> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not q:
            continue
        if q in ("/quit", "/q", "exit"):
            break
        if q == "/new":
            history = []
            print("(cleared)\n")
            continue
        if q.startswith("/think"):
            think = "off" not in q
            print("think =", "on" if think else "off", "\n")
            continue
        history.append(("user", q))
        ans = chat_turn(model, tok, device, history, think)
        history.append(("assistant", ans.strip()))
        sys.stdout.write("\n")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ssh", metavar="HOST", default="",
                    help="ssh host running the North ML API (e.g. mac-mini)")
    ap.add_argument("--model", default="",
                    help="remote: ink-1|fiber-1|fiber-1-pro  ·  local: HF repo id")
    ap.add_argument("--no-think", action="store_true")
    a = ap.parse_args()

    if a.ssh:
        remote_loop(a.ssh, a.model or "fiber-1-pro")
    else:
        local_loop(a.model or "North-ML1/fiber-pro", not a.no_think)


if __name__ == "__main__":
    main()