#!/usr/bin/env python3
"""
LMArena Proxy w/ Setup
"""

import asyncio
import json
import logging
import os
import sys
import time
import uuid
import re
import random
import http.server
import socketserver
import threading
from datetime import datetime

import uvicorn
import requests
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse, Response

# setup logging
logging.basicConfig(level=logging.INFO, format='%(message)s')
logger = logging.getLogger(__name__)

# debug flag - set to True to see detailed request/response info
DEBUG = True

def debug_log(message, data=None):
    """log debug information if debug mode is enabled"""
    if DEBUG:
        print(f"\n{'='*60}")
        print(f"[DEBUG] {message}")
        if data:
            if isinstance(data, (dict, list)):
                print(json.dumps(data, indent=2, ensure_ascii=False)[:2000])  # limit to 2000 chars
            else:
                print(str(data)[:2000])
        print(f"{'='*60}\n")

# global state
CONFIG = {}
browser_ws: WebSocket | None = None
browser_connected = False
response_channels: dict[str, asyncio.Queue] = {}
MODEL_MAP = {}
id_capture_server = None
captured_ids = {"session_id": None, "message_id": None}
id_capture_active = False

# ports
API_PORT = 5102
ID_CAPTURE_PORT = 5103

def clear_screen():
    """clear the console screen"""
    os.system('cls' if os.name == 'nt' else 'clear')

def print_banner():
    """print welcome banner"""
    print("=" * 60)
    print("  LMArena Proxy Setup")
    print("=" * 60)
    print()

def load_config():
    """load configuration from config.jsonc"""
    global CONFIG
    config_path = 'config.jsonc'

    if not os.path.exists(config_path):
        # create default config
        CONFIG = {
            "session_id": "",
            "message_id": "",
            "api_key": "",
            "stream_timeout": 360
        }
        save_config()
        return

    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            content = f.read()
            # remove comments
            lines = []
            for line in content.splitlines():
                if not line.strip().startswith('//'):
                    lines.append(line)
            CONFIG = json.loads('\n'.join(lines))
    except Exception as e:
        logger.error(f"Error loading config: {e}")
        CONFIG = {"session_id": "", "message_id": "", "api_key": "", "stream_timeout": 360}

def save_config():
    """save configuration to config.jsonc"""
    try:
        with open('config.jsonc', 'w', encoding='utf-8') as f:
            json.dump(CONFIG, f, indent=2)
    except Exception as e:
        logger.error(f"Error saving config: {e}")

def load_models():
    """load model mappings from models.json"""
    global MODEL_MAP

    if not os.path.exists('models.json'):
        MODEL_MAP = {}
        return

    try:
        with open('models.json', 'r', encoding='utf-8') as f:
            MODEL_MAP = json.load(f)
        logger.info(f"[OK] Loaded {len(MODEL_MAP)} models")
    except Exception as e:
        logger.error(f"Error loading models: {e}")
        MODEL_MAP = {}

def save_models(models_dict):
    """save model mappings to models.json"""
    try:
        with open('models.json', 'w', encoding='utf-8') as f:
            json.dump(models_dict, f, indent=2)
        logger.info(f"[OK] Saved {len(models_dict)} models to models.json")
    except Exception as e:
        logger.error(f"Error saving models: {e}")

def extract_models_from_html(html_content):
    """extract model list from lmarena html"""
    models = {}
    model_names = set()

    # find model json objects in html
    for start_match in re.finditer(r'\{\\"id\\":\\"[a-zA-Z0-9_-]+\\"', html_content):
        start_index = start_match.start()
        open_braces = 0
        end_index = -1
        search_limit = start_index + 10000

        for i in range(start_index, min(len(html_content), search_limit)):
            if html_content[i] == '{':
                open_braces += 1
            elif html_content[i] == '}':
                open_braces -= 1
                if open_braces == 0:
                    end_index = i + 1
                    break

        if end_index != -1:
            json_string_escaped = html_content[start_index:end_index]
            json_string = json_string_escaped.replace('\\"', '"').replace('\\\\', '\\')

            try:
                model_data = json.loads(json_string)
                model_name = model_data.get('publicName')
                model_id = model_data.get('id')

                if model_name and model_id and model_name not in model_names:
                    # only include text models (skip image models)
                    capabilities = model_data.get('capabilities', [])
                    if 'image' not in str(capabilities).lower() or 'text' in str(capabilities).lower():
                        models[model_name] = model_id
                        model_names.add(model_name)
            except:
                continue

    return models

