import os
import io
import sys
import json
import warnings
import traceback
from typing import Optional, Dict, Any, List
import requests
import pandas as pd
from dotenv import load_dotenv

warnings.filterwarnings("ignore", category=RuntimeWarning)
load_dotenv()

# LangChain & LangGraph
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_groq import ChatGroq
from langgraph.prebuilt import create_react_agent

try:
    from ddgs import DDGS
except ImportError:
    from duckduckgo_search import DDGS

# Scoring base URL
SCORING_API_BASE = "https://agents-course-unit4-scoring.hf.space"
DOWNLOAD_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "downloads")
os.makedirs(DOWNLOAD_DIR, exist_ok=True)


# ---------------------------------------------------------
# Tool Definitions
# ---------------------------------------------------------

@tool
def web_search(query: str = "", **kwargs) -> str:
    """Search the web for up-to-date facts, Wikipedia entries, sports statistics, paper details, and general knowledge.
    Always pass the search query string in the 'query' argument.
    """
    search_query = query.strip() if isinstance(query, str) else ""
    if not search_query and kwargs:
        search_query = str(kwargs.get("search_query") or kwargs.get("input") or next(iter(kwargs.values()), "")).strip()
    if not search_query:
        return "No search query provided."
    try:
        with DDGS() as ddgs:
            results = list(ddgs.text(search_query, max_results=2))
            if not results:
                return f"No results found for query: '{search_query}'."
            formatted = []
            for i, r in enumerate(results, 1):
                title = r.get("title", "")
                snippet = r.get("body", "")[:250]
                formatted.append(f"[{i}] {title}: {snippet}")
            return "\n".join(formatted)
    except Exception as e:
        return f"Web search encountered an error: {e}"


def _run_python_code(code: str) -> str:
    """Internal runner for Python code execution."""
    old_stdout = sys.stdout
    redirected_output = io.StringIO()
    sys.stdout = redirected_output

    exec_globals = {
        "__builtins__": __builtins__,
        "pd": pd,
        "pandas": pd,
        "os": os,
        "json": json,
        "re": __import__("re"),
        "math": __import__("math"),
        "datetime": __import__("datetime"),
        "collections": __import__("collections"),
        "itertools": __import__("itertools"),
    }
    exec_locals = {}

    try:
        # Strip markdown python fences if LLM included them
        clean_code = str(code).strip()
        if clean_code.startswith("```"):
            lines = clean_code.split("\n")
            if lines[0].startswith("```"):
                lines = lines[1:]
            if lines and lines[-1].strip() == "```":
                lines = lines[:-1]
            clean_code = "\n".join(lines)

        # Check if the last statement is an expression to print its value
        parsed = clean_code.split("\n")
        last_line = parsed[-1].strip() if parsed else ""
        
        # Execute the code block
        exec(clean_code, exec_globals, exec_locals)
        output = redirected_output.getvalue()

        # If stdout is empty and last expression evaluated to something in locals
        if not output and last_line and last_line in exec_locals:
            output = str(exec_locals[last_line])

        clean_output = output.strip() if output.strip() else "Code executed successfully with no output."
        return clean_output[:800]
    except Exception:
        return f"Error executing Python code:\n{traceback.format_exc()[:500]}"
    finally:
        sys.stdout = old_stdout


@tool
def code_interpreter(code: str = "", **kwargs) -> str:
    """Execute Python code in an environment with pandas, math, re, json, datetime, openpyxl, etc.
    Always pass the code string in the 'code' argument.
    Use this tool to calculate numbers, simulate algorithms, process text/strings, inspect dataframes, or run scripts.
    """
    script = code
    if not script and kwargs:
        script = kwargs.get("script") or kwargs.get("command") or kwargs.get("input") or str(next(iter(kwargs.values()), ""))
    return _run_python_code(script)


@tool
def python(code: str = "", **kwargs) -> str:
    """Execute Python code in an environment with pandas, math, re, json, datetime, openpyxl, etc."""
    return code_interpreter.invoke({"code": code, **kwargs})


@tool
def python_interpreter(code: str = "", **kwargs) -> str:
    """Alias for code execution tool."""
    return code_interpreter.invoke({"code": code, **kwargs})


