"""Send one prompt to the owner's ChatGPT session and return what comes back.

Drives the chat web app rather than the API: the owner's subscription already pays
for these turns, and image generation is a first-class part of that surface. Nothing
here touches metered API billing.
"""

from __future__ import annotations

import base64
import binascii
import filecmp
import os
import re
import shutil
import time
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
from urllib.parse import unquote, urlparse

from browser import DOWNLOAD_DIR, Marionette, MarionetteError, Session

STATE = Path(os.environ.get("GPTBRIDGE_STATE_DIR", Path.home() / ".overdeck" / "gptbridge"))
OUT_DIR = STATE / "out"

COMPOSER = "#prompt-textarea"
SEND = 'button[data-testid="send-button"]'
STOP = 'button[data-testid="stop-button"]'
UPLOAD = "input#upload-files"

# The composer's settings button carries no test id; it is labelled with the model on
# one line and the effort in force on the next, so it is matched by that trailing word.
# What the effort menu actually offers, keyed by the argument the CLI accepts.
EFFORT_LABELS = {
    "instant": "instant",
    "medium": "medium",
    "high": "high",
    "xhigh": "extra high",
    "pro": "pro",
}
EFFORTS = tuple(EFFORT_LABELS)
# The composer's settings button is labelled with the effort in force, sometimes below
# the model name. Both surfaces use vocabulary from this set, so it locates the button
# without depending on a test id the app does not provide.
SETTINGS_WORDS = ("instant", "light", "medium", "high", "extra high", "max", "ultra",
                  "pro", "auto", "thinking")
SETTINGS_BUTTON = r"(" + "|".join(SETTINGS_WORDS) + r")\s*$"


class ChatError(RuntimeError):
    pass


@dataclass
class Reply:
    text: str
    images: list[Path] = field(default_factory=list)
    files: list[Path] = field(default_factory=list)
    url: str = ""


MARK = """
document.querySelectorAll('[data-gptbridge]').forEach(
  e => e.removeAttribute('data-gptbridge'));
const want = arguments[0].toLowerCase();
const scope = arguments[1] || 'button,[role=menuitem],[role=option],[role=menuitemradio]';
let best = null;
for (const n of document.querySelectorAll(scope)) {
  const t = (n.innerText || '').trim().toLowerCase();
  if (!t || n.getBoundingClientRect().width < 1) continue;
  if (t === want || t.startsWith(want) || t.includes(want)) {
    if (!best || t.length < best.t.length) best = {n: n, t: t};
  }
}
if (!best) return null;
best.n.setAttribute('data-gptbridge', '1');
return best.t.slice(0, 80);
"""

MARK_RE = """
document.querySelectorAll('[data-gptbridge]').forEach(
  e => e.removeAttribute('data-gptbridge'));
const re = new RegExp(arguments[0], 'i');
for (const n of document.querySelectorAll(arguments[1])) {
  const t = (n.innerText || '').trim();
  if (t && re.test(t) && n.getBoundingClientRect().width > 1) {
    n.setAttribute('data-gptbridge', '1');
    return t.slice(0, 80);
  }
}
return null;
"""


MENU_OPTIONS = """
return Array.from(document.querySelectorAll(
  '[role=menuitem],[role=option],[role=menuitemradio]'))
  .map(e => (e.innerText || '').trim().split('\\n')[0])
  .filter(t => t);
"""


def _click_marked(m: Marionette) -> bool:
    return m.js_click("[data-gptbridge]")


def click_label(m: Marionette, label: str, scope: str | None = None) -> str | None:
    found = m.sync_script(MARK, [label, scope])
    if not found or not _click_marked(m):
        return None
    return str(found)


