from pathlib import Path
import tempfile
import unittest

from fastapi.testclient import TestClient
from PIL import Image

from demo.runtime import Pipeline, uploaded_resolution
from demo.ui import create_app


class UploadResolutionTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.root = Path(self.temp.name)

    def image(self, name, size):
        path = self.root / name
        Image.new('RGB', size, (40, 100, 130)).save(path)
        return path

    def test_default_caps_each_dimension_to_the_smallest_input_and_rounds_down(self):
        for i, (size, expected) in enumerate([
            ((1600, 1200), (896, 896)), ((800, 600), (798, 588)),
            ((1280, 720), (896, 714)), ((896, 896), (896, 896)),
            ((500, 1000), (490, 896)), ((14, 14), (14, 14)),
        ]):
            with self.subTest(size=size):
                image = self.image(f'{i}.png', size)
                original = image.read_bytes()
                self.assertEqual(uploaded_resolution([image]), expected)
                self.assertEqual(image.read_bytes(), original)
        wide = self.image('wide.png', (1600, 600))
        tall = self.image('tall.png', (800, 1200))
        self.assertEqual(uploaded_resolution([wide, tall]), (798, 588))
        self.assertEqual(uploaded_resolution([]), (896, 896))
        with self.assertRaisesRegex(ValueError, 'at least 14'):
            uploaded_resolution([self.image('tiny.png', (13, 100))])

    def test_upload_callback_recomputes_dimensions_and_clear_restores_default(self):
        app = create_app(Pipeline(device='cpu'), self.root / 'results', self.root / 'no-examples')
        dep = next(d for d in app.state.demo.config['dependencies'] if d.get('api_name') == 'preview_inputs')
        components = {c['id']: c['props'] for c in app.state.demo.config['components']}
        self.assertEqual([components[i]['label'] for i in dep['outputs'][-2:]], ['Width', 'Height'])
        with TestClient(app) as client:
            files = []
            for name, size in [('large.png', (1280, 720)), ('small.png', (800, 600))]:
                image = self.image(name, size)
                with image.open('rb') as stream:
                    upload = client.post('/gradio_api/upload', files={'files': (name, stream, 'image/png')})
                files.append(dict(path=upload.json()[0], meta={'_type': 'gradio.FileData'}))
            for selected, expected in [(files, [798, 588]), (files[:1], [896, 714]), ([], [896, 896])]:
                r = client.post('/gradio_api/run/predict', json=dict(fn_index=dep['id'], data=[selected], session_hash='upload-resolution'))
                self.assertEqual(r.status_code, 200, r.text)
                self.assertEqual(r.json()['data'][-2:], expected)


if __name__ == '__main__':
    unittest.main()