#!/usr/bin/env python3
"""
daemon.py — ChatGPT persistent browser daemon
----------------------------------------------
Launches camoufox + loads session ONCE, then exposes a JSON HTTP API on
localhost:9223 so any script can send prompts without re-launching the browser.

Usage:
    # Start (run in background):
    PYTHONUNBUFFERED=1 python3 /home/node/.claude/skills/chatgpt/scripts/daemon.py &

    # Check it's alive:
    curl http://localhost:9223/health

    # Send a prompt:
    curl -s -X POST http://localhost:9223/send \
         -H 'Content-Type: application/json' \
         -d '{"prompt": "what is 2+2?"}' | python3 -m json.tool

    # Stop:
    curl -s -X POST http://localhost:9223/stop

Architecture:
    HTTP server thread  →  _req_queue  →  Playwright main thread
                        ←  _resp_queue ←

The Playwright page object must only be used from the main thread; the HTTP
server runs in a background thread and communicates via thread-safe queues.
"""

import json
import os
import queue
import signal
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path

sys.path.insert(0, "/home/node/.claude/skills/chatgpt/scripts")
from load_session import open_chatgpt_session

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
PORT      = 9223
PID_FILE  = Path("/workspace/group/.chatgpt-daemon.pid")
IMAGE_DIR = Path("/workspace/group/chatgpt-images")

# ---------------------------------------------------------------------------
# Thread-safe IPC between HTTP thread and Playwright main thread
# ---------------------------------------------------------------------------
_req_queue  = queue.Queue(maxsize=1)
_resp_queue = queue.Queue(maxsize=1)
_shutdown   = threading.Event()


# ---------------------------------------------------------------------------
# HTTP handler
# ---------------------------------------------------------------------------
class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        pass  # suppress default access logs

    def _send_json(self, code: int, data: dict):
        body = json.dumps(data).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        if self.path == "/health":
            try:
                _req_queue.put({"action": "health"}, timeout=1)
                result = _resp_queue.get(timeout=10)
                self._send_json(200, result)
            except queue.Full:
                self._send_json(503, {"error": "daemon busy"})
            except queue.Empty:
                self._send_json(504, {"error": "timeout"})
        else:
            self._send_json(404, {"error": "not found"})

    def do_POST(self):
        length = int(self.headers.get("Content-Length", 0))
        body   = self.rfile.read(length) if length else b"{}"
        try:
            data = json.loads(body)
        except Exception:
            self._send_json(400, {"error": "invalid JSON"})
            return

        if self.path == "/send":
            try:
                _req_queue.put({"action": "send", "prompt": data.get("prompt", "")}, timeout=5)
            except queue.Full:
                self._send_json(503, {"error": "daemon busy, try again"})
                return
            try:
                result = _resp_queue.get(timeout=300)  # 5-min max for image gen
                self._send_json(200, result)
            except queue.Empty:
                self._send_json(504, {"error": "timeout waiting for browser"})

        elif self.path == "/stop":
            _shutdown.set()
            self._send_json(200, {"status": "stopping"})

        else:
            self._send_json(404, {"error": "not found"})


def _run_http_server(port: int):
    server = HTTPServer(("localhost", port), Handler)
    print(f"[daemon] HTTP API listening on http://localhost:{port}")
    while not _shutdown.is_set():
        server.handle_request()
    server.server_close()