# id capture server
class IDCaptureHandler(http.server.SimpleHTTPRequestHandler):
    def log_message(self, format, *args):
        pass  # suppress logs

    def do_OPTIONS(self):
        self.send_response(204)
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'POST, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', 'Content-Type')
        self.end_headers()

    def do_POST(self):
        if self.path == '/update':
            try:
                content_length = int(self.headers['Content-Length'])
                post_data = self.rfile.read(content_length)
                data = json.loads(post_data)

                session_id = data.get('sessionId')
                message_id = data.get('messageId')

                if session_id and message_id:
                    captured_ids["session_id"] = session_id
                    captured_ids["message_id"] = message_id

                    self.send_response(200)
                    self.send_header('Access-Control-Allow-Origin', '*')
                    self.end_headers()
                    self.wfile.write(b'{"status": "success"}')
                    return

                self.send_response(400)
                self.end_headers()
            except:
                self.send_response(500)
                self.end_headers()

def start_id_capture_server():
    """start the id capture server in background"""
    global id_capture_server

    def run_server():
        with socketserver.TCPServer(("127.0.0.1", ID_CAPTURE_PORT), IDCaptureHandler) as httpd:
            httpd.serve_forever()

    thread = threading.Thread(target=run_server, daemon=True)
    thread.start()

# fastapi app
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """handle tampermonkey script connection"""
    global browser_ws, browser_connected, id_capture_active
    await websocket.accept()
    browser_ws = websocket
    browser_connected = True
    logger.info("[OK] Browser connected")

    # if id capture was active before disconnect, reactivate it
    if id_capture_active:
        try:
            await websocket.send_text(json.dumps({"command": "activate_id_capture"}))
            logger.info("[INFO] Reactivated ID capture mode after reconnect")
        except:
            pass

    try:
        while True:
            message_str = await websocket.receive_text()
            message = json.loads(message_str)

            request_id = message.get("request_id")
            data = message.get("data")

            if request_id and request_id in response_channels:
                await response_channels[request_id].put(data)
    except WebSocketDisconnect:
        logger.info("[INFO] Browser disconnected (this is normal during page navigation)")
        browser_connected = False
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
        browser_connected = False
    finally:
        browser_ws = None
        # don't clear response channels immediately - give time for reconnect
        await asyncio.sleep(2)
        if not browser_connected:
            for queue in response_channels.values():
                await queue.put({"error": "Browser disconnected"})
            response_channels.clear()

@app.get("/v1/models")
async def get_models():
    """return available models"""
    if not MODEL_MAP:
        return JSONResponse(status_code=404, content={"error": "No models available"})

    return {
        "object": "list",
        "data": [
            {
                "id": model_name,
                "object": "model",
                "created": int(time.time()),
                "owned_by": "LMArena"
            }
            for model_name in MODEL_MAP.keys()
        ]
    }

@app.post("/internal/start_id_capture")
async def start_id_capture():
    """activate id capture mode"""
    global id_capture_active

    id_capture_active = True

    if not browser_ws:
        # browser not connected yet, but flag is set for when it connects
        return JSONResponse({"status": "waiting", "message": "Will activate when browser connects"})

    try:
        await browser_ws.send_text(json.dumps({"command": "activate_id_capture"}))
        return JSONResponse({"status": "success"})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/internal/update_models")
async def update_models_endpoint(request: Request):
    """receive html and extract models"""
    html_content = await request.body()
    if not html_content:
        return JSONResponse(status_code=400, content={"error": "No content"})

    models = extract_models_from_html(html_content.decode('utf-8'))
    if models:
        save_models(models)
        load_models()
        return JSONResponse({"status": "success", "count": len(models)})

    return JSONResponse(status_code=400, content={"error": "No models found"})

