"""MITM diagnostic: real llama_cpp.server 0.3.35 on 127.0.0.1:7861 with a tiny GGUF,
raw TCP proxy on 0.0.0.0:7860 logging both directions byte-exactly to stdout.
If external 400 reproduces through the proxy, the logs show whether the proxy
bytes differ from our local replay and what llama returns raw.
"""
import socket
import threading
import subprocess
import sys
import time
import urllib.request
import shutil


def log(*a):
    print("MITM:", *a, flush=True)


def pipe(src: socket.socket, dst: socket.socket, tag: str, addr: str):
    try:
        n = 0
        while True:
            data = src.recv(65536)
            if not data:
                break
            if n == 0:
                log(f"{tag} FIRST-PACKET from {addr} ({len(data)}B): {data[:1200]!r}")
            n += len(data)
            dst.sendall(data)
    except Exception as e:
        log(f"{tag} pipe err {e!r}")
    finally:
        try:
            dst.shutdown(socket.SHUT_WR)
        except Exception:
            pass


def handle(conn: socket.socket, addr):
    log(f"CONN open {addr}")
    try:
        upstream = socket.create_connection(("127.0.0.1", 7861), timeout=10)
        t = threading.Thread(target=pipe, args=(conn, upstream, "C->S", addr), daemon=True)
        t.start()
        pipe(upstream, conn, "S->C", addr)
    except Exception as e:
        log(f"CONN {addr} err {e!r}")
    finally:
        try:
            conn.close()
        except Exception:
            pass
        log(f"CONN close {addr}")


def download_tiny_model():
    paths = [
        ("ggml-org/models", "tinyllamas/stories260K.gguf"),
        ("ggml-org/models", "tinyllamas/stories15M-q4_0.gguf"),
        ("ggml-org/models", "tinyllamas/stories15M.gguf"),
    ]
    from huggingface_hub import hf_hub_download
    for repo, fn in paths:
        try:
            p = hf_hub_download(repo_id=repo, filename=fn, local_dir="/app/models")
            log(f"model ok: {p} ({__import__('os').path.getsize(p)} B)")
            return p
        except Exception as e:
            log(f"model fail {repo}/{fn}: {e}")
    raise RuntimeError("no tiny model")


def main():
    model = download_tiny_model()
    import llama_cpp, uvicorn, fastapi
    log(f"versions: llama_cpp={llama_cpp.__version__} uvicorn={uvicorn.__version__} fastapi={fastapi.__version__}")
    cmd = [sys.executable, "-m", "llama_cpp.server", "--model", model,
           "--host", "127.0.0.1", "--port", "7861", "--n_ctx", "512", "--n_gpu_layers", "0"]
    env = {**__import__("os").environ, "FORWARDED_ALLOW_IPS": "*"}
    # NOTE: llama server is NOT given FORWARDED_ALLOW_IPS here initially — proxy is on 7860;
    # but uvicorn inside reads env anyway: it will trust forwarded headers from 127.0.0.1 (the proxy).
    proc = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stderr)
    for i in range(120):
        try:
            urllib.request.urlopen("http://127.0.0.1:7861/v1/models", timeout=3)
            log(f"llama server ready after {i+1}s")
            break
        except Exception:
            time.sleep(1)
    else:
        log("llama server NOT ready; exiting")
        sys.exit(1)
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("0.0.0.0", 7860))
    srv.listen(64)
    log("mitm proxy listening 0.0.0.0:7860 -> 127.0.0.1:7861")
    while True:
        c, a = srv.accept()
        threading.Thread(target=handle, args=(c, a), daemon=True).start()


if __name__ == "__main__":
    main()