import base64
import io
from unittest.mock import patch

from fastapi.testclient import TestClient
from PIL import Image


def _png_data_url() -> str:
    buf = io.BytesIO()
    Image.new("RGB", (8, 8), "white").save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/png;base64,{b64}"


def _client() -> TestClient:
    from surya.endpoint.app import app

    return TestClient(app)


def _image_message():
    return [
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": _png_data_url()}},
                {"type": "text", "text": "extract the text"},
            ],
        }
    ]


def test_returns_openai_chat_completion_shape():
    body = {"model": "surya-ocr", "messages": _image_message()}
    fake = {"text_lines": "hello\nworld", "ocr_text_json": {"blocks": []}, "elapsed_seconds": 1.23}
    with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake):
        r = _client().post("/v1/chat/completions", json=body)
    assert r.status_code == 200
    d = r.json()
    assert d["object"] == "chat.completion"
    assert d["model"] == "surya-ocr"
    assert d["choices"][0]["message"]["content"] == "hello\nworld"
    assert d["choices"][0]["finish_reason"] == "stop"
    assert d["surya"]["ocr_text_json"] == {"blocks": []}
    assert d["surya"]["elapsed_seconds"] == 1.23
    assert "usage" in d


def test_mode_full_page_skips_text_detection():
    body = {"model": "m", "mode": "full_page", "messages": _image_message()}
    fake = {"text_lines": "x", "ocr_text_json": {}, "elapsed_seconds": None}
    with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake) as m:
        r = _client().post("/v1/chat/completions", json=body)
    assert r.status_code == 200
    assert m.call_args.kwargs["skip_text_detection"] is True


def test_mode_table_routes_to_table_image():
    body = {"model": "m", "mode": "table", "messages": _image_message()}
    fake = {"text_lines": "tbl", "ocr_text_json": {"x": 1}, "elapsed_seconds": None}
    with patch("surya.endpoint.service.table_image", return_value=fake) as m:
        r = _client().post("/v1/chat/completions", json=body)
    assert r.status_code == 200
    assert m.called
    assert r.json()["choices"][0]["message"]["content"] == "tbl"


def test_ocr_with_boxes_false_omits_structured_json():
    body = {"model": "m", "ocr_with_boxes": False, "messages": _image_message()}
    fake = {"text_lines": "x", "ocr_text_json": {"a": 1}, "elapsed_seconds": None}
    with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake):
        r = _client().post("/v1/chat/completions", json=body)
    assert r.json()["surya"]["ocr_text_json"] is None


def test_stream_true_returns_400():
    body = {"model": "m", "stream": True, "messages": _image_message()}
    r = _client().post("/v1/chat/completions", json=body)
    assert r.status_code == 400
    assert r.json()["error"]["type"] == "invalid_request_error"


def test_missing_image_returns_400():
    body = {"model": "m", "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]}
    r = _client().post("/v1/chat/completions", json=body)
    assert r.status_code == 400
    assert r.json()["error"]["type"] == "invalid_request_error"


def test_pipeline_error_returns_500():
    body = {"model": "m", "messages": _image_message()}
    with patch("surya.endpoint.service.ocr_via_batcher", side_effect=RuntimeError("boom")):
        r = _client().post("/v1/chat/completions", json=body)
    assert r.status_code == 500
    assert r.json()["error"]["type"] == "internal_error"