import os
import gradio as gr
import requests
import inspect
import pandas as pd
import re
import json
from openai import OpenAI

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

class BasicAgent:
    def __init__(self):
        print("BasicAgent initialized.")
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("OPENAI_API_KEY environment variable not found.")
        
        self.client = OpenAI(api_key=api_key)
        self.model = "gpt-4o"  # Best for function calling
        
        # Define available tools
        self.tools = [
            {
                "type": "function",
                "function": {
                    "name": "web_search",
                    "description": "Search the web for current information on any topic",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Search query"
                            }
                        },
                        "required": ["query"]
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "wikipedia_search",
                    "description": "Get factual information from Wikipedia",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "topic": {
                                "type": "string", 
                                "description": "Topic to search on Wikipedia"
                            }
                        },
                        "required": ["topic"]
                    }
                }
            }
        ]
    
    def web_search(self, query: str) -> str:
        """Search using DuckDuckGo API"""
        try:
            url = "https://api.duckduckgo.com/"
            params = {'q': query, 'format': 'json', 'no_html': '1'}
            response = requests.get(url, params=params, timeout=10)
            data = response.json()
            
            result = ""
            if data.get('AbstractText'):
                result += f"Summary: {data['AbstractText']}\n"
            if data.get('RelatedTopics'):
                for topic in data['RelatedTopics'][:3]:
                    if isinstance(topic, dict) and 'Text' in topic:
                        result += f"- {topic['Text']}\n"
            
            return result or f"No results found for: {query}"
        except Exception as e:
            return f"Search error: {str(e)}"
    
    def wikipedia_search(self, topic: str) -> str:
        """Get Wikipedia summary"""
        try:
            url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{topic.replace(' ', '_')}"
            response = requests.get(url, timeout=10)
            if response.status_code == 200:
                data = response.json()
                if 'extract' in data:
                    return f"Wikipedia: {data['extract'][:800]}..."
            return f"No Wikipedia entry found for: {topic}"
        except Exception as e:
            return f"Wikipedia error: {str(e)}"
    
    def execute_function(self, name: str, arguments: dict) -> str:
        """Execute the requested function"""
        if name == "web_search":
            return self.web_search(arguments.get("query", ""))
        elif name == "wikipedia_search":
            return self.wikipedia_search(arguments.get("topic", ""))
        return f"Unknown function: {name}"
    
    def extract_boxed_answer(self, text: str) -> str:
        """Extract answer from \\boxed{} or \\text{}"""
        # Try boxed first
        boxed_pattern = r'\\boxed\{([^{}]*(?:\{[^{}]*\}[^{}]*)*)\}'
        matches = re.findall(boxed_pattern, text)
        if matches:
            return matches[-1].strip()
        
        # Try text
        text_pattern = r'\\text\{([^{}]*(?:\{[^{}]*\}[^{}]*)*)\}'
        matches = re.findall(text_pattern, text)
        if matches:
            return matches[-1].strip()
        
        # Fallback patterns
        fallback_patterns = [
            r'(?:final answer|answer):\s*(.+?)(?:\n|$)',
            r'(?:the answer is):\s*(.+?)(?:\n|$)',
        ]
        
        for pattern in fallback_patterns:
            matches = re.findall(pattern, text, re.IGNORECASE)
            if matches:
                return matches[-1].strip()
        
        # Clean up any LaTeX and return
        final_text = text.strip()
        final_text = re.sub(r'\\text\{([^{}]*)\}', r'\1', final_text)
        final_text = re.sub(r'\\boxed\{([^{}]*)\}', r'\1', final_text)
        return final_text.strip()
    
    def __call__(self, question: str) -> str:
        print(f"Processing question: {question[:50]}...")
        
        try:
            system_prompt = """You are an expert problem solver with access to search tools. 

For questions requiring current info, facts, or research, use the available tools first.
Think step by step, then provide your final answer.

CRITICAL: End with your final answer in this format: \\boxed{your_answer}

Examples:
- "Based on my search, the answer is 42. \\boxed{42}"
- "According to Wikipedia, it's Paris. \\boxed{Paris}"

Only use \\boxed{} for your final answer."""

            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": question}
            ]
            
            # Allow up to 3 tool calls to prevent infinite loops
            for iteration in range(3):
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    tools=self.tools,
                    tool_choice="auto",
                    max_tokens=1500,
                    temperature=0.1
                )
                
                response_message = response.choices[0].message
                messages.append({
                    "role": "assistant",
                    "content": response_message.content,
                    "tool_calls": response_message.tool_calls
                })
                
                # If no tool calls, we're done
                if not response_message.tool_calls:
                    final_response = response_message.content
                    break
                
                # Execute tool calls
                for tool_call in response_message.tool_calls:
                    function_name = tool_call.function.name
                    function_args = json.loads(tool_call.function.arguments)
                    
                    print(f"Calling {function_name} with {function_args}")
                    result = self.execute_function(function_name, function_args)
                    
                    messages.append({
                        "tool_call_id": tool_call.id,
                        "role": "tool",
                        "name": function_name,
                        "content": result
                    })
            
            final_answer = self.extract_boxed_answer(final_response)
            print(f"Final answer: {final_answer}")
            return final_answer
            
        except Exception as e:
            return f"Error: {str(e)}"

# [Rest of the code remains the same - just the agent class changed]
def run_and_submit_all(profile: gr.OAuthProfile | None):
    """Fetches questions, runs agent, submits answers"""
    space_id = os.getenv("SPACE_ID")
    
    if profile:
        username = f"{profile.username}"
        print(f"User logged in: {username}")
    else:
        return "Please Login to Hugging Face.", None

    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    # 1. Instantiate Agent
    try:
        agent = BasicAgent()
    except Exception as e:
        return f"Error initializing agent: {e}", None
    
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    # 2. Fetch Questions
    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        print(f"Fetched {len(questions_data)} questions.")
    except Exception as e:
        return f"Error fetching questions: {e}", None

    # 3. Run Agent
    results_log = []
    answers_payload = []
    
    for item in questions_data:
        task_id = item.get("task_id")
        question_text = item.get("question")
        if not task_id or question_text is None:
            continue
            
        try:
            submitted_answer = agent(question_text)
            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
        except Exception as e:
            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"ERROR: {e}"})

    # 4. Submit
    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
    
    try:
        response = requests.post(submit_url, json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        
        final_status = (
            f"Submission Successful!\n"
            f"Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')})\n"
            f"Message: {result_data.get('message', 'No message')}"
        )
        
        return final_status, pd.DataFrame(results_log)
    except Exception as e:
        return f"Submission failed: {e}", pd.DataFrame(results_log)

# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# AI Agent with Tool Calling")
    gr.Markdown("""
    **Features:**
    - Web search via DuckDuckGo
    - Wikipedia lookup
    - Smart tool selection by GPT-4o
    - Robust answer extraction
    
    **Setup:** Add OPENAI_API_KEY to repository secrets
    """)

    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit")
    status_output = gr.Textbox(label="Status", lines=5, interactive=False)
    results_table = gr.DataFrame(label="Results", wrap=True)

    run_button.click(fn=run_and_submit_all, outputs=[status_output, results_table])

if __name__ == "__main__":
    print("Starting AI Agent with Tool Calling...")
    demo.launch(debug=True, share=False)