"""Page Designer's AI Assist (target_type "page"), POST /ai-assistant/sessions
+ POST .../messages: an attached screenshot is run through OCR (app/core/ocr.py)
rather than sent to the model as an image — the model only ever sees the
resulting JSON text. Both the OCR step and the model call are mocked here
(see tests/test_ocr.py for real-Tesseract coverage of the OCR step itself),
so this exercises the plumbing between them without depending on either the
tesseract binary or a real model endpoint being reachable."""

import io
import json

import app.routers.ai_assistant as ai_assistant_router
from app.core.ocr import OcrUnavailableError


async def _upload_png(client, headers) -> str:
    resp = await client.post("/uploads", headers=headers, files={"file": ("screenshot.png", io.BytesIO(b"\x89PNG\r\n\x1a\n" + b"0" * 32), "image/png")})
    assert resp.status_code == 201, resp.text
    return resp.json()["id"]


async def _new_page_session(client, headers) -> str:
    resp = await client.post("/ai-assistant/sessions", headers=headers, json={"target_type": "page"})
    assert resp.status_code == 201, resp.text
    return resp.json()["id"]


async def test_ocr_json_reaches_the_model_prompt_and_never_raw_image_bytes(client, auth_headers, monkeypatch):
    captured_messages: list[list[dict]] = []

    async def fake_chat_completion(messages):
        captured_messages.append(messages)
        return json.dumps({"message": "Drafted a form.", "blocks": [{"id": "b1", "type": "text", "text": "hi"}]})

    monkeypatch.setattr(
        ai_assistant_router,
        "screenshot_to_layout_json",
        lambda raw_bytes: {
            "engine": "tesseract",
            "image_width": 800,
            "image_height": 400,
            "elements": [
                {"text": "Applicant Name", "left": 0.04, "top": 0.1, "width": 0.3, "height": 0.06},
                {"text": "Submit", "left": 0.75, "top": 0.82, "width": 0.1, "height": 0.05},
            ],
        },
    )
    monkeypatch.setattr(ai_assistant_router, "chat_completion", fake_chat_completion)

    session_id = await _new_page_session(client, auth_headers)
    file_id = await _upload_png(client, auth_headers)

    resp = await client.post(
        f"/ai-assistant/sessions/{session_id}/messages",
        headers=auth_headers,
        json={"message": "Reproduce this form.", "image_file_id": file_id},
    )
    assert resp.status_code == 200, resp.text
    body = resp.json()
    assert body["ollama_error"] is None
    assert body["session"]["resulting_schema"]["blocks"][0]["id"] == "b1"

    assert len(captured_messages) == 1
    system_prompt = captured_messages[0][0]["content"]
    assert captured_messages[0][0]["role"] == "system"
    # the OCR-detected text and its normalized position both made it into the prompt
    assert "Applicant Name" in system_prompt
    assert '"left":0.04' in system_prompt or '"left": 0.04' in system_prompt

    # never the raw bytes: no base64/data-URI image ever appears anywhere sent to the model
    whole_payload = json.dumps(captured_messages)
    assert "data:image" not in whole_payload
    assert "\\x89PNG" not in whole_payload

    # the session's own stored history is a short note, not the OCR JSON or the image bytes
    stored_user_message = body["session"]["prompts"][0]["content"]
    assert "screenshot.png" in stored_user_message and "2 text region(s) read via OCR" in stored_user_message
    assert "Applicant Name" not in stored_user_message


async def test_a_refinement_turn_without_a_new_image_sends_no_ocr_elements(client, auth_headers, monkeypatch):
    captured_messages: list[list[dict]] = []

    async def fake_chat_completion(messages):
        captured_messages.append(messages)
        return json.dumps({"message": "ok", "blocks": []})

    monkeypatch.setattr(ai_assistant_router, "chat_completion", fake_chat_completion)
    session_id = await _new_page_session(client, auth_headers)

    resp = await client.post(f"/ai-assistant/sessions/{session_id}/messages", headers=auth_headers, json={"message": "Make the button say Save instead."})
    assert resp.status_code == 200, resp.text
    assert "(none)" in captured_messages[0][0]["content"]


async def test_ocr_failure_reports_ollama_error_and_leaves_the_session_untouched(client, auth_headers, monkeypatch):
    def fail(raw_bytes):
        raise OcrUnavailableError("Tesseract OCR isn't installed on the server.")

    monkeypatch.setattr(ai_assistant_router, "screenshot_to_layout_json", fail)

    session_id = await _new_page_session(client, auth_headers)
    file_id = await _upload_png(client, auth_headers)

    resp = await client.post(
        f"/ai-assistant/sessions/{session_id}/messages",
        headers=auth_headers,
        json={"message": "Reproduce this form.", "image_file_id": file_id},
    )
    assert resp.status_code == 200, resp.text
    body = resp.json()
    assert body["ollama_error"] and "Tesseract" in body["ollama_error"]
    assert body["session"]["prompts"] == []
    assert body["session"]["resulting_schema"] is None


async def test_non_image_attachment_is_rejected_before_any_ocr_attempt(client, auth_headers):
    other = await client.post("/uploads", headers=auth_headers, files={"file": ("notes.txt", io.BytesIO(b"hello"), "text/plain")})
    assert other.status_code == 201, other.text
    session_id = await _new_page_session(client, auth_headers)

    resp = await client.post(
        f"/ai-assistant/sessions/{session_id}/messages",
        headers=auth_headers,
        json={"message": "Reproduce this.", "image_file_id": other.json()["id"]},
    )
    assert resp.status_code == 400