from src.models.gemini import GeminiModel


def _resolution_options(model_id):
    inputs = GeminiModel.video_settings_inputs(model_id)
    resolution_input = next(i for i in inputs if i.name == "resolution")
    return resolution_input.kwargs["options"]


def test_veo_2_0_only_supports_720p():
    assert _resolution_options("veo-2.0-generate-001") == ["720"]


def test_veo_3_1_lite_has_no_4k():
    assert _resolution_options("veo-3.1-lite-generate-preview") == ["720", "1080"]


def test_veo_3_1_supports_4k():
    assert _resolution_options("veo-3.1-generate-preview") == ["720", "1080", "4K"]
    assert _resolution_options("veo-3.1-fast-generate-preview") == ["720", "1080", "4K"]


def test_veo_resolution_kwargs_maps_4k():
    kwargs = GeminiModel.video_kwargs(
        "veo-3.1-generate-preview",
        {"size": "16:9", "resolution": "4K", "seconds": 8, "person_generation": "allow_all"},
    )
    assert kwargs["resolution"] == "4k"