import os
import re
import gradio as gr
import requests
import pandas as pd
from langchain_openai import ChatOpenAI
from langchain_community.tools import DuckDuckGoSearchResults
from langchain_experimental.tools import PythonREPLTool
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage

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

SYSTEM_PROMPT = """You are a general AI assistant. Answer GAIA benchmark questions accurately.

Available tools:
- duckduckgo_search: returns search results WITH URLs. Each result has a link field.
- read_webpage: reads FULL text of a URL. ALWAYS call this after finding a relevant link in search results.
- Python_REPL: calculations and data analysis. ALWAYS use print() to output results.

MANDATORY research strategy:
1. Search with duckduckgo_search — look for a link to a Wikipedia page, database, or article.
2. Call read_webpage on the most relevant link from results — get full page content.
3. Extract the precise answer from page content.
4. If Wikipedia has an article: read it directly — e.g. read_webpage("https://en.wikipedia.org/wiki/Topic").

Special rules:
- Reversed/encoded text: decode it yourself, no tools needed.
- YouTube: search the video ID + key terms from the question.
- Attached files not available: search web for the answer instead.
- If question asks for IOC code: return the IOC code. If question asks for country name: return full name.

When done, output ONLY:
FINAL ANSWER: [your answer]

STRICT format rules (exact match):
- Numbers: digits only, no $, no commas, no units unless asked
- Strings: no surrounding quotes, no trailing punctuation, no articles (a/an/the)
- Lists: comma-separated, no spaces after commas
- Always give an answer — never output "No answer found\""""


@tool
def read_webpage(url: str) -> str:
    """Read the full text content of a webpage. Use after finding a relevant URL via search to get precise information."""
    try:
        headers = {"User-Agent": "Mozilla/5.0 (compatible; research-agent/1.0)"}
        resp = requests.get(url, headers=headers, timeout=15, allow_redirects=True)
        if resp.status_code != 200:
            return f"Could not fetch page: HTTP {resp.status_code}"
        text = re.sub(r"<[^>]+>", " ", resp.text)
        text = re.sub(r"\s+", " ", text).strip()
        return text[:6000]
    except Exception as e:
        return f"Error reading page: {e}"


class BasicAgent:
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
        self.tools = [
            DuckDuckGoSearchResults(num_results=5),
            read_webpage,
            PythonREPLTool(),
        ]
        self.tools_map = {t.name: t for t in self.tools}
        self.llm_with_tools = self.llm.bind_tools(self.tools, parallel_tool_calls=False)
        print("BasicAgent initialized with OpenAI (gpt-4o).")

    def __call__(self, question: str, task_id: str = "") -> str:
        full_question = f"[Task ID: {task_id}]\n\n{question}" if task_id else question
        print(f"Running agent on task {task_id}: {question[:80]}...")

        messages = [
            SystemMessage(content=SYSTEM_PROMPT),
            HumanMessage(content=full_question),
        ]

        last_response = None
        for iteration in range(10):
            response = self.llm_with_tools.invoke(messages)
            messages.append(response)
            last_response = response

            if not response.tool_calls:
                break

            for tool_call in response.tool_calls:
                tool_name = tool_call["name"]
                tool_args = tool_call["args"]
                tool_id = tool_call["id"]
                first_arg = str(list(tool_args.values())[0])[:60] if tool_args else ""
                print(f"  [{iteration+1}] Tool: {tool_name}({first_arg})")
                if tool_name in self.tools_map:
                    try:
                        result = self.tools_map[tool_name].invoke(tool_args)
                    except Exception as e:
                        result = f"Tool error: {e}"
                else:
                    result = f"Unknown tool: {tool_name}"
                messages.append(ToolMessage(content=str(result)[:3000], tool_call_id=tool_id))

        raw_answer = last_response.content if last_response else ""

        # If loop ended without FINAL ANSWER (hit limit or empty content), force one
        if "FINAL ANSWER:" not in raw_answer:
            messages.append(HumanMessage(
                content="Based on all information gathered above, give your FINAL ANSWER now. Format: FINAL ANSWER: [answer]"
            ))
            forced = self.llm.invoke(messages)
            raw_answer = forced.content

        if "FINAL ANSWER:" in raw_answer:
            answer = raw_answer.split("FINAL ANSWER:")[-1].strip()
        else:
            answer = raw_answer.strip()

        answer = self._clean_answer(answer)
        print(f"Answer for {task_id}: {answer[:100]}")
        return answer

    def _clean_answer(self, answer: str) -> str:
        # Strip surrounding quotes
        answer = answer.strip('"\'')
        # Strip trailing sentence punctuation
        answer = answer.rstrip('.')
        # Remove currency symbols
        answer = answer.replace('$', '').replace('€', '').replace('£', '')
        # Remove placeholder text
        if answer in ('[answer]', '[Answer]', '[YOUR ANSWER]', '[your answer]'):
            return ""
        # Normalize list spacing: "a, b, c" → "a,b,c"
        if ',' in answer and not any(c.isdigit() for c in answer.split(',')[0]):
            answer = ','.join(part.strip() for part in answer.split(','))
        # Strip surrounding brackets
        if answer.startswith('[') and answer.endswith(']') and answer.count('[') == 1:
            answer = answer[1:-1]
        return answer.strip()


