# -*- coding: utf-8 -*-
"""
plane_gw.py v2 - OpenAI-compatible gateway over Plane AI (pi.plane.so)
=======================================================================
9Router-style: account pool, round-robin, auto-disable dead accounts, local-no-auth.
NEW in v2:
- TRUE token streaming (SSE delta passthrough, chunked transfer)
- OpenAI tool-calling EMULATION: upstream pi.plane.so accepts NO client tools
  (tools are server-side: execute-action/artifacts/mcp). The gateway injects tool
  schemas into the prompt with a strict JSON protocol, parses model output, and
  returns real OpenAI tool_calls. role=tool results are folded back into the
  next prompt (full transcript per turn - stateless, reliable).
- system prompts, multi-turn history, thinking/websearch via model params

Endpoints: GET /v1/models | POST /v1/chat/completions (stream+sync) | /health
           /admin/pool | /admin/reload
Pool file: plane_accounts.jsonl (FULL_OK records from plane_full_combo.py)
Run: python plane_gw.py [port]   (default 20130)
CLI: base_url=http://127.0.0.1:20130/v1  api_key=any  model=plane/claude-sonnet-5
Env: PLANE_ACCOUNTS=<path>  GW_KEY=<bearer key for remote access>
"""
import json, os, sys, time, re, threading, hashlib
import queue as pyqueue
import urllib.request, urllib.error
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

PI = "https://pi.plane.so"
ACCOUNTS_FILE = os.getenv("PLANE_ACCOUNTS", os.path.join(os.path.dirname(os.path.abspath(__file__)), "plane_accounts.jsonl"))
GWKEY = os.getenv("GW_KEY", "")
POOL_LOCK = threading.Lock()

TOOL_PREAMBLE = """TOOL PROTOCOL:
You have access to these tools:
%s

When you want to call a tool, reply with ONLY a fenced json block in EXACTLY this format:
```json
{"tool_calls": [{"name": "<tool_name>", "arguments": {<json arguments per schema>}}]}
```
Nothing else in the message. You may request several tools in the list.
After the user returns tool results (lines starting with "Tool <name> result:"), continue: call another tool the same way, or give your FINAL answer in plain text (no json block).
"""

class Account:
    def __init__(self, rec):
        self.email = rec.get("email")
        self.cookie = rec.get("session_cookie", "")
        ws = rec.get("workspace") or {}
        self.ws_slug = ws.get("slug")
        self.ws_id = ws.get("id")
        self.user_id = rec.get("user_id")
        self.is_active = bool(self.cookie and self.ws_slug and self.ws_id)
        self.used = 0
        self.last_err = ""
        self.disabled_reason = ""

