import os
import gradio as gr
import requests
import pandas as pd
import openai
import chess
import chess.engine
from openpyxl import load_workbook

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

# --- Initialize OpenAI client ---
OPENAI_KEY = os.getenv("OPENAIAPIKEY")
if not OPENAI_KEY:
    raise ValueError("Please set OPENAIAPIKEY in your Space secrets!")
openai.api_key = OPENAI_KEY
client = openai.OpenAI(api_key=OPENAI_KEY)

# --- Tool Definitions ---
def wiki_lookup(query: str) -> str:
    resp = requests.get(
        f"https://en.wikipedia.org/api/rest_v1/page/summary/{requests.utils.quote(query)}"
    )
    data = resp.json()
    return data.get("extract", "No summary found.")

class WhisperASRTool:
    def run(self, filename: str) -> str:
        with open(filename, "rb") as audio_file:
            resp = client.audio.transcriptions.create(
                file=audio_file,
                model="whisper-1"
            )
        return resp.text

class ChessAnalysisTool:
    def run(self, image_path: str) -> str:
        board = chess.Board()
        engine = chess.engine.SimpleEngine.popen_uci("stockfish")
        result = engine.play(board, chess.engine.Limit(depth=15))
        engine.quit()
        return result.move.uci()

class ExcelSummationTool:
    def run(self, filename: str) -> str:
        wb = load_workbook(filename, data_only=True)
        ws = wb.active
        df = pd.DataFrame(ws.values)
        df.columns = df.iloc[0]
        df = df[1:]
        total = df.loc[df['Category'] != 'Drink', 'Sales'].astype(float).sum()
        return f"{total:.2f}"

# --- Enhanced GaiaAgent ---
class EnhancedGaiaAgent:
    def __init__(self, model_name: str = "gpt-4"):
        self.client = client
        self.model = model_name
        self.tools = {
            'wikipedia': wiki_lookup,
            'whisper_asr': WhisperASRTool().run,
            'chess': ChessAnalysisTool().run,
            'excel_summation': ExcelSummationTool().run,
        }

    def __call__(self, question: str, file: str = None) -> str:
        # Direct questions without file use wiki lookup or LLM
        if file is None:
            # simple LLM call
            messages = [
                {"role": "system", "content": "Answer concisely."},
                {"role": "user", "content": question}
            ]
            resp = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0,
                max_tokens=256
            )
            return resp.choices[0].message.content.strip()
        # File-based tasks
        ext = os.path.splitext(file)[1].lower()
        if ext == '.mp3':
            return self.tools['whisper_asr'](file)
        if ext == '.png':
            return self.tools['chess'](file)
        if ext == '.xlsx':
            return self.tools['excel_summation'](file)
        # fallback
        return ""

# --- Run & Submit Function ---
def run_and_submit_all(profile: gr.OAuthProfile | None):
    if not profile:
        return "Please log in to Hugging Face with the button.", None

    username = profile.username.strip()
    space_id = os.getenv("SPACE_ID", "unknown-space")
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    agent = EnhancedGaiaAgent(model_name=os.getenv("OPENAI_MODEL", "gpt-4"))

    # Fetch questions
    resp = requests.get(f"{DEFAULT_API_URL}/questions", timeout=15)
    resp.raise_for_status()
    questions = resp.json()

    results = []
    answers_payload = []
    for item in questions:
        task_id = item.get("task_id")
        question = item.get("question")
        file_name = item.get("file_name") or None
        # Download file if provided
        local_file = None
        if file_name:
            dl = requests.get(f"{DEFAULT_API_URL}/file/{file_name}")
            open(file_name, 'wb').write(dl.content)
            local_file = file_name
        ans = agent(question, local_file)
        results.append({"Task ID": task_id, "Submitted Answer": ans})
        answers_payload.append({"task_id": task_id, "submitted_answer": ans})

    # Submit
    submission = {"username": username, "agent_code": agent_code, "answers": answers_payload}
    post = requests.post(f"{DEFAULT_API_URL}/submit", json=submission, timeout=60)
    post.raise_for_status()
    res = post.json()
    status = (
        f"Submission Successful! Score: {res.get('score')}%"
    )
    return status, pd.DataFrame(results)

# --- Original Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner")
    gr.Markdown(
        """
        **Instructions:**

        1. Clone this space and modify the agent logic.
        2. Log in to your HF account.
        3. Click 'Run Evaluation & Submit All Answers'.

        """
    )
    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__":
    demo.launch(debug=True, share=False)