@tool
def download_task_file(task_id: str, file_name: str = "") -> str:
    """Download an attached file for a given task_id from the scoring server and inspect its contents.
    Supports python (.py), spreadsheets (.xlsx, .csv), and text files.
    """
    if not task_id:
        return "Error: task_id is required to download a file."

    task_dir = os.path.join(DOWNLOAD_DIR, task_id)
    os.makedirs(task_dir, exist_ok=True)
    target_filename = file_name if file_name else f"{task_id}.bin"
    target_path = os.path.join(task_dir, target_filename)

    url = f"{SCORING_API_BASE}/files/{task_id}"
    try:
        resp = requests.get(url, timeout=30)
        resp.raise_for_status()
        with open(target_path, "wb") as f:
            f.write(resp.content)
    except requests.exceptions.HTTPError as e:
        if resp.status_code == 404:
            return (
                f"Note: Server returned 404 for file '{target_filename}' (task_id: {task_id}). "
                "Answer using web search, general knowledge, or reasoning without the file."
            )
        return f"Failed to download file: {e}"
    except Exception as e:
        return f"Failed to download file: {e}"

    # Auto-inspect based on extension
    lower_name = target_path.lower()
    if lower_name.endswith((".py", ".txt", ".json", ".csv")):
        try:
            with open(target_path, "r", encoding="utf-8", errors="replace") as f:
                content = f.read()
            return f"Downloaded to '{target_path}'.\n\nFile contents:\n```\n{content[:800]}\n```"
        except Exception as e:
            return f"Downloaded to '{target_path}', but failed to read text: {e}"

    elif lower_name.endswith((".xlsx", ".xls")):
        try:
            excel_data = pd.ExcelFile(target_path)
            summary = [f"Downloaded to '{target_path}'. Sheets: {excel_data.sheet_names}"]
            for sheet in excel_data.sheet_names[:2]:
                df = pd.read_excel(target_path, sheet_name=sheet)
                summary.append(f"\nSheet {sheet} (Shape: {df.shape}), Columns: {list(df.columns[:10])}")
                summary.append("Head:\n" + df.head(2).to_string())
            return "\n".join(summary)[:800]
        except Exception as e:
            return f"Downloaded Excel to '{target_path}', but error parsing sheets: {e}"

    return f"File downloaded to '{target_path}' ({len(resp.content)} bytes)."


@tool
def read_local_file(file_path: str) -> str:
    """Read a local file from disk. Use this when inspecting downloaded task files or scripts."""
    try:
        if not os.path.exists(file_path):
            return f"File not found: {file_path}"
        with open(file_path, "r", encoding="utf-8", errors="replace") as f:
            content = f.read()
        return content[:800]
    except Exception as e:
        return f"Error reading file {file_path}: {e}"


# ---------------------------------------------------------
# System Prompt & Answer Cleaning
# ---------------------------------------------------------

GAIA_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].

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.

You have access to the following tools:
1. `web_search`: Search DuckDuckGo for factual information, Wikipedia articles, sports statistics, dates, or names. Pass arguments: {"query": "<search terms>"}
2. `code_interpreter`: Run Python code to perform arithmetic, data calculations, string manipulations, or execute scripts. Pass arguments: {"code": "<valid python code>"}
3. `download_task_file`: Download and inspect any attached file associated with a task_id (.py, .xlsx, .csv, etc.).
4. `read_local_file`: Read a local text file.

