#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
14_msg_proxy.py - a small logging / rewriting proxy between Claude Code and llama-server's /v1/messages endpoint.

  python 14_msg_proxy.py            listens on 127.0.0.1:18082, forwards to 127.0.0.1:18080, streams the answer back, and appends every
                                    /v1/messages request body to /content/proxy_log.jsonl.  fix() is the hook for rewriting a request
                                    (the chat template of Qwen3.8 rejects some shapes that Claude Code sends).
"""
import http.client
import http.server
import json
import time

UP = ("127.0.0.1", 18080)
CRLF = bytes([13, 10])


def fix(d):
    """Qwen3.8's chat template raises 'System message must be at the beginning' for a system-role entry in `messages`, which Claude Code sends
    after the user turn.  Fold every such entry into the last user message as an extra text block (nothing is dropped)."""
    msgs = d.get("messages") or []
    keep, extra = [], []
    for m in msgs:
        if m.get("role") == "system":
            c = m.get("content")
            text = c if isinstance(c, str) else chr(10).join(b.get("text", "") for b in c if isinstance(b, dict) and b.get("type") == "text")
            if text.strip():
                extra.append(text)
        else:
            keep.append(m)
    if len(keep) != len(msgs):
        block = dict(type="text", text=(chr(10) * 2).join(extra) if extra else "(no content)")
        for m in reversed(keep):
            if m.get("role") == "user":
                c = m.get("content")
                m["content"] = ([dict(type="text", text=c)] if isinstance(c, str) else list(c)) + [block]
                break
        else:
            keep.append(dict(role="user", content=[block]))
        d["messages"] = keep
    return d


class H(http.server.BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, *a):
        pass

    def handle_any(self):
        n = int(self.headers.get("Content-Length") or 0)
        body = self.rfile.read(n) if n else b""
        if self.path.startswith("/v1/messages") and body:
            try:
                d = json.loads(body)
                with open("/content/proxy_log.jsonl", "a") as f:
                    f.write(json.dumps(dict(t=time.time(), path=self.path, req=d))[:600000] + "\n")
                body = json.dumps(fix(d)).encode()
            except Exception as e:
                open("/content/proxy_err.log", "a").write(repr(e) + "\n")
        c = http.client.HTTPConnection(*UP, timeout=7200)
        hdr = dict((k, v) for k, v in self.headers.items() if k.lower() not in ("host", "content-length", "connection", "accept-encoding", "transfer-encoding"))
        c.request(self.command, self.path, body, hdr)
        r = c.getresponse()
        with open("/content/proxy_status.log", "a") as f:
            f.write(f"{time.time():.0f} {self.command} {self.path} {r.status}\n")
        self.send_response(r.status)
        for k, v in r.getheaders():
            if k.lower() not in ("transfer-encoding", "connection", "content-length"):
                self.send_header(k, v)
        self.send_header("Transfer-Encoding", "chunked")
        self.end_headers()
        while True:
            chunk = r.read1(65536)
            if not chunk:
                break
            self.wfile.write(b"%x" % len(chunk) + CRLF + chunk + CRLF)
            self.wfile.flush()
        self.wfile.write(b"0" + CRLF + CRLF)
        self.wfile.flush()

    do_GET = do_POST = do_PUT = do_DELETE = handle_any


if __name__ == "__main__":
    http.server.ThreadingHTTPServer(("127.0.0.1", 18082), H).serve_forever()