"""Does the tiny VLM echo the prompt - and does masking the logits fix it?

The card's open problem, as a measurement instead of an anecdote:

    python measure_vlm.py                # the built-in busy screen, CPU, no network
    python measure_vlm.py --screen live  # whatever is on this machine right now

It runs the same screen through the same model twice: once in the shipped `json`
mode (free text, where the 256M model answers by continuing the prompt), and once
in `choice` mode (digits only, so an echo is not representable). Then it prints
what each one did, and how long it took.
"""

from __future__ import annotations

import argparse
import time
from pathlib import Path

from tiny_computer_agent import (
    DecisionMenu, GrabScreen, ImageScreen, IndexedScreen, make_detector,
)
from tiny_computer_agent.models import (
    DEFAULT_TEXT_MODEL, DEFAULT_VISION_MODEL, TextPlannerBrain, VLMBrain,
)

GOAL = "click 'Save' to save the document"
OUT = Path("preview")


def busy_desktop(width: int = 1100, height: int = 700):
    """A deliberately noisy screen: menus, a dialog, prose, a status bar.

    The echo failure showed up on a *busy* live screen and not on a clean
    synthetic one, so the fixture has to be busy to reproduce it.
    """
    from PIL import Image, ImageDraw, ImageFont

    def font(size):
        for candidate in ("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
                          "C:/Windows/Fonts/arial.ttf"):
            try:
                return ImageFont.truetype(candidate, size)
            except Exception:
                continue
        return ImageFont.load_default()

    image = Image.new("RGB", (width, height), (238, 240, 244))
    draw = ImageDraw.Draw(image)
    draw.rectangle([0, 0, width, 30], fill=(38, 40, 48))
    for index, label in enumerate(["File", "Edit", "View", "Insert", "Format", "Tools",
                                   "Window", "Help"]):
        draw.text((16 + index * 74, 8), label, fill="white", font=font(14))

    draw.rectangle([0, 30, 220, height], fill=(226, 229, 234))
    for index, item in enumerate(["Documents", "Downloads", "Pictures", "Music",
                                  "Reports", "Archives", "Scratch"]):
        draw.text((18, 52 + index * 30), item, fill=(40, 44, 52), font=font(14))

    draw.rectangle([280, 90, 1000, 470], fill="white", outline=(120, 124, 132), width=2)
    draw.text((300, 106), "Untitled document - unsaved changes", fill=(28, 30, 34), font=font(16))
    draw.text((300, 150),
              "Your changes will be lost if you close this window without saving.",
              fill=(70, 74, 82), font=font(13))
    draw.text((300, 172), "The file has not been written to disk yet.", fill=(70, 74, 82),
              font=font(13))
    draw.text((300, 214), "File name", fill=(60, 62, 70), font=font(14))
    draw.rectangle([300, 238, 960, 272], fill=(252, 252, 253), outline=(120, 124, 132))
    draw.text((310, 247), "report.txt", fill=(30, 32, 36), font=font(14))
    draw.rectangle([300, 330, 470, 374], fill=(66, 111, 220), outline=(36, 62, 140), width=2)
    draw.text((330, 342), "Save", fill="white", font=font(16))
    draw.rectangle([500, 330, 670, 374], fill=(228, 230, 235), outline=(120, 124, 132))
    draw.text((520, 342), "Don't Save", fill=(30, 32, 36), font=font(15))
    draw.rectangle([700, 330, 860, 374], fill=(228, 230, 235), outline=(120, 124, 132))
    draw.text((730, 342), "Cancel", fill=(30, 32, 36), font=font(16))

    draw.rectangle([280, 510, 1000, 660], fill=(248, 249, 251), outline=(150, 154, 160))
    draw.text((300, 524), "Recent files", fill=(40, 44, 52), font=font(15))
    for index, name in enumerate(["report-final-v3.txt", "notes.md", "budget.xlsx",
                                  "draft.txt", "screenshot.png"]):
        draw.text((300, 552 + index * 20), name, fill=(66, 70, 78), font=font(13))

    draw.rectangle([0, height - 28, width, height], fill=(52, 56, 64))
    draw.text((12, height - 22), "report.txt  -  modified  -  UTF-8  -  line 42, column 7",
              fill=(206, 210, 216), font=font(13))
    return image