RULES:
- STRICT BUDGET: Aim to answer within 1 to 2 tool calls. Be decisive.
- When calling ANY tool, format arguments strictly as a valid JSON dictionary: {"code": "<string>"} or {"query": "<string>"}. Never output raw unquoted code.
- ALWAYS terminate your final response with: FINAL ANSWER: [YOUR FINAL ANSWER]
"""

def sanitize_answer(raw_answer: str) -> str:
    """Clean the raw LLM output to conform strictly to GAIA exact match expectations."""
    if not raw_answer:
        return ""

    import re
    ans = raw_answer.strip()

    # If the text contains 'FINAL ANSWER:', extract everything after the last occurrence
    upper_ans = ans.upper()
    if "FINAL ANSWER:" in upper_ans:
        idx = upper_ans.rfind("FINAL ANSWER:")
        ans = ans[idx + len("FINAL ANSWER:"):].strip()

    # Strip markdown code blocks if entire answer is wrapped
    if ans.startswith("```") and ans.endswith("```"):
        lines = ans.split("\n")
        ans = "\n".join(lines[1:-1]).strip()

    # If wrapped in square brackets like [42] or [Paris], remove them
    if ans.startswith("[") and ans.endswith("]"):
        ans = ans[1:-1].strip()

    # Regex to strip common prefixes (case insensitive)
    pattern = r"^(?:the\s+)?(?:final\s+)?answer\s*(?:is)?\s*[:\-]?\s*|^result\s*[:\-]?\s*|^output\s*[:\-]?\s*"
    ans = re.sub(pattern, "", ans, flags=re.IGNORECASE).strip()

    # If multiline, check if last non-empty line has the answer
    lines = [line.strip() for line in ans.split("\n") if line.strip()]
    if len(lines) > 1:
        last_line = lines[-1]
        last_line_cleaned = re.sub(pattern, "", last_line, flags=re.IGNORECASE).strip()
        if re.search(r"(?:answer|result|output)\s*[:\-]", last_line, flags=re.IGNORECASE) or len(last_line_cleaned) < len(ans) * 0.5:
            ans = last_line_cleaned

    # Strip surrounding bolding, quotes, backticks, brackets
    ans = ans.strip("`*\"' []")
    if (ans.startswith("**") and ans.endswith("**")) or (ans.startswith("__") and ans.endswith("__")):
        ans = ans[2:-2].strip()

    if ans.startswith("[") and ans.endswith("]"):
        ans = ans[1:-1].strip()

    # Strip trailing period if not a decimal number
    if ans.endswith(".") and not re.search(r"\d\.\d+$", ans):
        ans = ans[:-1].strip()

    return ans.strip("`*\"' []")


DIRECT_PROMPT = (
    "You are an expert providing the direct final answer to a benchmark question. "
    "Output ONLY the concise exact-match final answer string. "
    "Do NOT mention tools, do NOT attempt to call functions, do NOT write code, do NOT explain. "
    "Provide ONLY the final answer text or number."
)


# ---------------------------------------------------------
# LangGraph Agent Class
# ---------------------------------------------------------

class LangGraphGAIAAgent:
    """Autonomous LangGraph Agent powered by Groq's openai/gpt-oss-20b."""

    def __init__(self, groq_api_key: Optional[str] = None, model_name: str = "openai/gpt-oss-120b"):
        self.api_key = groq_api_key or os.getenv("GROQ_API_KEY")
        if not self.api_key:
            print("WARNING: GROQ_API_KEY is not set. Please provide it in the UI or environment.")
        
        self.model_name = model_name
        self.fallback_model = "qwen/qwen3.8-27b" if model_name != "qwen/qwen3.8-27b" else "openai/gpt-oss-120b"
        self.tools = [web_search, code_interpreter, python, python_interpreter, download_task_file, read_local_file]
        self._build_agent()

    def _build_agent(self):
        try:
            self.llm = ChatGroq(
                model=self.model_name,
                groq_api_key=self.api_key or "placeholder_key",
                temperature=0.0,
                max_retries=2
            )
            self.app = create_react_agent(
                model=self.llm,
                tools=self.tools,
                prompt=GAIA_SYSTEM_PROMPT
            )
        except Exception as e:
            print(f"Error initializing primary model {self.model_name}: {e}. Falling back to {self.fallback_model}...")
            self.llm = ChatGroq(
                model=self.fallback_model,
                groq_api_key=self.api_key or "placeholder_key",
                temperature=0.0,
                max_retries=2
            )
            self.app = create_react_agent(
                model=self.llm,
                tools=self.tools,
                prompt=GAIA_SYSTEM_PROMPT
            )

    def __call__(self, question: str, task_id: Optional[str] = None, file_name: Optional[str] = None) -> str:
        """Process a single question and return the sanitized exact answer."""
        if not self.api_key:
            return "ERROR: GROQ_API_KEY is required to run the agent."

        import time

        # Enhance question with file and task metadata if available
        augmented_prompt = question
        if task_id:
            augmented_prompt += f"\n\n[Context: task_id='{task_id}']"
        if file_name:
            augmented_prompt += f"\n[Context: attached file='{file_name}'. You can download and inspect it using the download_task_file tool with task_id='{task_id}']"

        messages = [
            SystemMessage(content=GAIA_SYSTEM_PROMPT),
            HumanMessage(content=augmented_prompt)
        ]

        # Attempt invocation with primary model
        try:
            result = self.app.invoke(
                {"messages": messages},
                config={"recursion_limit": 6}
            )
            last_message = result["messages"][-1]
            raw_answer = str(last_message.content) if last_message else ""

            # If agent stalled or hit LangGraph step limit ("need more steps" / "sorry"), ask directly
            lower_ans = raw_answer.lower()
            if (
                not raw_answer.strip()
                or "need more steps" in lower_ans
                or lower_ans.startswith("sorry")
                or "cannot answer" in lower_ans
                or "i cannot" in lower_ans
                or lower_ans.startswith("download_task_file")
                or lower_ans.startswith("web_search")
                or lower_ans.startswith("python")
            ):
                print(f"Agent stalled or reached step limit ('{raw_answer[:45]}...'). Querying model directly...", flush=True)
                direct_resp = self.llm.invoke([
                    SystemMessage(content=DIRECT_PROMPT),
                    HumanMessage(content=f"Question: {question}\n\nWhat is the concise exact-match final answer?")
                ])
                ans = sanitize_answer(direct_resp.content)
                if not ans and self.fallback_model:
                    try:
                        fallback_llm = ChatGroq(model=self.fallback_model, groq_api_key=self.api_key, temperature=0.0)
                        fb_resp = fallback_llm.invoke([
                            SystemMessage(content=DIRECT_PROMPT),
                            HumanMessage(content=f"Question: {question}\n\nWhat is the concise exact-match final answer?")
                        ])
                        return sanitize_answer(fb_resp.content)
                    except Exception:
                        pass
                return ans

            return sanitize_answer(raw_answer)

        except Exception as e:
            error_msg = str(e)
            print(f"Agent execution error with {self.model_name}: {error_msg}", flush=True)

            # 1. If rate limit (413 or 429) or token quota encountered:
            # Multi-turn messages overflowed token bucket; ask directly with a single tiny prompt (~100 tokens).
            if any(k in error_msg.lower() for k in ["413", "429", "rate_limit", "itpm", "tokens"]):
                print("Token quota or rate limit hit. Falling back to lightweight direct answer...", flush=True)
                time.sleep(4)
                try:
                    direct_resp = self.llm.invoke([
                        SystemMessage(content=DIRECT_PROMPT),
                        HumanMessage(content=f"Question: {question}\n\nWhat is the concise exact-match final answer?")
                    ])
                    return sanitize_answer(direct_resp.content)
                except Exception as retry_e:
                    print(f"Direct query hit limit too ({retry_e}). Waiting 10s for bucket refill...", flush=True)
                    time.sleep(10)
                    try:
                        direct_resp = self.llm.invoke([
                            SystemMessage(content=DIRECT_PROMPT),
                            HumanMessage(content=f"Question: {question}\n\nWhat is the concise exact-match final answer?")
                        ])
                        return sanitize_answer(direct_resp.content)
                    except Exception as final_e:
                        return f"AGENT ERROR: {final_e}"

            # 2. If recursion limit hit:
            if "recursion" in error_msg.lower():
                print("Recursion limit reached. Requesting direct answer from LLM...", flush=True)
                try:
                    recovery_resp = self.llm.invoke([
                        SystemMessage(content=DIRECT_PROMPT),
                        HumanMessage(content=f"Question: {question}\n\nWhat is the concise exact-match final answer?")
                    ])
                    return sanitize_answer(recovery_resp.content)
                except Exception as rec_err:
                    return f"AGENT ERROR: {rec_err}"

            # 3. Fallback to fallback_model if primary failed
            if self.model_name != self.fallback_model:
                print(f"Triggering automatic fallback to {self.fallback_model}...", flush=True)
                try:
                    fallback_llm = ChatGroq(
                        model=self.fallback_model,
                        groq_api_key=self.api_key,
                        temperature=0.0,
                        max_retries=2
                    )
                    fallback_app = create_react_agent(
                        model=fallback_llm,
                        tools=self.tools,
                        prompt=GAIA_SYSTEM_PROMPT
                    )
                    result = fallback_app.invoke(
                        {"messages": messages},
                        config={"recursion_limit": 6}
                    )
                    last_msg = result["messages"][-1]
                    ans = str(last_msg.content)
                    if "need more steps" in ans.lower() or ans.lower().startswith("sorry") or ans.lower().startswith("download_task_file"):
                        direct = fallback_llm.invoke([
                            SystemMessage(content=DIRECT_PROMPT),
                            HumanMessage(content=f"Question: {question}\n\nWhat is the concise exact-match final answer?")
                        ])
                        return sanitize_answer(direct.content)
                    return sanitize_answer(ans)
                except Exception as fallback_err:
                    print(f"Fallback model failed: {fallback_err}", flush=True)
                    error_msg = str(fallback_err)

            # 4. Final attempt: direct query
            try:
                final_direct = self.llm.invoke([
                    SystemMessage(content=DIRECT_PROMPT),
                    HumanMessage(content=f"Question: {question}\n\nWhat is the concise exact-match final answer?")
                ])
                return sanitize_answer(final_direct.content)
            except Exception:
                return f"AGENT ERROR: {error_msg}"