def wait_ready(m: Marionette, timeout: float = 120.0) -> None:
    """Wait for the composer, sitting out any bot-check interstitial."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        time.sleep(2)
        try:
            state = m.sync_script(
                "return {t: document.title, c: !!document.querySelector(arguments[0])};",
                [COMPOSER],
            )
        except MarionetteError as exc:
            if "unload" not in str(exc) and "detach" not in str(exc):
                raise
            continue
        if state.get("c"):
            return
    raise ChatError("the chat composer never appeared")


# The account also has a Work surface with its own, much smaller quota. Turns must go
# to the Chat surface, and the app remembers whichever was used last.
SURFACE = """
const hdr = document.querySelector('header') || document.body;
const tabs = Array.from(hdr.querySelectorAll('button,[role=tab]'))
  .filter(b => ['chat', 'work'].includes((b.innerText || '').trim().toLowerCase()));
const active = tabs.find(b => b.getAttribute('data-state') === 'on');
const chat = tabs.find(b => (b.innerText || '').trim().toLowerCase() === 'chat');
if (!chat) return {found: false};
if (active === chat) return {found: true, switched: false};
chat.setAttribute('data-gptbridge', '1');
return {found: true, switched: true};
"""


def use_chat_surface(m: Marionette) -> bool:
    """Put the conversation on the Chat surface rather than Work."""
    state = m.sync_script(SURFACE)
    if not state.get("found"):
        return False
    if not state.get("switched"):
        return True
    _click_marked(m)
    time.sleep(3)
    wait_ready(m)
    return True


def set_effort(m: Marionette, effort: str) -> str:
    """Pick a reasoning effort from the composer's settings menu."""
    if effort not in EFFORTS:
        raise ChatError(f"effort must be one of {EFFORTS}")
    if not m.sync_script(MARK_RE, [SETTINGS_BUTTON, "form button"]):
        raise ChatError("no effort control on the composer")
    if not _click_marked(m):
        raise ChatError("the composer settings menu would not open")
    time.sleep(2)
    if not click_label(m, "effort"):
        raise ChatError("the settings menu has no effort entry")
    time.sleep(2)
    wanted = EFFORT_LABELS[effort]
    chosen = click_label(m, wanted)
    if not chosen:
        offered = m.sync_script(MENU_OPTIONS)
        raise ChatError(
            f"effort {effort!r} is not offered here; this surface offers: "
            + ", ".join(offered)
        )
    time.sleep(1.5)
    # The menu stays open over the composer and would swallow the send click.
    m.press()
    time.sleep(1)
    return chosen


# A capped account states it in a notice beside the composer or in a dialog over it.
# The owner's account was not capped when this was written, so there is no class or
# test id to pin: the scan is anchored on the composer's own region plus the cap
# vocabulary, which returns None on anything it does not recognise rather than
# inventing a cap.
CAP_WORDS = r"(hit|reached|reach)ed?\s+(your|the)\s+.{0,24}limit|" \
            r"you've\s+(hit|reached)|usage\s+limit|message\s+limit|" \
            r"limit\s+(reached|resets)|out of\s+(messages|credits)"
CAP_SCAN = r"""
const words = new RegExp(arguments[1], 'i');
const form = document.querySelector('form');
const roots = [form && form.parentElement, form,
               document.querySelector('[role=dialog]')].filter(Boolean);
for (const root of roots) {
  for (const n of root.querySelectorAll('div,span,p')) {
    const t = (n.innerText || '').trim();
    if (!t || t.length > 400 || n.children.length > 3) continue;
    if (words.test(t)) return {text: t};
  }
}
return null;
"""
# "resets at 3:45 PM", "try again at 15:45" — the wording varies, the clock time does not.
RESET_TIME = re.compile(r"\b(\d{1,2}):(\d{2})\s*(am|pm)?\b", re.I)


