import os
import re
import operator
from typing import Literal

import gradio as gr
import pandas as pd
import requests
from dotenv import load_dotenv

from langchain_anthropic import ChatAnthropic
from langchain.messages import SystemMessage, HumanMessage, AnyMessage, ToolMessage
from typing_extensions import TypedDict, Annotated
from langgraph.graph import StateGraph, START, END
from tools.tools import get_tools

load_dotenv()

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

# --- Model ---
model = ChatAnthropic(model="claude-haiku-4-5-20251001", temperature=0)

# --- Tools ---
tools = get_tools()
tools_by_name = {tool.name: tool for tool in tools}
model_with_tools = model.bind_tools(tools)

# --- State ---
class State(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]
    llm_calls: int

# --- System Prompt ---
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: [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. "
    "If search results don't contain the answer, use visit_webpage to read the full page. Never guess. "
    "If a file path is mentioned in the question, use read_file_tool to read its contents. "
    "For Python files, read them first, then use Python_REPL to execute the code. "
    "For Excel/CSV files, use Python_REPL with pandas to read and analyze them."
)

# --- Nodes ---
def llm_call(state: dict):
    try:
        response = model_with_tools.invoke(
            [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
        )
    except Exception as e:
        response = model.invoke(
            [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
        )
    return {
        "messages": [response],
        "llm_calls": state.get("llm_calls", 0) + 1,
    }

def tool_node(state: dict):
    result = []
    for tc in state["messages"][-1].tool_calls:
        try:
            obs = tools_by_name[tc["name"]].invoke(tc["args"])
        except Exception as e:
            obs = f"Error: {e}"
        result.append(ToolMessage(content=str(obs), tool_call_id=tc["id"]))
    return {"messages": result}

def should_continue(state: State) -> Literal["tool_node", "__end__"]:
    if state.get("llm_calls", 0) >= MAX_LLM_CALLS:
        return END
    if state["messages"][-1].tool_calls:
        return "tool_node"
    return END

# --- Build & Compile Agent ---
g = StateGraph(State)
g.add_node("llm_call", llm_call)
g.add_node("tool_node", tool_node)
g.add_edge(START, "llm_call")
g.add_conditional_edges("llm_call", should_continue, ["tool_node", END])
g.add_edge("tool_node", "llm_call")
agent = g.compile()

# --- Output Parser ---
def extract_answer(text) -> str:
    if isinstance(text, list):
        text = " ".join([block.get("text", "") for block in text if isinstance(block, dict) and block.get("type") == "text"])
    if not isinstance(text, str):
        text = str(text)
    match = re.search(r'\[YOUR FINAL ANSWER\]:?\s*(.*)', text, re.DOTALL)
    return match.group(1).strip() if match else text.strip()

# --- Agent Wrapper ---
def run_agent(question: str) -> str:
    try:
        result = agent.invoke(
            {"messages": [HumanMessage(content=question)], "llm_calls": 0},
            {"recursion_limit": 25},
        )
        ai_message = result["messages"][-1]
        return extract_answer(ai_message.content)
    except Exception as e:
        print(f"Agent error: {e}")
        return f"AGENT ERROR: {e}"

# --- Submission Logic ---
def run_and_submit_all(profile: gr.OAuthProfile | None):
    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"

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

    # 1. 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:
            return "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except Exception as e:
        print(f"Error fetching questions: {e}")
        return f"Error fetching questions: {e}", None

    # 2. Run Agent on each question
    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 = run_agent(question_text)
            print(f"Task {task_id[:12]}... | Answer: {submitted_answer[:100]}")

            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:
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # 3. Submit
    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload,
    }
    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")

    try:
        response = requests.post(submit_url, json=submission_data, timeout=120)
        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 Exception:
            error_detail += f" Response: {e.response.text[:500]}"
        return f"Submission Failed: {error_detail}", 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("# Agent Evaluation Runner")
    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 your agent, submit answers, and see the score.

        ---
        **Note:** This will take several minutes as the agent processes all 20 questions.
        """
    )

    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 = os.getenv("SPACE_HOST")
    space_id = os.getenv("SPACE_ID")

    if space_host:
        print(f"✅ SPACE_HOST: {space_host}")
    else:
        print("ℹ️  Running locally (no SPACE_HOST).")

    if space_id:
        print(f"✅ SPACE_ID: {space_id}")
    else:
        print("ℹ️  Running locally (no SPACE_ID).")

    print("-" * 74 + "\n")
    print("Launching Gradio Interface...")
    demo.launch(debug=True, share=False)