"""WebUI session management against a fake OpenAI-compatible API.

  python3 scripts/e2e_webui_sessions.py      (CST_AGENT= picks the binary)

Checks: retry runs the last prompt again and edit-and-resend replaces it
(the model sees the new text, the session keeps one copy); retry is a 409
while a turn runs and when there is nothing to retry; export returns the
transcript as Markdown; rename, clear and delete work on the store, and
deleting the open session moves to a new one; history messages carry the
note / tool_result flags. HOME points at a temp dir.
"""
import json, os, shutil, subprocess, sys, tempfile, threading, time, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

BIN = os.environ.get("CST_AGENT", os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "cst-agent"))
SEEN = []


def last_prompt(msgs):
    """The prompts are one line; a note before one shares its user message."""
    for m in reversed(msgs):
        c = m.get("content")
        if m.get("role") == "user" and isinstance(c, str):
            return c.rstrip("\n").split("\n")[-1]
    return ""


class Api(BaseHTTPRequestHandler):
    """Answers "回答：<last prompt>"; a prompt with "slow" takes 2 s."""
    def log_message(self, *a):
        pass

    def do_GET(self):
        body = json.dumps({"data": [{"id": "deepseek-flash"}]}).encode()
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_POST(self):
        req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))) or b"{}")
        msgs = req.get("messages", [])
        p = last_prompt(msgs)
        tail = msgs[-1] if msgs else {}
        tc = tail.get("content") if isinstance(tail.get("content"), str) else ""
        answered = tail.get("role") == "tool" or tc.startswith("result ")
        if not answered:
            SEEN.append(p)
        if "slow" in p:
            time.sleep(2)
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.end_headers()
        if answered:
            chunks = ({"delta": {"role": "assistant", "content": "读完了"}},
                      {"delta": {}, "finish_reason": "stop"})
        elif p.startswith("native"):
            args = json.dumps({"path": "hello.txt"})
            chunks = ({"delta": {"role": "assistant", "tool_calls": [
                          {"index": 0, "id": "c1", "type": "function",
                           "function": {"name": "read_file", "arguments": args}}]}},
                      {"delta": {}, "finish_reason": "tool_calls"})
        elif p.startswith("ctl"):
            chunks = ({"delta": {"role": "assistant",
                                 "content": "我来读。\ncall t1 read_file\narg path 9\nhello.txtend\n"}},
                      {"delta": {}, "finish_reason": "stop"})
        else:
            chunks = ({"delta": {"role": "assistant", "content": "回答：" + p}},
                      {"delta": {}, "finish_reason": "stop"})
        for o in chunks:
            chunk = {"id": "x", "object": "chat.completion.chunk", "model": "deepseek-flash",
                     "choices": [dict(index=0, **o)]}
            self.wfile.write(b"data: " + json.dumps(chunk, ensure_ascii=False).encode() + b"\n\n")
            self.wfile.flush()
        self.wfile.write(b"data: [DONE]\n\n")
        self.wfile.flush()


fails = 0


def check(ok, msg):
    global fails
    print(("PASS " if ok else "FAIL ") + msg)
    fails += 0 if ok else 1


def wait(fn, t=15):
    end = time.time() + t
    while time.time() < end:
        v = fn()
        if v:
            return v
        time.sleep(0.1)
    return None


def call(base, path, obj=None, method=None, raw_reply=False):
    data = None if obj is None else json.dumps(obj).encode()
    req = urllib.request.Request(base + path, data=data, method=method or ("POST" if data else "GET"),
                                 headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=20) as r:
            raw = r.read()
            if raw_reply:
                return r.status, r.headers.get("Content-Type", ""), raw
            return r.status, json.loads(raw or b"{}")
    except urllib.error.HTTPError as e:
        raw = e.read()
        if raw_reply:
            return e.code, e.headers.get("Content-Type", ""), raw
        try:
            return e.code, json.loads(raw or b"{}")
        except ValueError:
            return e.code, {}


class Sse:
    def __init__(self, base):
        self.lines = []
        self.r = urllib.request.urlopen(base + "/api/events", timeout=60)
        threading.Thread(target=self.pump, daemon=True).start()

    def pump(self):
        try:
            for raw in self.r:
                self.lines.append(raw)
        except Exception:
            pass

    def count(self, name):
        return sum(1 for l in self.lines if l.strip() == b"event: " + name.encode())

    def close(self):
        try:
            self.r.close()
        except Exception:
            pass


api = ThreadingHTTPServer(("127.0.0.1", 0), Api)
threading.Thread(target=api.serve_forever, daemon=True).start()
ws = tempfile.mkdtemp(prefix="cst-sess-")
home = tempfile.mkdtemp(prefix="cst-home-")
env = dict(os.environ, HOME=home, XDG_CONFIG_HOME=home + "/.config",
           XDG_DATA_HOME=home + "/.local/share", XDG_CACHE_HOME=home + "/.cache")
for k in ("DEEPSEEK_API_KEY", "OPENAI_API_KEY", "CST_API_KEY"):
    env.pop(k, None)
