import os
from llama_index.llms.anthropic import Anthropic
from llama_index.core.tools import FunctionTool
from llama_index.core.agent import ReActAgent

from tools import compare_faces, compare_voices, scan_video

# Wrap functions as tools
face_tool = FunctionTool.from_defaults(
    fn=compare_faces,
    name="Compare_Faces",
    description="Compare two images to check if the faces match."
)

voice_tool = FunctionTool.from_defaults(
    fn=compare_voices,
    name="Compare_Voices",
    description="Compare two audio files to check if voices belong to the same person."
)

video_tool = FunctionTool.from_defaults(
    fn=scan_video,
    name="Scan_Video",
    description="Analyze a video using a reference image and optionally a YouTube link to detect face swaps."
)

# Anthropic LLM
llm = Anthropic(model="claude-3-opus-20240229")

# Create agent
agent = ReActAgent.from_tools(
    tools=[face_tool, voice_tool, video_tool],
    llm=llm,
    verbose=True,
)


# Chat function
def chat_with_agent(user_input, img1=None, img2=None, audio1=None, audio2=None, video=None, ref_img=None):
    context = "You are an AI assistant. Use the following tools ONLY:\n\n"
    context += "- 'Compare Faces' to compare two face images (requires img1_path, img2_path).\n"
    context += "- 'Compare Voices' to compare two audio samples (requires audio1, audio2).\n"
    context += "- 'Scan Video' to scan a video for deepfakes (requires video_file, ref_img).\n\n"
    context += "Be precise in tool naming. Do not invent tool names.\n\n"
    context += "User prompt:\n" + user_input + "\n"

    if img1 and img2:
        context += f"\nUploaded face images:\n - Image 1: {img1}\n - Image 2: {img2}"
    if audio1 and audio2:
        context += f"\nUploaded voice samples:\n - Audio 1: {audio1}\n - Audio 2: {audio2}"
    if video:
        context += f"\nUploaded video file:\n - Video: {video}"
    if ref_img:
        context += f"\nUploaded reference image:\n - Reference Image: {ref_img}\n"

    try:
        response = agent.chat(context)
        return response.response
    except Exception as e:
        return f"❌ Agent error: {str(e)}"