from pydantic import BaseModel
import json
from image_api import get_random_image
import base64
import requests


class TVCrewOutput(BaseModel):
    image_keyword: str


class TVCrew:
    def __init__(self, client, system_prompt: str):
        self.model = "google/gemma-3-12b-it"
        self.system_prompt = system_prompt
        self.context = ""
        self.client = client

    def suggest_image(self) -> tuple[str, str] | tuple[None, None]:
        """
        Finds a image related to conversation and returns the url and caption
        """
        messages_tv_crew = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": self.context},
        ]
        response_tv_crew = self.client.chat.completions.parse(
            model=self.model,
            messages=messages_tv_crew,
            response_format=TVCrewOutput,
        )
        response_tv_crew = response_tv_crew.choices[0].message.content
        response_tv_crew_json = json.loads(response_tv_crew)
        image_keyword = response_tv_crew_json["image_keyword"]
        if image_keyword == "none":
            return None, None
        image_url, alt = get_random_image(image_keyword)
        caption = None
        if image_url:
            caption = alt if len(alt) > 0 else image_keyword
            self.add_context(f"TV Shows: {caption}\n")
            img_bytes = requests.get(image_url).content
            image_base64 = base64.b64encode(img_bytes).decode("utf-8")
            return image_base64, caption
        return None, None

    def add_context(self, content):
        """Will be updated by main script"""
        self.context += content

    def clear_context(self):
        self.context = ""