import base64
from io import BytesIO
from unittest.mock import patch

import pytest
from fastapi.testclient import TestClient
from PIL import Image
from flux3_engine import Engine, GenerationRequest, LegacyRequest, Settings
from web import create_app


class FakeEngine:
    def __init__(self):
        self.pipe = True
        self.requests = []

    def health(self):
        return {'ready': True, 'quantization': 'bf16', 'default_steps': 16}

    def generate(self, request, image_format='PNG'):
        self.requests.append(request)
        data = BytesIO()
        Image.new('RGB', (request.width, request.height)).save(data, format=image_format)
        return {'image': base64.b64encode(data.getvalue()).decode(), 'metadata': {},
                'mime_type': 'image/jpeg' if image_format == 'JPEG' else 'image/png'}

    legacy = Engine.legacy


@pytest.fixture
def api(monkeypatch):
    monkeypatch.setenv('FLUX_API_TOKEN', 'test-only')
    engine = FakeEngine()
    with TestClient(create_app(engine)) as client:
        yield client, engine


def test_auth_required(api):
    client, engine = api
    for route in ('/sdxl', '/generate'):
        assert client.post(route, json={'inputs': 'portrait'}).status_code == 401
    assert client.get('/health').status_code == 401
    assert not engine.requests


@pytest.mark.parametrize('route', ['/docs', '/docs/oauth2-redirect', '/redoc', '/openapi.json'])
def test_documentation_disabled(api, route):
    client, engine = api
    assert client.get(route).status_code == 404
    assert client.get(route, headers={'Authorization': 'Bearer test-only'}).status_code == 404
    assert not engine.requests


@pytest.mark.parametrize('startup_error', [None, 'Model load failed'])
def test_unready_engine_returns_503(monkeypatch, startup_error):
    monkeypatch.setenv('FLUX_API_TOKEN', 'test-only')
    engine = FakeEngine()
    engine.pipe = True if startup_error else None
    with TestClient(create_app(engine, startup_error)) as client:
        headers = {'Authorization': 'Bearer test-only'}
        assert client.get('/health', headers=headers).status_code == 503
        for route in ('/sdxl', '/generate'):
            assert client.post(route, headers=headers, json={'inputs': 'portrait'}).status_code == 503
    assert not engine.requests


def test_legacy_payload_returns_jpeg_at_16_steps(api):
    client, engine = api
    response = client.post('/sdxl', headers={'Authorization': 'Bearer test-only'}, json={
        'inputs': 'portrait', 'num_inference_steps': 40, 'guidance_scale': 6,
        'negative_prompt': 'blurry', 'use_refiner': False, 'high_noise_frac': .8,
        'width': 512, 'height': 768})
    assert response.status_code == 200
    image = Image.open(BytesIO(base64.b64decode(response.json()['image'])))
    assert image.format == 'JPEG' and image.size == (512, 768)
    assert engine.requests[0].num_inference_steps == 16


def test_native_resolution_and_steps(api):
    client, engine = api
    headers = {'Authorization': 'Bearer test-only'}
    assert client.post('/generate', headers=headers,
                       json={'inputs': 'portrait', 'width': 768, 'height': 768}).status_code == 200
    assert engine.requests[-1].num_inference_steps == 16
    for payload in ({'inputs': ' '}, {'inputs': 'p', 'width': 1000},
                    {'inputs': 'p', 'num_inference_steps': 40}):
        assert client.post('/generate', headers=headers, json=payload).status_code == 422


def test_fail_closed_without_api_token(monkeypatch):
    monkeypatch.delenv('FLUX_API_TOKEN', raising=False)
    with pytest.raises(RuntimeError):
        create_app(FakeEngine())


def test_private_repo_does_not_reuse_upstream_commit(monkeypatch):
    monkeypatch.setenv('FLUX_MODEL_ID', 'example/private-model')
    monkeypatch.delenv('FLUX_MODEL_REVISION', raising=False)
    monkeypatch.delenv('FLUX_QUANTIZATION', raising=False)
    assert Settings.from_env() == Settings('example/private-model', 'main', 'q8_0')


def test_hf_handler_uses_mounted_repository():
    with patch('handler.Engine') as engine:
        from handler import EndpointHandler
        handler = EndpointHandler('/repository')
        payload = {'inputs': 'portrait'}
        handler(payload)
        engine.assert_called_once_with(model_path='/repository')
        engine.return_value.load.assert_called_once()
        engine.return_value.legacy.assert_called_once_with(payload)