import os
import replicate

from dotenv import load_dotenv


load_dotenv()
replicate_client = replicate.Client(api_token= os.getenv("REPLICATE_API_KEY"))

def caption_generation(
    video: bytes,
    caption_size: int,
    highlight_color: str):
    output = replicate_client.run(
        "shreejalmaharjan-27/tiktok-short-captions:46bf1c12c77ad1782d6f87828d4d8ba4d48646b8e1271b490cb9e95ccdbc4504",
        input={
            "model": "large-v3",
            "video": f"data:video/mp4;base64,{video}",
            "language": "auto",
            "temperature": 0,
            "caption_size": caption_size,
            "highlight_color": highlight_color,
            "suppress_tokens": "-1",
            "logprob_threshold": -1,
            "no_speech_threshold": 0.6,
            "condition_on_previous_text": True,
            "compression_ratio_threshold": 2.4,
            "temperature_increment_on_fallback": 0.2
        }
    )

    urls = []
    if isinstance(output, list) and output:
        first = output[0]
        url = getattr(first, "url", str(first))
        urls = [url]
    elif isinstance(output, str):
        urls = [output]
    elif hasattr(output, "url"):
        urls = [getattr(output, "url")]
    if urls:
        return urls[0]