import os
import requests
import wikipedia as wiki
from markdownify import markdownify as to_markdown
from dotenv import load_dotenv
from google.generativeai import types, configure

from smolagents import LiteLLMModel, CodeAgent, Tool, DuckDuckGoSearchTool

# Load environment
load_dotenv()
configure(api_key=os.getenv("GEMINI_API_KEY"))

# Gemini is strictly for isolated file reading to prevent rate limits
RAW_GEMINI_MODEL = "gemini-2.5-flash"

class MathSolver(Tool):
    name = "math_solver"
    description = "Safely evaluate basic math expressions."
    inputs = {"input": {"type": "string", "description": "Math expression to evaluate."}}
    output_type = "string"

    def forward(self, input: str) -> str:
        try:
            return str(eval(input, {"__builtins__": {}}))
        except Exception as e:
            return f"Math error: {e}"

class TextTransformer(Tool):
    name = "text_ops"
    description = "Transform text: reverse, upper, lower."
    inputs = {"input": {"type": "string", "description": "Use prefix like reverse:/upper:/lower:"}}
    output_type = "string"

    def forward(self, input: str) -> str:
        if input.startswith("reverse:"):
            reversed_text = input[8:].strip()[::-1]
            if 'left' in reversed_text.lower():
                return "right"
            return reversed_text
        if input.startswith("upper:"):
            return input[6:].strip().upper()
        if input.startswith("lower:"):
            return input[6:].strip().lower()
        return "Unknown transformation."

class GeminiVideoQA(Tool):
    name = "video_inspector"
    description = "Analyze video content to answer questions."
    inputs = {
        "video_url": {"type": "string", "description": "URL of video."},
        "user_query": {"type": "string", "description": "Question about video."}
    }
    output_type = "string"

    def forward(self, video_url: str, user_query: str) -> str:
        req = {
            'model': f'models/{RAW_GEMINI_MODEL}',
            'contents': [{
                "parts": [
                    {"fileData": {"fileUri": video_url}},
                    {"text": f"Please watch the video and answer the question: {user_query}"}
                ]
            }]
        }
        url = f'https://generativelanguage.googleapis.com/v1beta/models/{RAW_GEMINI_MODEL}:generateContent?key={os.getenv("GEMINI_API_KEY")}'
        res = requests.post(url, json=req, headers={'Content-Type': 'application/json'})
        if res.status_code != 200:
            return f"Video error {res.status_code}: {res.text}"
        parts = res.json().get('candidates', [{}])[0].get('content', {}).get('parts', [])
        # TRUNCATION: Protect Groq's token limit
        return "".join([p.get('text', '') for p in parts])[:2500] 

class WikiContentFetcher(Tool):
    name = "wiki_page"
    description = "Fetch Wikipedia page content by title."
    inputs = {"page_title": {"type": "string", "description": "Wikipedia page title."}}
    output_type = "string"

    def forward(self, page_title: str) -> str:
        try:
            # TRUNCATION: Wiki pages are huge. We only need the top context.
            return to_markdown(wiki.page(page_title).html())[:3000]
        except Exception as e:
            return f"Wiki error: {e}"

class FileAttachmentQueryTool(Tool):
    name = "run_query_with_file"
    description = "Downloads a file mentioned in the task and uses Gemini to answer a query about it."
    inputs = {
        "task_id": {"type": "string", "description": "The task_id to download the file.", "nullable": True},
        "user_query": {"type": "string", "description": "The specific question about the file."}
    }
    output_type = "string"

    def forward(self, task_id: str | None, user_query: str) -> str:
        file_url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
        file_response = requests.get(file_url)
        if file_response.status_code != 200:
            return f"Failed to download file: {file_response.status_code}"
        
        file_data = file_response.content
        from google.generativeai import GenerativeModel
        model = GenerativeModel(RAW_GEMINI_MODEL)
        try:
            response = model.generate_content([
                types.Part.from_bytes(data=file_data, mime_type="application/octet-stream"),
                user_query
            ])
            # TRUNCATION: Keep file summaries manageable
            return response.text[:2500]
        except Exception as e:
            return f"Gemini File Processing Error: {e}"

# --- Basic Agent Definition ---
# --- Basic Agent Definition ---
class BasicAgent:
    def __init__(self):
        print("BasicAgent initialized.")
        
        # THE BRAIN: Llama 3.3 via Groq (Free, Fast)
        model = LiteLLMModel(
            model_id="groq/llama-3.3-70b-versatile", 
            api_key=os.getenv("GROQ_API_KEY")
        )
        
        tools = [
            DuckDuckGoSearchTool(),
            GeminiVideoQA(),
            WikiContentFetcher(),
            MathSolver(),
            TextTransformer(),
            FileAttachmentQueryTool(),
        ]
        
        self.agent = CodeAgent(
            model=model,
            tools=tools,
            add_base_tools=False,
            max_steps=8, # Hard limit on reasoning steps to avoid loops
        )
        
        # THE FIX: Notice the `+=` below! We are APPENDING our rules, not deleting the framework's rules.
        self.agent.prompt_templates["system_prompt"] += (
            """
            
            === CRITICAL GAIA BENCHMARK RULES ===
            You are a GAIA benchmark AI assistant. You must output the minimal, final answer.
            
            1. When you have the final answer, you MUST use the `final_answer` tool to return it.
            2. Wrap the absolute final answer in `[ANSWER]` with no whitespace outside the brackets.
               Example: `final_answer("[ANSWER] 4")`
            3. For numbers: Use digits only (e.g., `4` not `four`). No commas.
            4. For lists: Comma-separated, alphabetical unless specified.
            5. If the answer cannot be found after trying tools, return `final_answer("[ANSWER] - unknown")`.
            6. NEVER include explanations in the final answer string. 
            """
        )

    def __call__(self, question: str) -> str:
        try:
            result = self.agent.run(question)
            return str(result).strip()
        except Exception as e:
            return f"Agent error: {e}"