@app.post("/internal/request_model_update")
async def request_model_update():
    """tell browser to send page source"""
    if not browser_ws:
        raise HTTPException(status_code=503, detail="Browser not connected")

    try:
        await browser_ws.send_text(json.dumps({"command": "send_page_source"}))
        return JSONResponse({"status": "success"})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    """handle chat completion requests"""
    if not browser_ws:
        raise HTTPException(status_code=503, detail="Browser not connected. Please open LMArena in your browser.")

    try:
        req_data = await request.json()
    except:
        raise HTTPException(status_code=400, detail="Invalid JSON")

    debug_log("Incoming Request", {
        "model": req_data.get("model"),
        "messages": req_data.get("messages", []),
        "stream": req_data.get("stream", False)
    })

    model_name = req_data.get("model")
    if not model_name or model_name not in MODEL_MAP:
        raise HTTPException(status_code=400, detail=f"Model '{model_name}' not found")

    session_id = CONFIG.get("session_id")
    message_id = CONFIG.get("message_id")

    if not session_id or not message_id:
        raise HTTPException(status_code=400, detail="Session not configured. Please run setup again.")

    request_id = str(uuid.uuid4())
    response_channels[request_id] = asyncio.Queue()

    # build lmarena payload with prefill support
    messages = req_data.get("messages", [])
    message_templates = []

    # detect if last message is assistant (prefill pattern)
    has_prefill = len(messages) >= 2 and messages[-1].get("role") == "assistant"

    for i, msg in enumerate(messages):
        role = msg.get("role")
        content = msg.get("content", "")

        if role == "developer":
            role = "system"

        is_last_message = (i == len(messages) - 1)

        # For prefill: the last assistant message stays as pending with its content
        # This makes the model continue from that content
        if has_prefill and is_last_message:
            # The prefill assistant message is "pending" with content already filled
            # The model will continue generating from this point
            status = "pending"
        else:
            # normal mode: only last message is pending
            status = "pending" if is_last_message else "success"

        message_templates.append({
            "role": role,
            "content": content if content else " ",
            "attachments": [],
            "participantPosition": "b" if role == "system" else "a",
            "status": status
        })

    payload = {
        "message_templates": message_templates,
        "target_model_id": MODEL_MAP[model_name],
        "session_id": session_id,
        "message_id": message_id
    }

    debug_log("Sending to LMArena", {
        "request_id": request_id[:8] + "...",
        "model": model_name,
        "model_id": MODEL_MAP[model_name],
        "num_messages": len(message_templates),
        "has_prefill": has_prefill,
        "message_templates": message_templates
    })

    message_to_browser = {
        "request_id": request_id,
        "payload": payload
    }

    try:
        await browser_ws.send_text(json.dumps(message_to_browser))

        if req_data.get("stream", False):
            return StreamingResponse(
                stream_generator(request_id, model_name),
                media_type="text/event-stream"
            )
        else:
            return await non_stream_response(request_id, model_name)
    except Exception as e:
        if request_id in response_channels:
            del response_channels[request_id]
        raise HTTPException(status_code=500, detail=str(e))

async def stream_generator(request_id: str, model: str):
    """generate streaming response"""
    response_id = f"chatcmpl-{uuid.uuid4()}"
    queue = response_channels.get(request_id)

    if not queue:
        yield f"data: {json.dumps({'error': 'Internal error'})}\n\n"
        return

    debug_log(f"Starting stream for request {request_id[:8]}...")

    buffer = ""
    text_pattern = re.compile(r'[ab]0:"((?:\\.|[^"\\])*)"')
    timeout = CONFIG.get("stream_timeout", 360)
    received_chunks = []

    try:
        while True:
            try:
                raw_data = await asyncio.wait_for(queue.get(), timeout=timeout)
            except asyncio.TimeoutError:
                break

            if raw_data == "[DONE]":
                break

            if isinstance(raw_data, dict) and 'error' in raw_data:
                error_chunk = {
                    "id": response_id,
                    "object": "chat.completion.chunk",
                    "created": int(time.time()),
                    "model": model,
                    "choices": [{
                        "index": 0,
                        "delta": {"content": f"\n\nError: {raw_data['error']}"},
                        "finish_reason": None
                    }]
                }
                yield f"data: {json.dumps(error_chunk)}\n\n"
                break

            buffer += raw_data if isinstance(raw_data, str) else "".join(raw_data)

            while (match := text_pattern.search(buffer)):
                try:
                    text_content = json.loads(f'"{match.group(1)}"')
                    if text_content:
                        received_chunks.append(text_content)
                        chunk = {
                            "id": response_id,
                            "object": "chat.completion.chunk",
                            "created": int(time.time()),
                            "model": model,
                            "choices": [{
                                "index": 0,
                                "delta": {"content": text_content},
                                "finish_reason": None
                            }]
                        }
                        yield f"data: {json.dumps(chunk)}\n\n"
                except:
                    pass
                buffer = buffer[match.end():]

        # send finish
        finish_chunk = {
            "id": response_id,
            "object": "chat.completion.chunk",
            "created": int(time.time()),
            "model": model,
            "choices": [{
                "index": 0,
                "delta": {},
                "finish_reason": "stop"
            }]
        }
        yield f"data: {json.dumps(finish_chunk)}\n\ndata: [DONE]\n\n"

        debug_log(f"Stream complete for {request_id[:8]}", {
            "total_chunks": len(received_chunks),
            "full_response": "".join(received_chunks)[:500] + ("..." if len("".join(received_chunks)) > 500 else "")
        })

    finally:
        if request_id in response_channels:
            del response_channels[request_id]

