import base64
import io
import time
from pathlib import Path
from typing import Any

from dotenv import load_dotenv
from PIL import Image
from together import Together

from api.flux import FluxAPI


class TogetherAPI(FluxAPI):
    def __init__(self):
        load_dotenv()
        self._client = Together()

    @property
    def name(self) -> str:
        return "together"

    def generate_image(self, prompt: str, save_path: Path) -> float:
        start_time = time.time()
        result = self._client.images.generate(
            prompt=prompt,
            model="black-forest-labs/FLUX.1-dev",
            width=1024,
            height=1024,
            steps=28,
            n=1,
            seed=0,
            response_format="b64_json",
        )
        end_time = time.time()
        if result and hasattr(result, "data") and len(result.data) > 0:
            self._save_image_from_result(result, save_path)
        else:
            raise Exception("No result returned from Together API")
        return end_time - start_time

    def _save_image_from_result(self, result: Any, save_path: Path):
        save_path.parent.mkdir(parents=True, exist_ok=True)
        b64_str = result.data[0].b64_json
        image_data = base64.b64decode(b64_str)
        image = Image.open(io.BytesIO(image_data))
        image.save(save_path)