import os
import base64
import requests
from openai import OpenAI
from duckduckgo_search import DDGS

# You can expand with image/audio/etc tools as needed

class BasicAgent:
    def __init__(self):
        self.llm = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        print("BasicAgent initialized.")

    def web_search(self, query: str, max_results: int = 5) -> str:
        """Search the web using DuckDuckGo for current information."""
        try:
            with DDGS() as ddgs:
                results = list(ddgs.text(query, max_results=max_results))
            if not results:
                return f"No results found for query: {query}"
            formatted_results = f"Web search results for '{query}':\n\n"
            for i, result in enumerate(results, 1):
                title = result.get('title', 'No title')
                body = result.get('body', 'No description')
                href = result.get('href', 'No URL')
                formatted_results += f"{i}. {title}\n"
                formatted_results += f"   URL: {href}\n"
                formatted_results += f"   Description: {body}\n\n"
            return formatted_results
        except Exception as e:
            return f"Error performing web search: {str(e)}"

    def fetch_file(self, task_id):
        DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
        try:
            url = f"{DEFAULT_API_URL}/files/{task_id}"
            r = requests.get(url, timeout=10)
            r.raise_for_status()
            return url, r.content, r.headers.get("Content-Type", "")
        except:
            return None, None, None

    def describe_image(self, img_path: str) -> str:
        # This is optional: image-based chess move recognition
        try:
            r = requests.get(img_path, timeout=10)
            image_bytes = r.content
            image_base64 = base64.b64encode(image_bytes).decode("utf-8")
            prompt = (
                "You're a chess assistant. Answer only with the best move in algebraic notation (e.g., Qd1#)."
            )
            # In this template, you'd need your LLM to accept vision input.
            # If not, just return an error or a stub.
            return "[Image analysis not implemented in this agent.]"
        except Exception as e:
            return f"Error extracting text: {str(e)}"

    def __call__(self, question: str, task_id: str = None) -> str:
        # Step 1: Always web search and build prompt
        search_snippet = self.web_search(question)
        # Step 2: Compose the mandated prompt exactly as specified
        full_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]. 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 comma to write your number neither use units such as $ or percent sign unless specified otherwise. "
            "If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. "
            "If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.\n\n"
            f"Here are web search results and the question:\n{search_snippet}\n\nQuestion: {question}"
        )

        # Step 3: Get answer from OpenAI, temperature 0 for max determinism
        response = self.llm.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": full_prompt},
            ],
            temperature=0.0,
            max_tokens=512,
        )
        answer = response.choices[0].message.content.strip()
        # You may want to extract just the "FINAL ANSWER: ..." line, or you can return as is
        return answer