import gradio as gr
import os
from dotenv import load_dotenv
from textwrap import dedent
from agno.agent import Agent
from agno.team.team import Team
from agno.media import Image
from agno.models.anthropic import Claude
from agno.tools.googlesearch import GoogleSearchTools

# Load API key
load_dotenv()
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
if not ANTHROPIC_API_KEY:
    raise ValueError("ANTHROPIC_API_KEY not found in .env file.")

# Define Agents
image_agent = Agent(
    name="Image Analyzer",
    role="Analyze uploaded car issue images to extract any visible mechanical problems or symptoms.",
    model=Claude(id="claude-3-5-sonnet-20240620", api_key=ANTHROPIC_API_KEY)
)

info_agent = Agent(
    name="Info Gatherer",
    role="Use online search to collect information about known issues for the given car brand, model, and year.",
    model=Claude(id="claude-3-haiku-20240307", api_key=ANTHROPIC_API_KEY),
    tools=[GoogleSearchTools()],
    show_tool_calls=False
)

agent_team = Team(
    name="Smart Repair Team",
    mode="coordinate",
    model=Claude(id="claude-3-haiku-20240307", api_key=ANTHROPIC_API_KEY),
    members=[image_agent, info_agent],
    description="You are a task router that routes tasks to a team of agents to help users diagnose and fix car issues.",
    instructions=dedent("""\
        - First, use the uploaded image to detect possible issues visually.
        - Then, use search tools to gather brand/model-specific info if available from the image extracted information.
        - Give the user a helpful, step-by-step repair guide.
        - Explain steps clearly for non-experts.
        - Mention safety precautions and signs when to visit a professional.
    """),
    expected_output=dedent("""\
        Final response should be structured with these sections:
        - Diagnosis Summary
        - Step-by-Step Repair Instructions
    """),
    success_criteria="The user receives a clear, accurate step-by-step guide for handling the car issue.",
    add_datetime_to_instructions=True,
    show_tool_calls=False,
    enable_agentic_context=True,
    show_members_responses=False,
    markdown=True
)

# Main chat function
def diagnose_chat(image_pil, brand, model, year):
    if image_pil is None:
        return "Please upload an image of the car issue to proceed."

    os.makedirs("temp_images", exist_ok=True)
    temp_path = os.path.join("temp_images", "temp_car_issue.jpg")
    image_pil.save(temp_path)

    image_media = Image(filepath=temp_path)

    # Construct message
    extra_info = []
    if brand:
        extra_info.append(f"Brand: {brand}")
    if model:
        extra_info.append(f"Model: {model}")
    if year:
        extra_info.append(f"Year: {year}")
    
    info_text = "\n".join(extra_info)
    message = "Diagnose the car issue based on the uploaded image."
    if info_text:
        message += f"\n\nAdditional details:\n{info_text}"

    response = agent_team.run(message=message, images=[image_media])
    return response.content

# Gradio UI setup
with gr.Blocks() as demo:
    gr.Markdown("# 🔧 SmartCarDoc\nUpload a car issue image to quickly get a repair plan.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Image of Car Issue")

            brand_input = gr.Textbox(label="Car Brand (optional)")
            model_input = gr.Textbox(label="Car Model (optional)")
            year_input = gr.Textbox(label="Year (optional)")

            submit_button = gr.Button("Diagnose")

        with gr.Column():
            response_text = gr.Textbox(label="Repair Plan", lines=30, interactive=False)

    submit_button.click(
        fn=diagnose_chat,
        inputs=[image_input, brand_input, model_input, year_input],
        outputs=response_text
    )

if __name__ == "__main__":
    demo.launch()