class Pool:
    def __init__(self, path):
        self.path = path
        self.accounts = []
        self.idx = 0
        self.reload()
    def reload(self):
        accs = []
        if os.path.exists(self.path):
            with open(self.path, encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line: continue
                    try: rec = json.loads(line)
                    except Exception: continue
                    if rec.get("result") == "FULL_OK": accs.append(Account(rec))
        with POOL_LOCK:
            self.accounts = accs
            self.idx = 0
        return len(accs)
    def next(self):
        with POOL_LOCK:
            live = [a for a in self.accounts if a.is_active]
            if not live: return None
            a = live[self.idx % len(live)]
            self.idx += 1
            return a
    def disable(self, acc, reason):
        with POOL_LOCK:
            acc.is_active = False
            acc.disabled_reason = reason
    def state(self):
        with POOL_LOCK:
            return [{"email": a.email, "active": a.is_active, "used": a.used,
                     "err": a.last_err, "disabled": a.disabled_reason} for a in self.accounts]

POOL = Pool(ACCOUNTS_FILE)
MODELS_CACHE = {"ts": 0, "models": []}

# ---------------- upstream ----------------

def _pi_request(path, acc, data=None, timeout=60):
    url = PI + path
    hdrs = {"Cookie": acc.cookie, "User-Agent": "Mozilla/5.0",
            "Origin": "https://app.plane.so", "Referer": "https://app.plane.so/"}
    body = None
    if data is not None:
        hdrs["Content-Type"] = "application/json"
        body = json.dumps(data).encode()
    req = urllib.request.Request(url, data=body, headers=hdrs,
                                 method="POST" if data is not None else "GET")
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return r.status, r.read().decode(errors="ignore")

def get_models(acc):
    now = time.time()
    if MODELS_CACHE["models"] and now - MODELS_CACHE["ts"] < 3600:
        return MODELS_CACHE["models"]
    try:
        st, body = _pi_request("/api/v1/chat/get-models/?workspace_slug=" + str(acc.ws_slug), acc)
        js = json.loads(body)
        models = [m for m in js.get("models", []) if m.get("type") == "language_model"]
        if models:
            MODELS_CACHE["ts"] = now
            MODELS_CACHE["models"] = models
        return models
    except Exception:
        return MODELS_CACHE["models"]

def _start_stream(acc, prompt, model):
    st, body = _pi_request("/api/v1/chat/initialize-chat/", acc,
                           {"workspace_in_context": False, "workspace_slug": acc.ws_slug})
    chat_id = json.loads(body).get("chat_id")
    if not chat_id:
        raise RuntimeError("init_chat %s: %s" % (st, body[:150]))
    payload = {"query": prompt, "llm": model, "is_new": True, "is_temp": True,
               "workspace_in_context": False, "mode": "ask",
               "workspace_id": acc.ws_id, "workspace_slug": acc.ws_slug,
               "chat_id": chat_id, "context": {}}
    st, body = _pi_request("/api/v1/chat/queue-answer/", acc, payload)
    tok = json.loads(body).get("stream_token")
    if not tok:
        raise RuntimeError("queue %s: %s" % (st, body[:150]))
    acc.used += 1
    return tok

def _iter_sse(tok, acc, timeout=180):
    """Yield raw text chunks from the upstream SSE (delta events only)."""
    url = PI + "/api/v1/chat/stream-answer/" + tok
    req = urllib.request.Request(url, headers={"Cookie": acc.cookie, "User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        event = None
        for raw in r:
            line = raw.decode(errors="ignore").rstrip("\n")
            if line.startswith("event:"):
                event = line[6:].strip()
            elif line.startswith("data:"):
                data = line[5:].strip()
                if event == "delta":
                    try:
                        chunk = json.loads(data).get("chunk", "")
                        if chunk: yield chunk
                    except Exception: pass
                elif event == "error":
                    raise RuntimeError("stream error: " + data[:200])
                elif event == "done":
                    break

def plane_chat(acc, prompt, model, timeout=180):
    tok = _start_stream(acc, prompt, model)
    return "".join(_iter_sse(tok, acc, timeout))

# ---------------- OpenAI request -> Plane prompt ----------------

def _msg_text(m):
    c = m.get("content", "")
    if isinstance(c, list):
        return " ".join(str(x.get("text", "")) for x in c if isinstance(x, dict))
    if c is None:
        # assistant tool_calls message
        calls = m.get("tool_calls") or []
        return "Assistant requested tool calls: " + json.dumps(
            [{"name": (t.get("function") or {}).get("name"),
              "arguments": (t.get("function") or {}).get("arguments")} for t in calls])
    return str(c)

def build_prompt(req):
    msgs = req.get("messages", [])
    tools = req.get("tools") or []
    sys_parts, conv = [], []
    for m in msgs:
        role = m.get("role", "user")
        text = _msg_text(m)
        if role == "system":
            sys_parts.append(text)
        elif role == "tool":
            name = m.get("name", "tool")
            conv.append("Tool %s result: %s" % (name, text))
        elif role == "assistant":
            tc = m.get("tool_calls")
            if tc:
                brief = [{"name": (t.get("function") or {}).get("name"),
                          "arguments": (t.get("function") or {}).get("arguments")} for t in tc]
                conv.append("Assistant: " + json.dumps({"tool_calls": brief}))
            elif text:
                conv.append("Assistant: " + text)
        else:
            conv.append("User: " + text)
    prompt = ""
    if sys_parts:
        prompt += "System instructions:\n" + "\n".join(sys_parts) + "\n\n"
    if tools:
        schemas = []
        for t in tools:
            fn = t.get("function", t)
            schemas.append(json.dumps({"name": fn.get("name"), "description": fn.get("description", ""),
                                       "parameters": fn.get("parameters", {})}, ensure_ascii=False))
        prompt += TOOL_PREAMBLE % ("\n".join(schemas)) + "\n"
    prompt += "Conversation:\n" + "\n".join(conv) + "\n\nAssistant:"
    return prompt

# ---------------- tool-call parsing ----------------

TOOL_BLOCK = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.S)

def parse_tool_calls(text):
    """Return (clean_text, [calls]) - calls extracted from fenced json blocks."""
    calls, clean = [], text
    for mm in list(TOOL_BLOCK.finditer(text)):
        try:
            js = json.loads(mm.group(1))
        except Exception:
            continue
        if isinstance(js, dict) and isinstance(js.get("tool_calls"), list) and js["tool_calls"]:
            ok = True
            for c in js["tool_calls"]:
                if not isinstance(c, dict) or "name" not in c:
                    ok = False; break
            if ok:
                calls.extend(js["tool_calls"])
                clean = clean.replace(mm.group(0), "")
    return clean.strip(), calls

class StreamToolFilter:
    """Incremental filter: pass through normal text, hold potential tool blocks."""
    def __init__(self):
        self.buf = ""
        self.in_fence = False
        self.fence_buf = ""
        self.calls = []
        self.clean_parts = []
    def feed(self, chunk):
        """Return text safe to emit now."""
        out = []
        self.buf += chunk
        while True:
            if not self.in_fence:
                idx = self.buf.find("```")
                if idx == -1:
                    keep = max(0, len(self.buf) - 3)  # hold back possible partial fence
                    if keep:
                        out.append(self.buf[:keep]); self.clean_parts.append(self.buf[:keep])
                        self.buf = self.buf[keep:]
                    break
                else:
                    if idx:
                        out.append(self.buf[:idx]); self.clean_parts.append(self.buf[:idx])
                    self.in_fence = True
                    self.fence_buf = ""
                    self.buf = self.buf[idx+3:]
            else:
                end = self.buf.find("```")
                if end == -1:
                    self.fence_buf += self.buf; self.buf = ""
                    if len(self.fence_buf) > 20000:  # runaway fence = not a tool block
                        out.append("```" + self.fence_buf); self.clean_parts.append("```" + self.fence_buf)
                        self.fence_buf = ""; self.in_fence = False
                    break
                else:
                    block = self.fence_buf + self.buf[:end]
                    self.buf = self.buf[end+3:]
                    self.in_fence = False
                    js = None
                    try:
                        js = json.loads(block.strip())
                        is_tool = isinstance(js, dict) and isinstance(js.get("tool_calls"), list) and js["tool_calls"] \
                                  and all(isinstance(c, dict) and "name" in c for c in js["tool_calls"])
                    except Exception:
                        is_tool = False
                    if is_tool:
                        self.calls.extend(js["tool_calls"])
                    else:
                        txt = "```" + block + "```"
                        out.append(txt); self.clean_parts.append(txt)
        return "".join(out)
    def finish(self):
        """Flush held-back content at stream end. Returns (final_clean_text, calls)."""
        tail = ""
        if self.in_fence:
            tail = "```" + self.fence_buf
        tail += self.buf
        if tail:
            self.clean_parts.append(tail)
        return "".join(self.clean_parts).strip(), self.calls

# ---------------- streaming worker ----------------

def plane_stream_worker(acc, prompt, model, out_q, timeout=180):
    try:
        tok = _start_stream(acc, prompt, model)
        flt = StreamToolFilter()
        for chunk in _iter_sse(tok, acc, timeout):
            emit = flt.feed(chunk)
            if emit: out_q.put(("delta", emit))
        clean, calls = flt.finish()
        out_q.put(("final", (clean, calls)))
    except Exception as e:
        out_q.put(("err", str(e)[:200]))
    out_q.put(("end", None))

# ---------------- HTTP handler ----------------

def strip_model_prefix(m):
    return m.split("/", 1)[1] if m.startswith("plane/") else m

def mk_tool_calls_message(calls, model):
    tcs = []
    for i, c in enumerate(calls):
        args = c.get("arguments", {})
        if not isinstance(args, str): args = json.dumps(args, ensure_ascii=False)
        tcs.append({"id": "call_plane_%d_%d" % (int(time.time()), i), "type": "function",
                    "function": {"name": c.get("name", ""), "arguments": args}})
    return {"role": "assistant", "content": None, "tool_calls": tcs}

class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    def log_message(self, *a): pass
    def _local(self): return self.client_address[0] in ("127.0.0.1", "::1")
    def _authed(self):
        if not GWKEY: return True
        if self._local(): return True
        auth = self.headers.get("Authorization", "")
        if auth.startswith("Bearer ") and auth[7:] == GWKEY: return True
        if self.headers.get("x-api-key") == GWKEY: return True
        return False
    def _send(self, code, obj):
        data = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(data)
    def do_OPTIONS(self):
        self.send_response(204)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Headers", "*")
        self.send_header("Access-Control-Allow-Methods", "*")
        self.send_header("Content-Length", "0")
        self.end_headers()
    def do_GET(self):
        path = self.path.split("?")[0]
        if path == "/health":
            return self._send(200, {"ok": True, "pool_live": len([a for a in POOL.accounts if a.is_active]), "v": 2})
        if path == "/admin/pool":
            if not self._local(): return self._send(403, {"error": "local only"})
            return self._send(200, POOL.state())
        if path == "/admin/reload":
            if not self._local(): return self._send(403, {"error": "local only"})
            return self._send(200, {"loaded": POOL.reload()})
        if not self._authed():
            return self._send(401, {"error": {"message": "API key required for remote API access", "type": "auth"}})
        if path == "/v1/models":
            acc = POOL.next()
            models = get_models(acc) if acc else []
            data = [{"id": "plane/" + m["id"], "object": "model", "created": int(time.time()),
                     "owned_by": m.get("provider", "plane")} for m in models]
            return self._send(200, {"object": "list", "data": data})
        return self._send(404, {"error": "not found"})

    def do_POST(self):
        path = self.path.split("?")[0]
        if path != "/v1/chat/completions":
            return self._send(404, {"error": "not found"})
        if not self._authed():
            return self._send(401, {"error": {"message": "API key required for remote API access", "type": "auth"}})
        n = int(self.headers.get("Content-Length", 0))
        try:
            req = json.loads(self.rfile.read(n) or b"{}")
        except Exception:
            return self._send(400, {"error": "bad json"})
        model = strip_model_prefix(req.get("model", "claude-sonnet-5"))
        stream = bool(req.get("stream"))
        prompt = build_prompt(req)
        tried, last_err = 0, ""
        while tried < 3:
            acc = POOL.next()
            if acc is None:
                return self._send(503, {"error": {"message": "No active accounts in pool", "type": "pool"}})
            tried += 1
            try:
                if stream:
                    return self._stream_response(acc, prompt, model)
                text = plane_chat(acc, prompt, model)
                clean, calls = parse_tool_calls(text)
                msg = mk_tool_calls_message(calls, model) if calls else {"role": "assistant", "content": clean}
                resp = {"id": "chatcmpl-plane-" + hashlib.sha1(prompt.encode()).hexdigest()[:12],
                        "object": "chat.completion", "created": int(time.time()), "model": model,
                        "choices": [{"index": 0, "message": msg,
                                     "finish_reason": "tool_calls" if calls else "stop"}],
                        "usage": {"prompt_tokens": len(prompt)//4, "completion_tokens": len(text)//4,
                                  "total_tokens": (len(prompt)+len(text))//4}}
                return self._send(200, resp)
            except urllib.error.HTTPError as e:
                last_err = "HTTP %s" % e.code
                acc.last_err = last_err
                if e.code in (401, 403): POOL.disable(acc, "auth %s" % e.code)
                elif e.code == 429: POOL.disable(acc, "credits/ratelimit 429")
            except Exception as e:
                last_err = str(e)[:150]
                acc.last_err = last_err
        return self._send(502, {"error": {"message": "upstream failed: " + last_err, "type": "upstream"}})

    def _stream_response(self, acc, prompt, model):
        q = pyqueue.Queue()
        t = threading.Thread(target=plane_stream_worker, args=(acc, prompt, model, q), daemon=True)
        t.start()
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Transfer-Encoding", "chunked")
        self.end_headers()
        cid = "chatcmpl-plane-" + hashlib.sha1((prompt + str(time.time())).encode()).hexdigest()[:12]
        def send_chunk(payload):
            b = payload.encode()
            self.wfile.write(("%X\r\n" % len(b)).encode() + b + b"\r\n")
            self.wfile.flush()
        def base(delta, finish=None):
            return {"id": cid, "object": "chat.completion.chunk", "created": int(time.time()),
                    "model": model, "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}
        send_chunk("data: " + json.dumps(base({"role": "assistant", "content": ""})) + "\n\n")
        calls = []
        while True:
            kind, payload = q.get()
            if kind == "delta":
                send_chunk("data: " + json.dumps(base({"content": payload})) + "\n\n")
            elif kind == "final":
                clean, calls = payload
                if calls:
                    for i, c in enumerate(calls):
                        args = c.get("arguments", {})
                        if not isinstance(args, str): args = json.dumps(args, ensure_ascii=False)
                        d = {"tool_calls": [{"index": i, "id": "call_plane_%d_%d" % (int(time.time()), i),
                                              "type": "function",
                                              "function": {"name": c.get("name", ""), "arguments": args}}]}
                        send_chunk("data: " + json.dumps(base(d)) + "\n\n")
            elif kind == "err":
                send_chunk("data: " + json.dumps({"error": {"message": payload}}) + "\n\n")
            elif kind == "end":
                break
        finish = "tool_calls" if calls else "stop"
        try:
            send_chunk("data: " + json.dumps(base({}, finish)) + "\n\n")
            send_chunk("data: [DONE]\n\n")
            self.wfile.write(b"0\r\n\r\n")
            self.wfile.flush()
        except Exception:
            pass

def main():
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 20130
    n = POOL.reload()
    print("[plane_gw v2] pool: %d accounts from %s" % (n, ACCOUNTS_FILE))
    print("[plane_gw v2] listening on http://127.0.0.1:%d/v1 (auth=%s, tools=emulated, streaming=true)" % (port, "on" if GWKEY else "off"))
    ThreadingHTTPServer(("127.0.0.1", port), Handler).serve_forever()

if __name__ == "__main__":
    main()