def parse_reset(text: str, *, now: datetime | None = None) -> str | None:
    """The next occurrence of the clock time a cap notice names, as ISO 8601."""
    match = RESET_TIME.search(text)
    if not match:
        return None
    hour, minute = int(match.group(1)), int(match.group(2))
    meridiem = (match.group(3) or "").lower()
    if meridiem == "pm" and hour < 12:
        hour += 12
    elif meridiem == "am" and hour == 12:
        hour = 0
    if not (0 <= hour <= 23 and 0 <= minute <= 59):
        return None
    now = now or datetime.now().astimezone()
    resume = now.replace(hour=hour, minute=minute, second=0, microsecond=0)
    if resume <= now:
        resume += timedelta(days=1)
    return resume.isoformat()


def detect_cap(m: Marionette) -> dict | None:
    """The usage-cap state if the session is capped, else None. Reads only."""
    hit = m.sync_script(CAP_SCAN, [COMPOSER, CAP_WORDS])
    if not hit:
        return None
    text = str(hit.get("text", ""))
    return {"kind": "usage-cap", "resume_at": parse_reset(text), "text": text}


def attach(m: Marionette, paths: list[Path]) -> None:
    for path in paths:
        resolved = path.expanduser().resolve()
        if not resolved.is_file():
            raise ChatError(f"attachment not found: {resolved}")
        if not m.send_keys(UPLOAD, str(resolved)):
            raise ChatError("the chat page exposes no file input")
    if paths:
        wait_uploads(m)


def wait_uploads(m: Marionette, timeout: float = 180.0) -> None:
    """Uploads must finish before send, or the turn goes without them."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        time.sleep(2)
        busy = m.sync_script(
            "return document.querySelectorAll("
            "'[data-testid*=attachment] [role=progressbar],"
            " form [role=progressbar]').length;"
        )
        if not busy:
            return
    raise ChatError("attachments never finished uploading")


TYPE = """
const ed = document.querySelector(arguments[0]);
if (!ed) return false;
ed.focus();
document.execCommand('selectAll', false, null);
document.execCommand('insertText', false, arguments[1]);
return (ed.innerText || '').length > 0;
"""


def type_prompt(m: Marionette, prompt: str) -> None:
    # Keystroke-by-keystroke entry is far too slow for a multi-line prompt, and the
    # editor ignores a plain textContent assignment.
    if not m.sync_script(TYPE, [COMPOSER, prompt]):
        raise ChatError("could not enter the prompt into the composer")
    time.sleep(1)


# Text replies carry data-message-author-role="assistant", but image replies do not —
# they render as a bare conversation turn. Count turns and treat any turn without a
# user message inside it as the model's.
TURN = '[data-testid^="conversation-turn"]'
USER = '[data-message-author-role="user"]'
TURN_COUNT = """
let n = 0;
for (const t of document.querySelectorAll(arguments[0])) {
  if (!t.querySelector(arguments[1])) n++;
}
return n;
"""


PROGRESS = r"""
let n = 0, imgs = 0, text = 0;
for (const t of document.querySelectorAll(arguments[0])) {
  if (t.querySelector(arguments[1])) continue;
  n++;
  imgs = t.querySelectorAll('img').length;
  text = (t.innerText || '').replace(/^ChatGPT said:\s*/, '').trim().length;
}
return {turns: n, imgs: imgs, text: text,
        stop: !!document.querySelector(arguments[2])};
