from smolagents import CodeAgent, tool, load_tool, HfApiModel
from youtube_transcript_api import YouTubeTranscriptApi
import datetime, pytz, yaml, re
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# ---------- TOOLS ----------
@tool
def get_transcript(url: str) -> str:
    """
    Télécharge la transcription (FR ou EN) d’une vidéo YouTube.
    Args:
        url: lien complet YouTube
    """
    video_id = re.search(r"(?:v=|youtu\.be/)([^&\n?#]+)", url)
    if not video_id:
        return "Impossible de détecter l’ID vidéo."
    video_id = video_id.group(1)

    try:
        transcript = YouTubeTranscriptApi.get_transcript(video_id, languages=['fr', 'en'])
        return " ".join(seg["text"] for seg in transcript)
    except Exception as e:
        return f"Erreur transcript : {e}"

@tool
def summarize(text: str, max_chars: int = 1500) -> str:
    """
    Résume le texte fourni en 8-10 bullet points clairs.

    Args:
        text (str): Le texte à résumer.
        max_chars (int, optional): Nombre maximum de caractères à considérer dans le texte. Par défaut : 1500.

    Returns:
        str: Un prompt de résumé à envoyer au modèle.
    """
    snippet = text[:max_chars]
    prompt = (
        "Résume le texte suivant en 8‑10 bullet points clairs :\n\n"
        f"{snippet}\n\n# Résumé :"
    )
    return prompt


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Renvoie l’heure actuelle dans un fuseau horaire donné.

    Args:
        timezone (str): Nom du fuseau horaire (ex: 'Europe/Paris', 'America/New_York').

    Returns:
        str: Heure locale actuelle dans le fuseau horaire spécifié.
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time: {str(e)}"


# ---------- AGENT ----------

model = HfApiModel(
    model_id="mistralai/Mistral-7B-Instruct-v0.1",
    temperature=0.5,
    max_tokens=2048,
)

final_answer = FinalAnswerTool()

# Charge les prompts
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[final_answer, get_transcript, summarize, get_current_time_in_timezone],
    max_steps=6,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()