import os
import requests
from smolagents import CodeAgent, LiteLLMModel, DuckDuckGoSearchTool, tool
# ==========================================
# 1. DOWNLOAD TOOL (For GAIA Datasets)
# ==========================================
@tool
def download_file(url: str, file_name: str) -> str:
    """Downloads a file from a given URL and saves it locally.
    Args:
        url: The direct URL to the file.
        file_name: The local path and filename to save it as.
    """
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()
        directory = os.path.dirname(file_name)
        if directory:
            os.makedirs(directory, exist_ok=True)
        with open(file_name, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return f"File successfully downloaded and saved as: {file_name}"
    except Exception as e:
        return f"File System Error: {str(e)}"
# ==========================================
# 2. MODEL & AGENT INITIALIZATION (UPDATED)
# ==========================================
# LiteLLMModel மூலம் Hugging Face மாடலை இணைக்கிறோம்
model = LiteLLMModel(
    model_id="huggingface/Qwen/Qwen2.5-Coder-32B-Instruct",
    api_key=os.environ.get("HF_TOKEN")
)
agent = CodeAgent(
    tools=[DuckDuckGoSearchTool(), download_file],
    model=model,
    additional_authorized_imports=["requests", "pandas", "numpy", "json", "os", "math"]
)

# ==========================================
# 3. LANGGRAPH WRAPPER (To keep app.py happy)
# ==========================================
# உங்களின் app.py-ல் எந்த மாற்றமும் செய்யாமல் இருக்க, 
# LangGraph போலவே செயல்படும் ஒரு போலி கிளாஸை (Wrapper) உருவாக்குகிறோம்.
class GaiaAppWrapper:
    def invoke(self, state: dict) -> dict:
        # Extract question from LangGraph state format
        question = state["messages"][0].content
        file_name = state.get("file_name")
        
        # Strict Prompting
        prompt = (
            "You are a GAIA AI Assistant. You must answer the following question. "
            "For complex logic, math, data analysis, or current events, write Python code or use web search. "
            "CRITICAL: Output ONLY the exact entity, number, or short phrase requested as your final answer. "
            "Do not include any conversational filler like 'The answer is'.\n\n"
            f"Question: {question}"
        )
        
        if file_name:
            prompt += f"\n\nContext: The user provided a file at '{file_name}'. You can analyze it using pandas."
        
        try:
            # smolagents-ஐ ரன் செய்கிறோம்
            result = agent.run(prompt)
            return {"final_answer": str(result)}
        except Exception as e:
            return {"final_answer": f"Agent Error: {str(e)}"}

# app.py இதைத்தான் Import செய்யும்
gaia_app = GaiaAppWrapper()