def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the BasicAgent on them, submits all answers,
    and displays the results.
    """
    space_id = os.getenv("SPACE_ID")

    if profile:
        username = f"{profile.username}"
        print(f"User logged in: {username}")
    else:
        print("User not logged in.")
        return "Please Login to Hugging Face with the button.", 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:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(agent_code)

    # 2. Fetch Questions
    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
            print("Fetched questions list is empty.")
            return "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except requests.exceptions.RequestException as e:
        print(f"Error fetching questions: {e}")
        return f"Error fetching questions: {e}", None
    except requests.exceptions.JSONDecodeError as e:
        print(f"Error decoding JSON response from questions endpoint: {e}")
        print(f"Response text: {response.text[:500]}")
        return f"Error decoding server response for questions: {e}", None
    except Exception as e:
        print(f"An unexpected error occurred fetching questions: {e}")
        return f"An unexpected error occurred fetching questions: {e}", None

    # 3. Run Agent
    results_log = []
    answers_payload = []
    print(f"Running agent on {len(questions_data)} questions...")
    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:
            print(f"Skipping item with missing task_id or question: {item}")
            continue
        try:
            submitted_answer = agent(question_text, task_id)
            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:
            print(f"Error running agent on task {task_id}: {e}")
            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})

    if not answers_payload:
        print("Agent did not produce any answers to submit.")
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # 4. Submit
    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
    print(status_update)

    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
    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"User: {result_data.get('username')}\n"
            f"Overall Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
            f"Message: {result_data.get('message', 'No message received.')}"
        )
        print("Submission successful.")
        return final_status, pd.DataFrame(results_log)
    except requests.exceptions.HTTPError as e:
        error_detail = f"Server responded with status {e.response.status_code}."
        try:
            error_json = e.response.json()
            error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
        except requests.exceptions.JSONDecodeError:
            error_detail += f" Response: {e.response.text[:500]}"
        status_message = f"Submission Failed: {error_detail}"
        print(status_message)
        return status_message, pd.DataFrame(results_log)
    except requests.exceptions.Timeout:
        status_message = "Submission Failed: The request timed out."
        print(status_message)
        return status_message, pd.DataFrame(results_log)
    except requests.exceptions.RequestException as e:
        status_message = f"Submission Failed: Network error - {e}"
        print(status_message)
        return status_message, pd.DataFrame(results_log)
    except Exception as e:
        status_message = f"An unexpected error occurred during submission: {e}"
        print(status_message)
        return status_message, pd.DataFrame(results_log)


# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# Agent Evaluation Runner — Groq + Tool Binding")
    gr.Markdown(
        """
        **Instructions:**

        1. Log in to your Hugging Face account using the button below.
        2. Click 'Run Evaluation & Submit All Answers' to fetch questions, run the agent, and submit.

        **Agent:** Custom ReAct loop — OpenAI gpt-4o
        **Tools:** DuckDuckGo search, Python REPL, File fetcher (text + Excel)

        ---
        *Note: Running 20 questions takes several minutes.*
        """
    )

    gr.LoginButton()

    run_button = gr.Button("Run Evaluation & Submit All Answers")

    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)

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

if __name__ == "__main__":
    print("\n" + "-" * 30 + " App Starting " + "-" * 30)
    space_host_startup = os.getenv("SPACE_HOST")
    space_id_startup = os.getenv("SPACE_ID")

    if space_host_startup:
        print(f"✅ SPACE_HOST found: {space_host_startup}")
        print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")
    else:
        print("ℹ️  SPACE_HOST not found (running locally?).")

    if space_id_startup:
        print(f"✅ SPACE_ID found: {space_id_startup}")
        print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
    else:
        print("ℹ️  SPACE_ID not found (running locally?).")

    print("-" * (60 + len(" App Starting ")) + "\n")
    demo.launch(debug=True, share=False)