"""


USER_COUNT = "return document.querySelectorAll(arguments[0]).length;"


def submit(m: Marionette, timeout: float = 120.0) -> None:
    """Click send until the prompt is actually accepted.

    A click lands on an enabled send button and is still dropped while a large
    attachment is being processed, and the page looks identical either way — so the
    user turn appearing is the only proof, and it is clicked again until it does.
    """
    before = int(m.sync_script(USER_COUNT, [USER]) or 0)
    deadline = time.monotonic() + timeout
    clicked = False
    while True:
        if m.click(SEND) or m.js_click(SEND):
            clicked = True
        elif not clicked:
            raise ChatError("the send button never became available")
        for _ in range(6):
            time.sleep(1)
            if int(m.sync_script(USER_COUNT, [USER]) or 0) > before:
                return
        if time.monotonic() >= deadline:
            break
    raise ChatError("the composer never accepted the prompt"
                    f" after {timeout:.0f}s of clicking send{_timeout_shot(m)}")


def send(m: Marionette, before: int, timeout: float = 900.0) -> None:
    # One budget covers submission and the reply: how long the app takes to accept a
    # large attachment is the caller's wait, not a constant unrelated to it.
    deadline = time.monotonic() + timeout
    submit(m, timeout=deadline - time.monotonic())
    # The stop button flickers away between streamed chunks, so one quiet poll is not
    # proof the turn ended.
    quiet = 0
    state: dict = {}
    while time.monotonic() < deadline:
        time.sleep(2.5)
        state = m.sync_script(PROGRESS, [TURN, USER, STOP])
        # An image turn stays textless, so accept either kind of content.
        done = (not state.get("stop")
                and state.get("turns", 0) > before
                and (state.get("imgs", 0) or state.get("text", 0)))
        quiet = quiet + 1 if done else 0
        if quiet >= 2:
            return
    raise ChatError(f"the reply did not finish in time; {describe(state, before)}"
                    f"{_timeout_shot(m)}")


def describe(state: dict, before: int) -> str:
    """What the page looked like on the last poll — a bare timeout names no cause."""
    if not state:
        return "the page was never polled"
    return (f"reply turns {state.get('turns', 0)} (was {before} before sending), "
            f"{state.get('text', 0)} chars, {state.get('imgs', 0)} images, stop button "
            f"{'still up' if state.get('stop') else 'gone'}")


def _timeout_shot(m: Marionette) -> str:
    try:
        return f"; page saved to {m.screenshot(STATE / 'timeout.png')}"
    except (MarionetteError, OSError):
        return ""


HARVEST = r"""
const turns = Array.from(document.querySelectorAll(arguments[0]))
  .filter(t => !t.querySelector(arguments[1]));
const last = turns[turns.length - 1];
if (!last) return null;
const imgs = [];
for (const img of last.querySelectorAll('img')) {
  const src = img.currentSrc || img.src || '';
  const r = img.getBoundingClientRect();
  if (src && !src.startsWith('data:') && r.width >= 96 && r.height >= 96) {
    if (!imgs.includes(src)) imgs.push(src);
  }
}
const files = [];
for (const link of last.querySelectorAll('a[href]')) {
  const raw = link.getAttribute('href') || '';
  const href = link.href || raw;
  const downloadable = link.hasAttribute('download') || raw.startsWith('sandbox:') ||
    raw.startsWith('blob:') || href.includes('/backend-api/files/') ||
    href.includes('files.oaiusercontent.com') || href.includes('/mnt/data/');
  if (!downloadable || files.some(item => item.src === href)) continue;
  files.push({
    src: href,
    name: link.getAttribute('download') || link.getAttribute('title') ||
      link.getAttribute('aria-label') || (link.innerText || '').trim(),
  });
}
const downloads = Array.from(last.querySelectorAll('button'))
  .map(button => (button.innerText || '').trim())
  .filter(text => /^download\s+.+/i.test(text));
const text = (last.innerText || '').replace(/^ChatGPT said:\s*/, '').trim();
return {text: text, images: imgs, files: files, downloads: downloads, url: location.href};
"""

CLICK_DOWNLOAD = r"""
const turns = Array.from(document.querySelectorAll(arguments[0]))
  .filter(t => !t.querySelector(arguments[1]));
const last = turns[turns.length - 1];
if (!last) return false;
const buttons = Array.from(last.querySelectorAll('button'))
  .filter(button => /^download\s+.+/i.test((button.innerText || '').trim()));
