from __future__ import annotations

from attrs import define

from griptape.artifacts import ImageArtifact
from griptape.drivers.image_generation.base_image_generation_driver import BaseImageGenerationDriver


@define
class MockImageGenerationDriver(BaseImageGenerationDriver):
    model: str = "test-model"

    def try_text_to_image(self, prompts: list[str], negative_prompts: list[str] | None = None) -> ImageArtifact:
        return ImageArtifact(value="mock image", width=512, height=512, format="png")

    def try_image_variation(
        self, prompts: list[str], image: ImageArtifact, negative_prompts: list[str] | None = None
    ) -> ImageArtifact:
        return ImageArtifact(value="mock image", width=512, height=512, format="png")

    def try_image_inpainting(
        self,
        prompts: list[str],
        image: ImageArtifact,
        mask: ImageArtifact,
        negative_prompts: list[str] | None = None,
    ) -> ImageArtifact:
        return ImageArtifact(value="mock image", width=512, height=512, format="png")

    def try_image_outpainting(
        self,
        prompts: list[str],
        image: ImageArtifact,
        mask: ImageArtifact,
        negative_prompts: list[str] | None = None,
    ) -> ImageArtifact:
        return ImageArtifact(value="mock image", width=512, height=512, format="png")