async def non_stream_response(request_id: str, model: str):
    """generate non-streaming response"""
    response_id = f"chatcmpl-{uuid.uuid4()}"
    queue = response_channels.get(request_id)

    if not queue:
        return JSONResponse(status_code=500, content={"error": "Internal error"})

    debug_log(f"Starting non-stream for request {request_id[:8]}...")

    full_content = []
    buffer = ""
    text_pattern = re.compile(r'[ab]0:"((?:\\.|[^"\\])*)"')
    timeout = CONFIG.get("stream_timeout", 360)

    try:
        while True:
            try:
                raw_data = await asyncio.wait_for(queue.get(), timeout=timeout)
            except asyncio.TimeoutError:
                break

            if raw_data == "[DONE]":
                break

            if isinstance(raw_data, dict) and 'error' in raw_data:
                return JSONResponse(
                    status_code=500,
                    content={"error": {"message": raw_data['error']}}
                )

            buffer += raw_data if isinstance(raw_data, str) else "".join(raw_data)

            while (match := text_pattern.search(buffer)):
                try:
                    text_content = json.loads(f'"{match.group(1)}"')
                    if text_content:
                        full_content.append(text_content)
                except:
                    pass
                buffer = buffer[match.end():]
    finally:
        if request_id in response_channels:
            del response_channels[request_id]

    complete_response = "".join(full_content)

    debug_log(f"Non-stream complete for {request_id[:8]}", {
        "response_length": len(complete_response),
        "response_preview": complete_response[:500] + ("..." if len(complete_response) > 500 else "")
    })

    response_data = {
        "id": response_id,
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [{
            "index": 0,
            "message": {
                "role": "assistant",
                "content": complete_response
            },
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0
        }
    }

    return JSONResponse(content=response_data)

def run_server():
    """run the fastapi server"""
    uvicorn.run(app, host="0.0.0.0", port=API_PORT, log_level="error")

def wait_for_browser_connection(timeout=30):
    """wait for browser to connect"""
    print("[WAIT] Waiting for browser to connect...")
    start_time = time.time()

    while time.time() - start_time < timeout:
        if browser_ws is not None:
            print("[OK] Browser connected!")
            return True
        time.sleep(0.5)

    return False

def update_models_from_browser():
    """request browser to send page source and update models"""
    print("\n[INFO] Updating model list from LMArena...")

    if not browser_ws:
        print("[ERROR] Browser not connected")
        return False

    try:
        # send request to browser
        requests.post(f"http://127.0.0.1:{API_PORT}/internal/request_model_update", timeout=5)

        # wait for models to be updated
        print("[WAIT] Extracting models from page...")
        time.sleep(3)

        load_models()
        if MODEL_MAP:
            print(f"[OK] Found {len(MODEL_MAP)} models")
            return True
        else:
            print("[WARN] No models found. Make sure you're on the LMArena page.")
            return False
    except Exception as e:
        print(f"[ERROR] Error updating models: {e}")
        return False

def capture_session_ids():
    """guide user through capturing session ids"""
    global id_capture_active

    print("\n" + "=" * 60)
    print("  STEP 2: Capture Session IDs")
    print("=" * 60)
    print()
    print("Now we need to capture your session IDs.")
    print()
    print("Please follow these steps:")
    print()
    print("1. Go to: https://lmarena.ai/c/new?mode=direct")
    print("2. Send a message to ANY model")
    print("3. Wait for the model to respond")
    print("4. Click the RETRY button (circular arrows)")
    print()
    print("[WAIT] Waiting for you to click Retry...")
    print("       (ID capture will auto-activate even after reconnects)")
    print()

    # activate capture mode
    try:
        requests.post(f"http://127.0.0.1:{API_PORT}/internal/start_id_capture", timeout=5)
    except:
        print("[ERROR] Error activating capture mode")
        return False

    # wait for ids to be captured
    start_time = time.time()
    timeout = 180  # 3 minutes (more time for navigation)
    last_status_time = time.time()

    while time.time() - start_time < timeout:
        if captured_ids["session_id"] and captured_ids["message_id"]:
            print("\n[OK] Session IDs captured successfully!")
            print(f"     Session ID: {captured_ids['session_id'][:20]}...")
            print(f"     Message ID: {captured_ids['message_id'][:20]}...")

            # save to config
            CONFIG["session_id"] = captured_ids["session_id"]
            CONFIG["message_id"] = captured_ids["message_id"]
            save_config()

            # deactivate capture mode
            id_capture_active = False

            return True

        # print status update every 15 seconds
        if time.time() - last_status_time > 15:
            if browser_connected:
                print("       Still waiting... (browser connected)")
            else:
                print("       Still waiting... (reconnecting browser...)")
            last_status_time = time.time()

        time.sleep(0.5)

    print("\n[ERROR] Timeout waiting for session IDs")
    print("        Please try again and make sure to click the Retry button")
    id_capture_active = False
    return False

def interactive_setup():
    """run interactive setup"""
    clear_screen()
    print_banner()

    print("For Idiots Edition")
    print()
    print("=" * 60)
    print("  STEP 1: Connect Browser")
    print("=" * 60)
    print()
    print("Please:")
    print("1. Make sure Tampermonkey script is installed and enabled")
    print("See: https://hibbard.eu/tampermonkey-tutorial/")
    print("Copy contents of LMArenaProxy.js for the scripts")
    print("2. Open https://lmarena.ai/ in your browser")
    print()

    # wait for browser (automatic, no prompt)
    if not wait_for_browser_connection():
        print("\n[ERROR] Browser didn't connect. Please check:")
        print("        - Tampermonkey script is installed")
        print("        - Script is enabled")
        print("        - LMArena page is open")
        return False

    # update models
    time.sleep(1)
    if not update_models_from_browser():
        print("\n[WARN] Couldn't update models, but continuing...")

    # capture ids
    time.sleep(1)
    if not capture_session_ids():
        return False

    print()
    print("=" * 60)
    print("  Setup Complete!")
    print("=" * 60)
    print()
    print(f"Your proxy is now running at: http://127.0.0.1:{API_PORT}/v1")
    print()
    print("You can now use it with any OpenAI-compatible client such as SillyTavern:")
    print(f"  - API Base: http://127.0.0.1:{API_PORT}/v1")
    print(f"  (Full path http://127.0.0.1:{API_PORT}/v1/chat/completions)")
    print("  - API Key: (anything)")
    print(f"  - Models: {len(MODEL_MAP)} available")
    print(f"Error 500? You may need to do the captcha in browser.")
    print(f"Blank replies? LMArena rate limited you.")
    print()
    print("Press Ctrl+C to stop the server")
    print()

    return True

def main():
    """main entry point"""
    # load existing config
    load_config()
    load_models()

    # start id capture server
    start_id_capture_server()

    # start api server in background
    server_thread = threading.Thread(target=run_server, daemon=True)
    server_thread.start()

    # give server time to start
    time.sleep(2)

    # check if already configured
    if CONFIG.get("session_id") and CONFIG.get("message_id") and MODEL_MAP:
        clear_screen()
        print_banner()
        print("[OK] Already configured!")
        print()
        print(f"Proxy running at: http://127.0.0.1:{API_PORT}/v1")
        print(f"Models available: {len(MODEL_MAP)}")
        print()
        print("Options:")
        print("  1. Continue with existing setup")
        print("  2. Run setup again")
        print()
        choice = input("Choice (1/2): ").strip()

        if choice == "2":
            if not interactive_setup():
                print("\n[ERROR] Setup failed")
                return
        else:
            print("\nServer is running. Press Ctrl+C to stop.")
    else:
        # run interactive setup
        if not interactive_setup():
            print("\n[ERROR] Setup failed")
            return

    # keep server running
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n\nShutting down...")

if __name__ == "__main__":
    main()

// ==UserScript==
// @name         LMArena Proxy
// @namespace    http://tampermonkey.net/
// @version      3.0
// @description  Simple proxy for LMArena - text chat only
// @author       Your Name
// @match        https://lmarena.ai/*
// @match        https://*.lmarena.ai/*
// @icon         https://www.google.com/s2/favicons?sz=64&domain=lmarena.ai
// @grant        none
// @run-at       document-end
// ==/UserScript==

(function () {
    'use strict';

    const SERVER_URL = "ws://localhost:5102/ws";
    let socket;
    let isCaptureModeActive = false;

    function connect() {
        console.log(`[LMArena Proxy] Connecting to ${SERVER_URL}...`);
        socket = new WebSocket(SERVER_URL);

        socket.onopen = () => {
            console.log("[LMArena Proxy] ✅ Connected");
            document.title = "✅ " + document.title;
        };

        socket.onmessage = async (event) => {
            try {
                const message = JSON.parse(event.data);

                if (message.command) {
                    console.log(`[LMArena Proxy] Command: ${message.command}`);

                    if (message.command === 'refresh' || message.command === 'reconnect') {
                        location.reload();
                    } else if (message.command === 'activate_id_capture') {
                        console.log("[LMArena Proxy] ✅ ID capture mode activated");
                        isCaptureModeActive = true;
                        document.title = "🎯 " + document.title;
                    } else if (message.command === 'send_page_source') {
                        console.log("[LMArena Proxy] Sending page source...");
                        sendPageSource();
                    }
                    return;
                }

                const { request_id, payload } = message;

                if (!request_id || !payload) {
                    console.error("[LMArena Proxy] Invalid message:", message);
                    return;
                }

                console.log(`[LMArena Proxy] Request ${request_id.substring(0, 8)}`);
                await handleRequest(request_id, payload);

            } catch (error) {
                console.error("[LMArena Proxy] Error:", error);
            }
        };

        socket.onclose = () => {
            console.warn("[LMArena Proxy] Disconnected. Reconnecting in 5s...");
            if (document.title.startsWith("✅ ")) {
                document.title = document.title.substring(2);
            }
            setTimeout(connect, 5000);
        };

        socket.onerror = (error) => {
            console.error("[LMArena Proxy] Error:", error);
            socket.close();
        };
    }

    async function handleRequest(requestId, payload) {
        const { message_templates, target_model_id, session_id, message_id } = payload;

        if (!session_id || !message_id) {
            const errorMsg = "Session IDs missing. Please run setup again.";
            console.error(`[LMArena Proxy] ${errorMsg}`);
            sendToServer(requestId, { error: errorMsg });
            sendToServer(requestId, "[DONE]");
            return;
        }

        if (!message_templates || message_templates.length === 0) {
            const errorMsg = "No messages to send.";
            console.error(`[LMArena Proxy] ${errorMsg}`);
            sendToServer(requestId, { error: errorMsg });
            sendToServer(requestId, "[DONE]");
            return;
        }

        const apiUrl = `/nextjs-api/stream/retry-evaluation-session-message/${session_id}/messages/${message_id}`;

        const newMessages = [];
        let lastMsgId = null;

        for (let i = 0; i < message_templates.length; i++) {
            const template = message_templates[i];
            const currentMsgId = crypto.randomUUID();
            const parentIds = lastMsgId ? [lastMsgId] : [];

            // use status from template if provided, otherwise default behavior
            const status = template.status || ((i === message_templates.length - 1) ? 'pending' : 'success');

            // pending assistant messages need modelId to trigger generation
            const modelId = (status === 'pending' && template.role === 'assistant') ? target_model_id : null;

            newMessages.push({
                role: template.role,
                content: template.content,
                id: currentMsgId,
                evaluationId: null,
                evaluationSessionId: session_id,
                parentMessageIds: parentIds,
                experimental_attachments: [],
                failureReason: null,
                metadata: null,
                modelId: modelId,
                participantPosition: template.participantPosition || "a",
                createdAt: new Date().toISOString(),
                updatedAt: new Date().toISOString(),
                status: status,
            });
            lastMsgId = currentMsgId;
        }

        const body = {
            messages: newMessages,
            modelId: target_model_id,
        };

        console.log("[LMArena Proxy] Sending to LMArena API");

        window.isProxyRequest = true;
        try {
            const response = await fetch(apiUrl, {
                method: 'PUT',
                headers: {
                    'Content-Type': 'text/plain;charset=UTF-8',
                    'Accept': '*/*',
                },
                body: JSON.stringify(body),
                credentials: 'include'
            });

            if (!response.ok || !response.body) {
                const errorBody = await response.text();
                throw new Error(`Response error: ${response.status}. ${errorBody}`);
            }

            const reader = response.body.getReader();
            const decoder = new TextDecoder();

            while (true) {
                const { value, done } = await reader.read();
                if (done) {
                    console.log(`[LMArena Proxy] ✅ Request ${requestId.substring(0, 8)} complete`);
                    sendToServer(requestId, "[DONE]");
                    break;
                }
                const chunk = decoder.decode(value);
                sendToServer(requestId, chunk);
            }

        } catch (error) {
            console.error(`[LMArena Proxy] ❌ Error:`, error);
            sendToServer(requestId, { error: error.message });
        } finally {
            window.isProxyRequest = false;
        }
    }

    function sendToServer(requestId, data) {
        if (socket && socket.readyState === WebSocket.OPEN) {
            const message = {
                request_id: requestId,
                data: data
            };
            socket.send(JSON.stringify(message));
        } else {
            console.error("[LMArena Proxy] Cannot send data - not connected");
        }
    }

    // intercept fetch to capture session ids
    const originalFetch = window.fetch;
    window.fetch = function(...args) {
        const urlArg = args[0];
        let urlString = '';

        if (urlArg instanceof Request) {
            urlString = urlArg.url;
        } else if (urlArg instanceof URL) {
            urlString = urlArg.href;
        } else if (typeof urlArg === 'string') {
            urlString = urlArg;
        }

        if (urlString) {
            const match = urlString.match(/\/nextjs-api\/stream\/retry-evaluation-session-message\/([a-f0-9-]+)\/messages\/([a-f0-9-]+)/);

            if (match && !window.isProxyRequest && isCaptureModeActive) {
                const sessionId = match[1];
                const messageId = match[2];
                console.log(`[LMArena Proxy] 🎯 Captured IDs!`);

                isCaptureModeActive = false;
                if (document.title.startsWith("🎯 ")) {
                    document.title = document.title.substring(2);
                }

                fetch('http://127.0.0.1:5103/update', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ sessionId, messageId })
                })
                .then(response => {
                    if (!response.ok) throw new Error(`Status: ${response.status}`);
                    console.log(`[LMArena Proxy] ✅ IDs sent to server`);
                })
                .catch(err => {
                    console.error('[LMArena Proxy] Error sending IDs:', err.message);
                });
            }
        }

        return originalFetch.apply(this, args);
    };

    async function sendPageSource() {
        try {
            const htmlContent = document.documentElement.outerHTML;
            await fetch('http://localhost:5102/internal/update_models', {
                method: 'POST',
                headers: {
                    'Content-Type': 'text/html; charset=utf-8'
                },
                body: htmlContent
            });
            console.log("[LMArena Proxy] Page source sent");
        } catch (e) {
            console.error("[LMArena Proxy] Error sending page source:", e);
        }
    }

    console.log("========================================");
    console.log("  LMArena Proxy v3.0");
    console.log("  Connected to: ws://localhost:5102");
    console.log("========================================");

    connect();

})();