# ---------------------------------------------------------------------------
# Completion detection (Playwright main thread)
# ---------------------------------------------------------------------------
def _wait_for_completion(page, timeout_s: int = 180) -> str:
    """
    Wait for ChatGPT to finish responding.
    Returns 'image' | 'text' | 'timeout'.

    Strategy: poll the last assistant message's text. When it exists and
    remains stable for two consecutive 0.8s checks, streaming is done.
    Also watches for image download buttons.
    """
    time.sleep(0.5)
    print(f"[daemon] waiting for completion, url={page.url}")
    deadline  = time.time() + timeout_s
    iteration = 0
    prev_text = None
    prev_text_iter = -99

    # Phase 1: wait for URL to move to /c/ (up to 10s)
    url_deadline = time.time() + 10
    while time.time() < url_deadline:
        if "/c/" in page.url:
            break
        time.sleep(0.15)
    print(f"[daemon] conversation url={page.url}")

    while time.time() < deadline:
        # --- Image generation complete ---
        if page.query_selector('button[aria-label*="Download"]'):
            print(f"[daemon] detected: image download button (iter {iteration})")
            # Wait for all images to fully load (not just partially rendered)
            _wait_for_images_loaded(page)
            return "image"

        # --- Text: last assistant message stabilised ---
        try:
            msgs = page.query_selector_all('[data-message-author-role="assistant"]')
            cur_text = msgs[-1].inner_text().strip() if msgs else ""
        except Exception:
            cur_text = ""

        if iteration % 20 == 0:
            snippet = cur_text[:60].replace('\n', ' ') if cur_text else "(empty)"
            print(f"[daemon] iter={iteration} url={page.url} msg={snippet!r}")

        if cur_text and cur_text == prev_text and (iteration - prev_text_iter) >= 1:
            # Text hasn't changed since last check — streaming is done
            print(f"[daemon] detected: message stable (iter {iteration})")
            return "text"

        if cur_text != prev_text:
            prev_text = cur_text
            prev_text_iter = iteration

        time.sleep(0.5)
        iteration += 1

    print(f"[daemon] timed out after {timeout_s}s")
    return "timeout"


# ---------------------------------------------------------------------------
# Image load detection (Playwright main thread)
# ---------------------------------------------------------------------------
def _wait_for_images_loaded(page, timeout_s: int = 15):
    """Wait until all generated images in the assistant message are fully loaded.

    Checks img.complete and img.naturalWidth — a partially loaded progressive
    image will have complete=false or a small naturalWidth until fully decoded.
    """
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        all_loaded = page.evaluate('''() => {
            const imgs = document.querySelectorAll('img[alt*="Generated"]');
            if (imgs.length === 0) return false;
            for (const img of imgs) {
                if (!img.complete || img.naturalWidth < 512) return false;
            }
            return true;
        }''')
        if all_loaded:
            print(f"[daemon] all images fully loaded")
            return
        time.sleep(0.5)
    print(f"[daemon] image load wait timed out after {timeout_s}s")


# ---------------------------------------------------------------------------
# Image extraction (Playwright main thread)
# ---------------------------------------------------------------------------
def _extract_images(page, ctx) -> list:
    """Download generated images, return saved file paths.

    Filters out preview/thumbnail images by:
    1. Preferring images near a download button (the final render)
    2. Downloading all candidates, then keeping only the largest file
       per unique resolution (discards blurry previews)
    """
    IMAGE_DIR.mkdir(parents=True, exist_ok=True)

    # Extract only the final rendered images.
    # ChatGPT creates two file IDs per generation: the final render has
    # alt="Generated image: ..." while the blurry preview has alt="".
    # Prefer images with the "Generated" alt tag; fall back to all if none match.
    urls = page.evaluate('''() => {
        // First pass: only images with alt containing "Generated" (final renders)
        let imgs = document.querySelectorAll('img[alt*="Generated"]');
        const result = [];
        const seen = new Set();
        for (const img of imgs) {
            if (img.complete && img.naturalWidth >= 512 && img.src && !seen.has(img.src)) {
                seen.add(img.src);
                result.push(img.src);
            }
        }
        if (result.length > 0) return result;

        // Fallback: all estuary/oaiusercontent images
        imgs = document.querySelectorAll(
            'img[src*="oaiusercontent"], img[src*="estuary"]'
        );
        for (const img of imgs) {
            if (img.complete && img.naturalWidth >= 512 && img.src && !seen.has(img.src)) {
                seen.add(img.src);
                result.push(img.src);
            }
        }
        return result;
    }''')

    if not urls:
        return []

    print(f"[daemon] found {len(urls)} candidate image URLs")

    import requests
    cookies = {c["name"]: c["value"] for c in ctx.cookies()}
    sess = requests.Session()
    sess.cookies.update(cookies)
    sess.headers["User-Agent"] = "Mozilla/5.0"

    # Download all candidates
    downloaded = []
    for i, url in enumerate(urls):
        try:
            resp = sess.get(url, timeout=60)
            resp.raise_for_status()
            ext = "webp" if "webp" in resp.headers.get("content-type", "") else "png"
            ts  = int(time.time())
            out = IMAGE_DIR / f"image-{ts}-{i+1}.{ext}"
            out.write_bytes(resp.content)
            downloaded.append((str(out), len(resp.content)))
            print(f"[daemon] downloaded {out.name}: {len(resp.content):,} bytes")
        except Exception as e:
            print(f"[daemon] image download failed: {e}")

    if not downloaded:
        return []

    return [p for p, _ in downloaded]