const button = buttons[arguments[2]];
if (!button) return false;
button.click();
return true;
"""


# Fetching from the page keeps the session cookies on the request; the asset hosts
# reject an unauthenticated download.
FETCH = """
const done = arguments[arguments.length - 1];
fetch(arguments[0], {credentials: 'include'})
  .then(r => {
    if (!r.ok) throw new Error('http ' + r.status);
    const type = r.headers.get('content-type') || '';
    const disposition = r.headers.get('content-disposition') || '';
    return r.blob().then(blob => ({blob, type, disposition}));
  })
  .then(item => new Promise(res => {
    const fr = new FileReader();
    fr.onload = () => res({
      type: item.type || item.blob.type,
      disposition: item.disposition,
      data: fr.result,
    });
    fr.readAsDataURL(item.blob);
  }))
  .then(done)
  .catch(e => done({error: String(e)}));
"""

EXT = {
    "image/png": ".png",
    "image/jpeg": ".jpg",
    "image/webp": ".webp",
    "image/gif": ".gif",
}


def slugify(text: str, limit: int = 40) -> str:
    slug = re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")
    return (slug[:limit].rstrip("-")) or "reply"


def _content_disposition_name(value: str) -> str:
    encoded = re.search(r"filename\*=UTF-8''([^;]+)", value, re.IGNORECASE)
    if encoded:
        return unquote(encoded.group(1))
    plain = re.search(r'filename="?([^";]+)', value, re.IGNORECASE)
    return plain.group(1).strip() if plain else ""


def _safe_file_name(candidate: str, src: str, disposition: str, index: int) -> str:
    disposition_name = _content_disposition_name(disposition)
    path_name = unquote(urlparse(src).path.rsplit("/", 1)[-1])
    candidates = ([candidate] if Path(candidate).suffix else []) + [
        disposition_name, path_name, candidate]
    for value in candidates:
        value = value.strip().splitlines()[0] if value.strip() else ""
        if not value or value in {".", ".."} or "/" in value or "\\" in value:
            continue
        if any(ord(char) < 32 for char in value):
            continue
        cleaned = re.sub(r'[<>:"|?*]', "_", value).strip(" .")
        if cleaned:
            return cleaned
    return f"generated-{index}"


def _fetch_blob(m: Marionette, src: str, kind: str, index: int) -> dict:
    blob = m.script(FETCH, [src])
    if not isinstance(blob, dict) or "data" not in blob:
        raise ChatError(f"could not download {kind} {index}: {blob}")
    return blob


def _blob_bytes(blob: dict) -> bytes:
    _, separator, payload = str(blob["data"]).partition(",")
    if not separator:
        raise ChatError("download did not return a data URL")
    try:
        return base64.b64decode(payload, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise ChatError("download returned invalid base64") from exc


def _append_unique_file(reply: Reply, path: Path) -> None:
    for existing in reply.files:
        if existing.stat().st_size == path.stat().st_size and filecmp.cmp(
            existing, path, shallow=False
        ):
            path.unlink()
            return
    reply.files.append(path)


def _download_snapshot() -> dict[Path, tuple[int, int]]:
    DOWNLOAD_DIR.mkdir(mode=0o700, parents=True, exist_ok=True)
    return {
        path: (path.stat().st_mtime_ns, path.stat().st_size)
        for path in DOWNLOAD_DIR.iterdir()
        if path.is_file()
    }


def _download_button(
    m: Marionette, index: int, label: str, out_dir: Path, timeout: float = 300.0,
) -> Path:
    before = _download_snapshot()
    clicked = m.sync_script(CLICK_DOWNLOAD, [TURN, USER, index])
    if not clicked:
        raise ChatError(f"could not click generated file {index + 1}: {label}")
    deadline = time.monotonic() + timeout
    stable: tuple[Path, tuple[int, int]] | None = None
    stable_samples = 0
    while time.monotonic() < deadline:
        time.sleep(0.5)
        snapshot = _download_snapshot()
        active_parts = [
            path for path, signature in snapshot.items()
            if path.suffix == ".part" and before.get(path) != signature
        ]
        if active_parts:
            stable = None
            stable_samples = 0
            continue
        changed = [
            path for path, signature in snapshot.items()
            if path.suffix != ".part" and before.get(path) != signature
        ]
        if not changed:
            continue
        if len(changed) != 1:
            raise ChatError(f"generated file click produced multiple downloads: {changed}")
        downloaded = changed[0]
        signature = snapshot[downloaded]
        if signature[1] == 0:
            continue
        candidate = (downloaded, signature)
        if candidate == stable:
            stable_samples += 1
        else:
            stable = candidate
            stable_samples = 1
        if stable_samples < 3:
            continue
        suggested = re.sub(r"^download\s+", "", label, flags=re.IGNORECASE).strip()
        name = _safe_file_name(suggested, downloaded.as_uri(), "", index + 1)
        destination = out_dir / name
        if destination.exists():
            raise ChatError(f"refusing to overwrite generated file: {destination}")
        shutil.move(downloaded, destination)
        return destination
    raise ChatError(f"generated file download timed out: {label}")


def harvest(m: Marionette, out_dir: Path, stamp: str, slug: str) -> Reply:
    data = m.sync_script(HARVEST, [TURN, USER])
    if not data:
        raise ChatError("no reply turn appeared in the conversation")
    reply = Reply(text=str(data.get("text", "")), url=str(data.get("url", "")))
    out_dir.mkdir(parents=True, exist_ok=True)
    for index, src in enumerate(data.get("images") or [], start=1):
        blob = _fetch_blob(m, str(src), "image", index)
        suffix = EXT.get(str(blob.get("type", "")).split(";")[0], ".png")
        name = f"{stamp}-{slug}{'' if index == 1 else f'-{index}'}{suffix}"
        dest = out_dir / name
        if dest.exists():
            raise ChatError(f"refusing to overwrite generated image: {dest}")
        dest.write_bytes(_blob_bytes(blob))
        reply.images.append(dest)
    for index, item in enumerate(data.get("files") or [], start=1):
        if not isinstance(item, dict) or not item.get("src"):
            raise ChatError(f"invalid generated file metadata {index}: {item}")
        src = str(item["src"])
        blob = _fetch_blob(m, src, "file", index)
        name = _safe_file_name(
            str(item.get("name") or ""), src, str(blob.get("disposition") or ""), index)
        dest = out_dir / name
        if dest.exists():
            raise ChatError(f"refusing to overwrite generated file: {dest}")
        dest.write_bytes(_blob_bytes(blob))
        _append_unique_file(reply, dest)
    if not data.get("files"):
        for index, label in enumerate(data.get("downloads") or []):
            _append_unique_file(reply, _download_button(m, index, str(label), out_dir))
    return reply


def ask_on(
    m: Marionette,
    prompt: str,
    *,
    effort: str | None = None,
    attachments: list[Path] | None = None,
    out_dir: Path = OUT_DIR,
    stamp: str,
    timeout: float = 900.0,
) -> Reply:
    """One ask on an already-open conversation, so a caller that owns the browser
    (solwebd) runs the same path as the standalone CLI."""
    if effort:
        set_effort(m, effort)
    attach(m, attachments or [])
    type_prompt(m, prompt)
    before = int(m.sync_script(TURN_COUNT, [TURN, USER]) or 0)
    send(m, before, timeout=timeout)
    return harvest(m, out_dir, stamp, slugify(prompt))


def ask(
    prompt: str,
    *,
    effort: str | None = None,
    attachments: list[Path] | None = None,
    out_dir: Path = OUT_DIR,
    stamp: str,
    mode: str = "virtual",
    reseed: bool = False,
    timeout: float = 900.0,
) -> Reply:
    with Session(mode=mode, reseed=reseed) as m:
        m.navigate("https://chatgpt.com/")
        wait_ready(m)
        use_chat_surface(m)
        return ask_on(m, prompt, effort=effort, attachments=attachments,
                      out_dir=out_dir, stamp=stamp, timeout=timeout)