import os
import gradio as gr
import requests
import pandas as pd
from io import BytesIO

# --- LangChain & Groq Imports ---
from groq import Groq
from langchain_groq import ChatGroq
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.prompts import ChatPromptTemplate
from langchain.tools import Tool


# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"


# --- Custom Tool Definition using Groq ---
def transcribe_audio_from_task_id(task_id: str) -> str:
    """
    Downloads an audio file for a given task_id from the scoring server,
    transcribes it using the GROQ API with Whisper, and returns the text.
    Use this tool ONLY when a question explicitly mentions an audio file or recording.
    The task_id MUST be provided as the input.
    """
    print(f"Tool 'transcribe_audio_from_task_id' (using Groq) called with task_id: {task_id}")
    try:
        # Step 1: Download the file
        file_url = f"{DEFAULT_API_URL}/files/{task_id}"
        print(f"Downloading audio file from: {file_url}")
        audio_response = requests.get(file_url)
        audio_response.raise_for_status()

        # Step 2: Prepare the file for the Groq API
        audio_bytes = BytesIO(audio_response.content)
        audio_bytes.name = f"{task_id}.mp3" # Give the file-like object a name

        # Step 3: Initialize the Groq client and transcribe
        print("Initializing Groq client for transcription...")
        client = Groq(api_key=os.getenv("GROQ_API_KEY"))

        print("Transcribing audio with Groq's Whisper...")
        transcription = client.audio.transcriptions.create(
          file=audio_bytes,
          model="whisper-large-v3",
          response_format="text",
        )
        
        transcribed_text = str(transcription)
        print(f"Transcription successful. Result: {transcribed_text}")
        return transcribed_text

    except Exception as e:
        error_message = f"Error in Groq audio transcription tool: {e}"
        print(error_message)
        return error_message


# --- Agent Definition ---
class LangChainAgent:
    def __init__(self, groq_api_key: str, tavily_api_key: str):
        print("Initializing LangChainAgent...")
        
        # THIS IS THE CORRECTED LINE
        self.llm = ChatGroq(model_name="llama3-70b-8192", groq_api_key=groq_api_key, temperature=0.0)

        # Define all available tools
        audio_tool = Tool(
            name="audio_transcriber",
            func=transcribe_audio_from_task_id,
            description="Use this tool to transcribe an audio file. The input must be the task_id of the question.",
        )
        self.tools = [
            TavilySearchResults(max_results=3, tavily_api_key=tavily_api_key, name="web_search"),
            audio_tool,
        ]
        
        # Define the strict system prompt
        prompt = ChatPromptTemplate.from_messages([
            ("system", (
                "You are a powerful problem-solving agent. Your goal is to answer the user's question accurately. "
                "You have access to the following tools: a web search tool and an audio transcription tool.\n"
                "RULES:\n"
                "- Carefully analyze the user's question to determine if a tool is needed.\n"
                "- For questions requiring current information or facts, use the 'web_search' tool.\n"
                "- For questions that mention an audio file (.mp3, recording, voice memo, etc.), use the 'audio_transcriber' tool with the provided 'task_id'.\n"
                "- Once you have all the necessary information, you MUST provide ONLY THE FINAL ANSWER to the user's question. Do not include any extra conversation, explanations, apologies, or introductory phrases."
            )),
            ("human", "Question: {input}\nTask ID: {task_id}"),
            ("placeholder", "{agent_scratchpad}"),
        ])

        agent = create_tool_calling_agent(self.llm, self.tools, prompt)
        self.agent_executor = AgentExecutor(agent=agent, tools=self.tools, verbose=True, handle_parsing_errors=True)
        print("LangChainAgent initialized.")

    def __call__(self, question: str, task_id: str) -> str:
        print(f"Agent received question (ID: {task_id}): {question[:50]}...")
        try:
            response = self.agent_executor.invoke({"input": question, "task_id": task_id})
            answer = response.get("output", "Agent failed to produce an answer.")
        except Exception as e:
            answer = f"Agent execution failed with an error: {e}"
        print(f"Agent generated answer: {answer}")
        return answer


# --- Main Application Logic ---
def run_and_submit_all(profile: gr.OAuthProfile | None):
    space_id = os.getenv("SPACE_ID")
    if not profile:
        return "Please Login to Hugging Face with the button.", None
    username = profile.username
    print(f"User logged in: {username}")

    try:
        groq_api_key = os.getenv("GROQ_API_KEY")
        tavily_api_key = os.getenv("TAVILY_API_KEY")
        if not all([groq_api_key, tavily_api_key]):
            raise ValueError("An API key secret (GROQ or TAVILY) is missing.")
        agent = LangChainAgent(groq_api_key=groq_api_key, tavily_api_key=tavily_api_key)
    except Exception as e:
        return f"Error initializing agent: {e}", None

    questions_url = f"{DEFAULT_API_URL}/questions"
    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=20)
        response.raise_for_status()
        questions_data = response.json()
    except Exception as e:
        return f"Error fetching questions: {e}", None

    results_log, answers_payload = [], []
    for item in questions_data:
        task_id, question_text = item.get("task_id"), item.get("question")
        if not task_id or not question_text: continue
        submitted_answer = agent(question=question_text, task_id=task_id)
        answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
        results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    submission_data = {"username": username, "agent_code": agent_code, "answers": answers_payload}
    submit_url = f"{DEFAULT_API_URL}/submit"
    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
    try:
        response = requests.post(submit_url, json=submission_data, timeout=90) # Increased timeout
        response.raise_for_status()
        result_data = response.json()
        final_status = (f"Submission Successful!\nUser: {result_data.get('username')}\n"
                        f"Overall Score: {result_data.get('score', 'N/A')}% "
                        f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
                        f"Message: {result_data.get('message', 'No message received.')}")
        return final_status, pd.DataFrame(results_log)
    except Exception as e:
        return f"Submission Failed: {e}", pd.DataFrame(results_log)


# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# Advanced Agent Evaluation Runner (Search + Groq Audio)")
    gr.Markdown("This agent can search the web with Tavily and transcribe audio with Groq's Whisper.")
    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)
    run_button.click(fn=run_and_submit_all, outputs=[status_output, results_table])

if __name__ == "__main__":
    print("\n" + "-"*30 + " App Starting " + "-"*30)
    for key in ["GROQ_API_KEY", "TAVILY_API_KEY"]:
        print(f"✅ {key} secret is set." if os.getenv(key) else f"⚠️ WARNING: {key} secret is not set.")
    print("-"*(60 + len(" App Starting ")) + "\n")
    demo.launch(debug=True, share=False)