import os
import re
from pathlib import Path
import tempfile

import openai
import pandas as pd
from tabulate import tabulate
from smolagents.tools import Tool, PipelineTool

# Load your OpenAI API key
openai.api_key = os.getenv("OPENAI_API_KEY")

SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template:
FINAL ANSWER: [YOUR FINAL ANSWER].

If you have enough confidence (≥70%) in your internal knowledge, answer directly without calling any other tool.
Otherwise, decide which tool (wiki, web search, python, audio, excel) is best and use it.

YOUR FINAL ANSWER should be:
- a number OR as few words as possible OR a comma separated list of numbers and/or strings.
- If you are asked for a number, don't use commas, symbols or units (e.g. %, $, km) unless explicitly asked.
- If you are asked for a string, don't use articles ("a", "the"), abbreviations (e.g. "NYC"), or extra words; write digits in plain text unless specified otherwise.
- If you are asked for a comma separated list, apply the above rules to each element.
"""

class AnswerTool(Tool):
    name = "answer_tool"
    description = "Answer GAIA Level 1 questions in exact-match format."
    inputs = {"question": {"type": "string", "description": "The GAIA question text."}}
    output_type = "string"

    def forward(self, question: str) -> str:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": question},
        ]
        resp = openai.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            temperature=0.0,
            max_tokens=64,
        )
        text = resp.choices[0].message.content.strip()
        if "FINAL ANSWER:" in text:
            return text.split("FINAL ANSWER:")[-1].strip()
        return text

class SpeechToTextTool(PipelineTool):
    """Transcribes audio files using OpenAI Whisper."""
    name = "transcriber"
    description = "Transcribes a local audio file to text using Whisper API."
    inputs = {"audio": {"type": "string", "description": "Path to audio file."}}
    output_type = "string"
    default_checkpoint = "openai/whisper-1"

    def __call__(self, audio: str) -> str:
        return self._transcribe(audio)

    @staticmethod
    def _transcribe(audio_path: str) -> str:
        path = Path(audio_path).expanduser().resolve()
        if not path.is_file():
            raise FileNotFoundError(f"Audio file not found: {path}")
        with path.open("rb") as fp:
            response = openai.audio.transcriptions.create(
                file=fp,
                model="whisper-1",
                response_format="text"
            )
        return response

class ExcelToTextTool(Tool):
    """Renders an Excel worksheet as Markdown text."""
    name = "excel_to_text"
    description = "Read an Excel file and return a Markdown table."
    inputs = {
        "excel_path": {"type": "string", "description": "Path to Excel file."},
        "sheet_name": {"type": "string", "description": "Sheet name or index as string.", "nullable": True},
    }
    output_type = "string"

    def forward(self, excel_path: str, sheet_name: str = None) -> str:
        path = Path(excel_path).expanduser().resolve()
        if not path.exists():
            return f"Error: Excel file not found at {path}"
        sheet = 0 if not sheet_name or sheet_name == "" else (
            int(sheet_name) if sheet_name.isdigit() else sheet_name
        )
        df = pd.read_excel(path, sheet_name=sheet)
        if hasattr(df, "to_markdown"):
            return df.to_markdown(index=False)
        return tabulate(df, headers="keys", tablefmt="github", showindex=False)