from phi.assistant import Assistant
from phi.tools.wikipedia import WikipediaTools
from google import genai
import os

client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))

class GeminiWrapper:
    """Custom wrapper to use new Gemini SDK inside phi"""

    def __init__(self, model="gemini-2.0-flash"):
        self.model = model

    def __call__(self, prompt, **kwargs):
        response = client.models.generate_content(
            model=self.model,
            contents=prompt
        )
        return response.text


def get_content_enhancer_agent(knowledge_base=None):
    return Assistant(
        llm=GeminiWrapper(),
        tools=[WikipediaTools()],
        knowledge_base=knowledge_base,
        search_knowledge=True,
        description="You are an expert multimedia content creator and writer.",
        instructions=[
            "You will receive a summary from a video analysis.",
            "Enhance it using knowledge + Wikipedia.",
            "Create a cohesive narrative.",
            "Insert [IMAGE: caption=\"...\"] placeholders.",
            "Make prompts vivid and visual.",
            "Return only final text."
        ],
        show_tool_calls=True,
    )