# ---------------------------------------------------------------------------
# Prompt handling (Playwright main thread)
# ---------------------------------------------------------------------------
def _do_send(page, ctx, prompt: str) -> dict:
    """Send a prompt and wait for the response."""
    print(f"[daemon] _do_send start, url={page.url}")
    page.set_viewport_size({"width": 1280, "height": 900})

    # Navigate to fresh new chat
    if "/c/" in page.url or page.url.rstrip("/") != "https://chatgpt.com":
        print("[daemon] navigating to chatgpt.com...")
        page.goto("https://chatgpt.com/", wait_until="domcontentloaded")
        time.sleep(1)
    print(f"[daemon] on url={page.url}")

    # Focus and type into the composer
    # Note: fill() doesn't trigger ChatGPT's input handlers reliably, so we
    # use type() with a fast delay. 15ms is ~3x faster than the original 40ms.
    print("[daemon] clicking composer...")
    composer = page.locator(
        'div[contenteditable="true"]#prompt-textarea, '
        '#prompt-textarea, '
        'div[contenteditable="true"][data-id="root"]'
    ).first
    composer.click()
    time.sleep(0.3)
    print("[daemon] typing prompt...")
    composer.type(prompt, delay=15)
    time.sleep(0.1)

    # Find and click send button
    sb = page.query_selector('button[data-testid="send-button"]')
    print(f"[daemon] send-button found={sb is not None}, enabled={sb.is_enabled() if sb else 'n/a'}")
    if not sb:
        return {"error": "send button not found", "url": page.url}

    sb.click()
    print("[daemon] prompt sent, waiting for response...")

    kind = _wait_for_completion(page)
    print(f"[daemon] response complete, kind={kind}")

    if kind == "image":
        paths = _extract_images(page, ctx)
        return {
            "kind": "image",
            "images": paths,
            "conv_url": page.url,
        }
    elif kind == "text":
        # Read last assistant message
        text = ""
        try:
            msgs = page.query_selector_all('[data-message-author-role="assistant"]')
            if msgs:
                text = msgs[-1].inner_text().strip()
        except Exception as e:
            print(f"[daemon] text extraction error: {e}")
        return {
            "kind": "text",
            "text": text,
            "conv_url": page.url,
        }
    else:
        return {"kind": "timeout", "conv_url": page.url}


def _do_health(page) -> dict:
    return {"status": "ok", "url": page.url}


# ---------------------------------------------------------------------------
# Main loop (Playwright main thread)
# ---------------------------------------------------------------------------
def main():
    PID_FILE.write_text(str(os.getpid()))
    print(f"[daemon] PID {os.getpid()} written to {PID_FILE}")

    # Start HTTP server in background thread
    http_thread = threading.Thread(target=_run_http_server, args=(PORT,), daemon=True)
    http_thread.start()

    # Graceful shutdown on SIGTERM
    signal.signal(signal.SIGTERM, lambda *_: _shutdown.set())

    print(f"[daemon] loading session from /workspace/group/chatgpt-session.json...")
    with open_chatgpt_session() as (ctx, page):
        print(f"[daemon] ready at {page.url}")

        while not _shutdown.is_set():
            try:
                cmd = _req_queue.get(timeout=1)
            except queue.Empty:
                continue

            action = cmd.get("action")
            try:
                if action == "health":
                    result = _do_health(page)
                elif action == "send":
                    result = _do_send(page, ctx, cmd["prompt"])
                else:
                    result = {"error": f"unknown action: {action}"}
            except Exception as e:
                print(f"[daemon] error in action={action}: {e}")
                result = {"error": str(e)}

            _resp_queue.put(result)

    print("[daemon] shutdown complete")
    PID_FILE.unlink(missing_ok=True)


if __name__ == "__main__":
    main()