def _open_gtk_window():
    """Open the GTK demo window on a fresh session bus, so AT-SPI can see it."""
    import os
    import subprocess
    import sys
    import time

    address = subprocess.run(["dbus-daemon", "--session", "--fork", "--print-address=1"],
                             capture_output=True, text=True).stdout.strip()
    os.environ.setdefault("DISPLAY", ":99")
    os.environ["DBUS_SESSION_BUS_ADDRESS"] = address
    os.environ["NO_AT_BRIDGE"] = "0"
    ready = Path("/tmp/measure-gtk.ready")
    ready.unlink(missing_ok=True)
    process = subprocess.Popen(
        [sys.executable, "showcase_scenes/gtk_dialog.py", "/tmp/measure-gtk.marker", str(ready)],
        env=dict(os.environ), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    for _ in range(60):
        if ready.exists():
            break
        time.sleep(0.25)
    time.sleep(2.5)
    print("GTK window mapped:", ready.exists())
    return process


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--screen", default="builtin", choices=("builtin", "live"))
    parser.add_argument("--detector", default="ocr",
                        choices=("ocr", "auto", "best", "atspi", "grid"),
                        help="how the menu is built; 'best' uses the accessibility "
                             "tree first, so only real widgets are offered")
    parser.add_argument("--open-gtk", action="store_true",
                        help="open the GTK window (with a session bus) and index that")
    parser.add_argument("--brain", default="vlm", choices=("vlm", "text-planner"),
                        help="which tiny model to ask; text-planner reads the index only")
    parser.add_argument("--model", default=None)
    parser.add_argument("--goal", default=GOAL)
    parser.add_argument("--modes", default="json,choice,scan,rank")
    parser.add_argument("--sweep", action="store_true",
                        help="ask the same four items about four different goals, and "
                             "count how often each mode gets it right")
    parser.add_argument("--max-new-tokens", type=int, default=64)
    args = parser.parse_args(argv)

    OUT.mkdir(exist_ok=True)
    gtk = None
    if args.open_gtk:
        gtk = _open_gtk_window()
    try:
        if args.screen == "live":
            source = GrabScreen(monitor=1, max_width=1100)
            image = source.capture().frame
        else:
            source = ImageScreen(busy_desktop())
            image = source.capture().frame
        image.save(OUT / "vlm_echo_screen.png")
        observation = IndexedScreen(source, make_detector(args.detector)).capture()
    finally:
        if gtk is not None:
            gtk.terminate()
    menu = DecisionMenu(observation)
    print(f"brain: {args.brain} | screen: {args.screen} ({image.width}x{image.height}) "
          f"| detector: {args.detector}")
    print(f"index: {len(observation.elements)} elements -> menu of {len(menu)}")
    print(menu.render())
    print()

    brain_class = VLMBrain if args.brain == "vlm" else TextPlannerBrain
    model_id = args.model or (DEFAULT_VISION_MODEL if args.brain == "vlm" else DEFAULT_TEXT_MODEL)

    if args.sweep:
        # every item, asked about every goal: the same shot, the same weights,
        # only the question changes - which is the whole point of the item
        labels = [element.label for element in observation.elements][:4]
        print(f"sweep: {len(labels)} items, {len(labels)} goals, brain {args.brain}")
        print()
        print(f"{'mode':8} {'right':>7}  seconds  what it chose")
        for mode in args.modes.split(","):
            brain = brain_class(model_id, decision_mode=mode, retries=3, votes=1,
                                max_new_tokens=args.max_new_tokens)
            right, chosen, start = 0, [], time.time()
            for label in labels:
                goal = f"click '{label}' to do that"
                try:
                    action = brain.decide(goal, observation)
                    got = action.label if action.kind == "click" else action.kind
                except Exception as exc:
                    got = f"<{type(exc).__name__}>"
                chosen.append(got)
                right += int(got == label)
            print(f"{mode:8} {right}/{len(labels):<5} "
                  f"{round(time.time() - start, 1):8}  {chosen}")
        return 0

    results = []
    for mode in args.modes.split(","):
        brain = brain_class(model_id, decision_mode=mode, retries=3, votes=1,
                            max_new_tokens=args.max_new_tokens)
        start = time.time()
        try:
            action = brain.decide(args.goal, observation)
            outcome = action.describe()
        except Exception as exc:
            outcome = f"{type(exc).__name__}: {str(exc)[:110]}"
        elapsed = time.time() - start
        reply = (brain.last_reply or "").strip()
        echoed = bool(reply) and (reply == (brain.last_prompt or "").strip()
                                  or "You are controlling a computer" in reply)
        results.append((mode, echoed, outcome, round(elapsed, 1)))
        print(f"--- decision_mode={mode} ---")
        print(f"  echoed the prompt: {echoed}")
        print(f"  reply   : {reply[:160]!r}")
        print(f"  outcome : {outcome}")
        print(f"  seconds : {elapsed:.1f}")

    print()
    print(f"{'mode':8} {'echoed':7} {'seconds':8} outcome")
    for mode, echoed, outcome, elapsed in results:
        print(f"{mode:8} {str(echoed):7} {elapsed:8} {outcome[:90]}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())