import os
import gradio as gr
import requests
import pandas as pd
from typing import List, Dict

from smolagents import CodeAgent, DuckDuckGoSearchTool, Tool

from wikipedia_searcher import WikipediaSearcher
from audio_transcriber import AudioTranscriptionTool
from image_analyzer import ImageAnalysisTool


class WikipediaSearchTool(Tool):
    name = "wikipedia_search"
    description = "Search Wikipedia for a given query."
    inputs = {
        "query": {
            "type": "string",
            "description": "The search query string"
        }
    }
    output_type = "string"

    def __init__(self):
        super().__init__()
        self.searcher = WikipediaSearcher()

    def forward(self, query: str) -> str:
        return self.searcher.search(query)


# Hugging Face Inference API wrapper for chat completion
class HFChatModel:
    def __init__(self, model_id: str):
        self.model_id = model_id
        self.api_url = f"https://api-inference.huggingface.co/models/{model_id}"
        self.headers = {"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}"}
        self.system_prompt = """
You are an agent solving the GAIA benchmark and you are required to provide exact answers.
Rules to follow:
1. Return only the exact requested answer: no explanation and no reasoning.
2. For yes/no questions, return exactly "Yes" or "No".
3. For dates, use the exact format requested.
4. For numbers, use the exact number, no other format.
5. For names, use the exact name as found in sources.
6. If the question has an associated file, download the file first using the task ID.
Examples of good responses:
- "42"
- "Yes"
- "October 5, 2001"
- "Buenos Aires"
Never include phrases like "the answer is..." or "Based on my research".
Only return the exact answer.
"""

    def generate(self, messages: List[Dict[str, str]]) -> str:
        # Prepend system prompt as first message
        all_messages = [{"role": "system", "content": self.system_prompt}] + messages

        payload = {
            "inputs": {
                "past_user_inputs": [],
                "generated_responses": [],
                "text": "\n".join(m["content"] for m in all_messages if m["role"] != "system")
            }
        }

        # Some HF chat models expect just a string prompt; adjust accordingly per your model's requirements

        response = requests.post(self.api_url, headers=self.headers, json=payload)
        if response.status_code == 200:
            output = response.json()
            # Output format depends on model; adjust as needed
            if isinstance(output, list) and len(output) > 0 and "generated_text" in output[0]:
                return output[0]["generated_text"].strip()
            elif isinstance(output, dict) and "generated_text" in output:
                return output["generated_text"].strip()
            else:
                # fallback to raw text
                return str(output).strip()
        else:
            raise RuntimeError(f"Hugging Face API error {response.status_code}: {response.text}")


class MyAgent:
    def __init__(self):
        self.model = HFChatModel(model_id="gpt-4o-mini")  # Or any HF chat model you want

        self.agent = CodeAgent(
            tools=[
                DuckDuckGoSearchTool(),
                WikipediaSearchTool(),
                AudioTranscriptionTool(),
                ImageAnalysisTool(),
            ],
            model=self,  # We'll route calls via __call__ below
        )

    def __call__(self, prompt: str) -> str:
        # Construct chat message for HF model
        messages = [{"role": "user", "content": prompt}]
        return self.model.generate(messages)


def run_and_submit_all(profile: gr.OAuthProfile | None):
    space_id = os.getenv("SPACE_ID")

    if profile:
        username = profile.username
    else:
        return "Please Login to Hugging Face with the button.", None

    DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
    questions_url = f"{DEFAULT_API_URL}/questions"
    submit_url = f"{DEFAULT_API_URL}/submit"

    try:
        agent = MyAgent()
    except Exception as e:
        return f"Error initializing agent: {e}", None

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    try:
        response = requests.get(questions_url, timeout=15)
        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 = item.get("task_id")
        if not task_id:
            continue
        try:
            answer = agent(item.get("question", ""))
            answers_payload.append({"task_id": task_id, "submitted_answer": answer})
            results_log.append({
                "Task ID": task_id,
                "Question": item.get("question", ""),
                "Submitted Answer": answer
            })
        except Exception as e:
            results_log.append({
                "Task ID": task_id,
                "Question": item.get("question", ""),
                "Submitted Answer": f"Error: {e}"
            })

    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload
    }

    try:
        response = requests.post(submit_url, json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        final_status = (
            f"Submission Successful!\n"
            f"User: {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)


with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner (HF API)")
    gr.LoginButton()
    run_btn = gr.Button("Run Evaluation & Submit All Answers")
    status_out = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    results_df = gr.DataFrame(label="Questions and Agent Answers")

    run_btn.click(fn=run_and_submit_all, outputs=[status_out, results_df])

if __name__ == "__main__":
    demo.launch(debug=True, share=False)