#!/usr/bin/env python3
"""Stub test for app.py _stream_core/_chat_stream fixes.

No Gradio launch, no Ollama, no Qdrant: the retriever and the LLM
stream are monkeypatched.  Proves:
  bug 1: _stream_core writes the module-global chunks_holder and
         _chat_stream's gallery receives the same chunks (page PNGs).
  bug 2: _chat_stream yields history + a growing assistant bubble,
         and the final history contains the full answer.
"""
from __future__ import annotations

import ast
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import app
import config


def main() -> int:
    failures = []

    # --- bug 1 precondition: a single shared module-global -----------
    tree = ast.parse((Path(__file__).resolve().parents[1] / "app.py").read_text())
    assigns = [n.targets[0].id for n in ast.walk(tree)
               if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name)
               and n.targets[0].id == "chunks_holder"]
    if len(assigns) != 1:
        failures.append(f"bug1: chunks_holder assigned {len(assigns)}x at module level, expected 1")
    if id(app._chat_stream.__globals__["chunks_holder"]) != id(app.chunks_holder):
        failures.append("bug1: _chat_stream globals chunks_holder is not app.chunks_holder")
    if id(app._stream_core.__globals__["chunks_holder"]) != id(app.chunks_holder):
        failures.append("bug1: _stream_core globals chunks_holder is not app.chunks_holder")

    # --- stub the retriever + LLM stream -----------------------------
    fake_chunks = [{"sha": "deadbeef", "page": 1, "kind": "text",
                    "text": "torque 45 Nm part KX40-8817 setpoint 2.4 bar",
                    "score": 0.9}]

    class FakeRetriever:
        def retrieve(self, query, top_n):
            return {"chunks": fake_chunks}

        def close(self):
            pass

    app._get_retriever = lambda: FakeRetriever()
    app._ollama_stream = lambda messages: iter(["The tor", "que is ", "45 Nm."])
    config.ENABLE_QUERY_ENHANCEMENT = False

    # --- bug 2: streamed assistant bubble grows and lands in history -
    history = [{"role": "user", "content": "What is the torque?"}]
    seen = []
    for hist, sources, gallery in app._chat_stream("What is the torque?", history):
        seen.append((hist, gallery))
    last_hist, _ = seen[-1]
    if not last_hist or last_hist[-1].get("role") != "assistant":
        failures.append("bug2: final history does not end with an assistant message")
    else:
        text = last_hist[-1]["content"]
        if text != "The torque is 45 Nm.":
            failures.append(f"bug2: assistant text {text!r} != streamed answer")
    if not all(len(h) > len(history) for h, _ in seen):
        failures.append("bug2: intermediate yields do not include the assistant bubble")
    if last_hist[:-1] != history:
        failures.append("bug2: user history was mutated")

    # --- bug 1: gallery receives page PNG of the retrieved chunk -----
    with tempfile.TemporaryDirectory() as td:
        page_png = Path(td) / "deadbeef" / "p0001.png"
        page_png.parent.mkdir()
        page_png.write_bytes(b"png")
        config.ENABLE_VISION_ANSWERS = True
        config.PAGES_DIR = Path(td)
        hist, sources, gallery = list(app._chat_stream("q", history))[-1]
        if gallery != [str(page_png)]:
            failures.append(f"bug1: gallery {gallery!r} != {[str(page_png)]}")
        if app.chunks_holder["chunks"] != fake_chunks:
            failures.append("bug1: module-global chunks_holder not populated by _stream_core")

    if failures:
        print("FAIL")
        for f in failures:
            print(" -", f)
        return 1
    print("PASS: bug1 (shared chunks_holder -> gallery) and bug2 (assistant bubble in history)")
    return 0


if __name__ == "__main__":
    sys.exit(main())