log = open(ws + "/server.log", "w+")
proc = subprocess.Popen([BIN, "--web", "--web-port", "0", "--workspace", ws, "--new", "--api-key", "sk-x",
                         "--provider", "deepseek", "--model", "deepseek-flash",
                         "--api-base", "http://127.0.0.1:%d/v1" % api.server_address[1],
                         "--protocol", "native-eos", "--approval", "auto"],
                        stdout=log, stderr=subprocess.STDOUT, env=env)


def port():
    log.seek(0)
    for l in log.read().splitlines():
        if l.startswith("webui: http://127.0.0.1:"):
            return l.split(":")[3].split("/")[0]


def messages(base):
    return call(base, "/api/session")[1].get("messages", [])


def prompts(base):
    return [m["content"] for m in messages(base)
            if m["role"] == "user" and not m.get("note") and not m.get("tool_result")]


def answers(base):
    return [m["content"] for m in messages(base) if m["role"] == "assistant"]


def turn(base, sse, path, body, want=202):
    n = sse.count("turn_end")
    code, d = call(base, path, body)
    if code == want == 202:
        wait(lambda: sse.count("turn_end") > n, 20)
    return code, d


try:
    base = "http://127.0.0.1:%s" % wait(port)
    sse = Sse(base)
    time.sleep(0.3)

    # 1. retry with nothing to retry, then a first turn
    code, d = call(base, "/api/retry", {})
    check(code == 409 and "nothing to retry" in d.get("error", ""), "retry on an empty session: 409 (%s)" % code)
    turn(base, sse, "/api/turn", {"text": "第一句"})
    check(prompts(base) == ["第一句"] and answers(base) == ["回答：第一句"], "first turn stored")

    # 2. retry: same prompt again, one copy in the session
    code, d = turn(base, sse, "/api/retry", {})
    check(code == 202 and d.get("text") == "第一句", "retry accepted (%s %s)" % (code, d))
    check(SEEN[-1] == "第一句", "the model saw the same prompt again")
    check(prompts(base) == ["第一句"] and answers(base) == ["回答：第一句"],
          "after retry: one prompt, one answer (%s)" % prompts(base))

    # 3. edit and resend: the new text replaces the old
    code, d = turn(base, sse, "/api/retry", {"text": "改过的问题"})
    check(code == 202 and d.get("text") == "改过的问题", "edit-and-resend accepted")
    check(SEEN[-1] == "改过的问题", "the model saw the edited prompt")
    check(prompts(base) == ["改过的问题"] and answers(base) == ["回答：改过的问题"],
          "after edit: %s / %s" % (prompts(base), answers(base)))
    check(call(base, "/api/retry", {"text": ""})[0] == 400, "an empty edit is a 400")
    check(call(base, "/api/retry", {"text": 5})[0] == 400, "a non-string edit is a 400")

    # 4. a second turn, then retry only touches the last one
    turn(base, sse, "/api/turn", {"text": "第二句"})
    turn(base, sse, "/api/retry", {})
    check(prompts(base) == ["改过的问题", "第二句"] and len(answers(base)) == 2,
          "retry keeps earlier turns: %s" % prompts(base))

    # 5. retry while a turn runs is a 409
    n = sse.count("turn_end")
    call(base, "/api/turn", {"text": "slow 请慢一点"})
    code, d = call(base, "/api/retry", {})
    check(code == 409, "retry during a turn: 409 (%s)" % code)
    wait(lambda: sse.count("turn_end") > n, 20)

    # 6. a note (plan mode changed mid-session) is flagged, and retry skips it
    call(base, "/api/config", {"plan_mode": True})
    turn(base, sse, "/api/turn", {"text": "计划一下"})
    ms = messages(base)
    notes = [m for m in ms if m.get("note")]
    check(len(notes) >= 1 and all(m["role"] == "user" for m in notes),
          "a mid-session note is flagged note:true (%d)" % len(notes))
    turn(base, sse, "/api/retry", {})
    check(SEEN[-1] == "计划一下" and prompts(base)[-1] == "计划一下", "retry after a note reruns the prompt")
    call(base, "/api/config", {"plan_mode": False})

    # 7. export
    code, ctype, raw = call(base, "/api/export", raw_reply=True)
    md = raw.decode("utf-8", "replace")
    check(code == 200 and ctype.startswith("text/markdown"), "export: %s %s" % (code, ctype))
    check(md.startswith("# ") and "## User" in md and "改过的问题" in md and "## Assistant" in md
          and "回答：第二句" in md and "第一句" not in md, "export has the transcript")

    # 8. rename and clear
    st = call(base, "/api/state")[1]
    cur = st.get("session")
    code, _ = call(base, "/api/sessions/rename", {"title": "我的会话"})
    check(code == 200, "rename the open session")
    ss = call(base, "/api/sessions")[1]
    check(any(s["id"] == cur and s["title"] == "我的会话" for s in ss.get("sessions", [])),
          "the list shows the new title")
    code, _, raw = call(base, "/api/export", raw_reply=True)
    check(raw.decode("utf-8", "replace").startswith("# 我的会话"), "export uses the title")
    long = "长" * 100
    call(base, "/api/sessions/rename", {"id": cur, "title": long})
    t = [s["title"] for s in call(base, "/api/sessions")[1]["sessions"] if s["id"] == cur][0]
    check(0 < len(t.encode()) <= 120 and t == "长" * len(t), "a long title is cut on a character (%d bytes)"
          % len(t.encode()))
    call(base, "/api/sessions/rename", {"title": ""})
    t = [s["title"] for s in call(base, "/api/sessions")[1]["sessions"] if s["id"] == cur][0]
    check(t == "改过的问题", "an empty title falls back to the first prompt (%r)" % t)
    check(call(base, "/api/sessions/rename", {"id": "nope", "title": "x"})[0] == 404, "rename unknown: 404")
    check(call(base, "/api/sessions/rename", {"id": "../x", "title": "x"})[0] == 404, "rename ../x: 404")
    check(call(base, "/api/sessions/rename", {})[0] == 400, "rename without a title: 400")

    # 9. delete another session, then the open one
    code, d = call(base, "/api/sessions/new", {})
    second = d.get("id")
    check(code == 200 and second and second != cur, "a second session")
    code, d = call(base, "/api/sessions/delete", {"id": cur})
    check(code == 200 and d.get("current") == second, "delete a closed session (%s %s)" % (code, d))
    ids = [s["id"] for s in call(base, "/api/sessions")[1]["sessions"]]
    check(cur not in ids and not os.path.exists(os.path.join(ws, ".agent", "sessions", cur + ".jsonl")),
          "it is gone from the list and the disk")
    n = sse.count("session")
    code, d = call(base, "/api/sessions/delete", {"id": second})
    check(code == 200 and d.get("current") and d["current"] != second, "delete the open session (%s)" % d)
    check(call(base, "/api/state")[1].get("session") == d.get("current"), "the state moved to the new one")
    check(wait(lambda: sse.count("session") > n, 5), "a session event was sent")
    check(call(base, "/api/sessions/delete", {"id": "nope"})[0] == 404, "delete unknown: 404")
    check(call(base, "/api/sessions/delete", {})[0] == 400, "delete without an id: 400")

    # 10. history: calls and results come back apart, for native and ctl
    with open(os.path.join(ws, "hello.txt"), "w", encoding="utf-8") as f:
        f.write("你好世界\n")
    turn(base, sse, "/api/turn", {"text": "native 读文件"})
    ms = messages(base)
    a = [m for m in ms if m["role"] == "assistant" and m.get("calls")]
    r = [m for m in ms if m.get("tool_result")]
    check(len(a) == 1 and a[0]["calls"][0].get("name") == "read_file"
          and a[0]["calls"][0].get("args", {}).get("path") == "hello.txt" and a[0]["calls"][0].get("id") == "c1",
          "native: the call is listed with its args (%s)" % (a[0].get("calls") if a else ms))
    check(len(r) == 1 and r[0]["role"] == "tool" and r[0].get("tool_call_id") == "c1"
          and "你好世界" in r[0].get("out", "") and not r[0]["out"].startswith("[tool ")
          and r[0].get("exit") == 0 and r[0].get("truncated") is False,
          "native: the result carries out/exit without the header (%s)" % r)
    check(answers(base)[-1] == "读完了", "native: the turn finished")

    check(call(base, "/api/settings", {"protocol": "ctl"})[0] == 200, "switch to the ctl protocol")
    turn(base, sse, "/api/turn", {"text": "ctl 读文件"})
    ms = messages(base)
    a = [m for m in ms if m["role"] == "assistant" and m.get("calls") and m["calls"][0].get("id") == "t1"]
    r = [m for m in ms if m.get("tool_result") and m.get("tool_call_id") == "t1"]
    check(len(a) == 1 and a[0].get("text", "").strip() == "我来读。" and "call t1" in a[0]["content"]
          and a[0]["calls"][0].get("args", {}).get("path") == "hello.txt"
          and "call t1" not in a[0].get("html", ""),
          "ctl: prose and calls are split (%s)" % (a[0] if a else ms[-4:]))
    check(len(r) == 1 and r[0]["role"] == "user" and "你好世界" in r[0].get("out", "")
          and not r[0]["out"].startswith("result ") and not r[0]["out"].endswith("end")
          and r[0].get("exit") == 0,
          "ctl: the result is flagged with out/exit (%s)" % r)
    check(r and r[0]["content"] not in prompts(base), "ctl: the result is not a prompt")
    call(base, "/api/settings", {"protocol": "native-eos"})
    sse.close()
finally:
    proc.terminate()
    try:
        proc.wait(20)
    except subprocess.TimeoutExpired:
        proc.kill()
        check(False, "server hung on SIGTERM")
    log.close()
    shutil.rmtree(ws, ignore_errors=True)
    shutil.rmtree(home, ignore_errors=True)
    api.shutdown()
print("FAILED:", fails)
sys.exit(1 if fails else 0)