import re
from typing import Literal

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.rate_limiters import InMemoryRateLimiter
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode
from langgraph.types import Command
from pydantic import BaseModel, Field

from tools import (
    arxiv_search,
    calculator,
    describe_image_file,
    download_and_read_pdf,
    execute_python_code,
    execute_python_file,
    fetch_webpage,
    get_youtube_transcript,
    read_excel_file,
    transcript_audio_file,
    web_search,
    wikipedia_search,
    wikipedia_section,
)

tools = [
    calculator,
    wikipedia_search,
    wikipedia_section,
    arxiv_search,
    web_search,
    fetch_webpage,
    read_excel_file,
    get_youtube_transcript,
    execute_python_file,
    execute_python_code,
    transcript_audio_file,
    describe_image_file,
    download_and_read_pdf,
]

SYSTEM_PROMPT = """
You are a GAIA benchmark problem-solving assistant specializing in multi-step reasoning tasks.

CRITICAL FIRST STEPS:

1. **File Attachment Check**: If the question mentions or includes an attached file,
   IMMEDIATELY inspect it using the appropriate tool BEFORE any other actions.
   - Excel/CSV: Use read_excel_file or appropriate reader
   - Image: Describe what you need to analyze
   - Always verify the file path and extension

2. **Question Analysis**: Identify:
   - Required precision (decimal places, units, format)
   - Temporal constraints ("as of DATE", "current", "latest")
   - Inclusion/exclusion criteria (be explicit about what counts)
   - Expected answer type (number, string, list, date)

3. **Text Capitalization Preservation**:
   - When answering with text from audio transcripts, videos, or source materials,
     preserve the EXACT capitalization as it appears in the source
   - "Extremely" from transcript → answer "Extremely" (not "extremely")
   - For reversed/transformed text, maintain original capitalization unless explicitly asked to change it
   - Single-word answers: match the capitalization from the context

TOOL SELECTION STRATEGY:

- **Wikipedia tasks**:
  - Use `wikipedia_search` to discover relevant pages and sections
  - Use `wikipedia_section` to retrieve detailed content
  - NEVER use for historical "as of [past date]" questions - Wikipedia shows current content

- **Sports questions** (player stats, rosters, uniform numbers, team info):
  - Be specific: "Taishō Tamai uniform number Hokkaido Nippon-Ham Fighters July 2023"
  - The tool automatically fetches and parses pages - look for patterns like "#18" or "No. 18"
  - Try these strategies:
    * Try `web_search` with different query phrasing
    * Use `fetch_webpage` on specific URLs from sports_statistics results
    * Look for official team roster pages or player profile pages
  - Be persistent: try 2-3 different searches before giving up
  - For player bios without numbers, `wikipedia_search` may work
  - **For "number before and after" questions**:
    1. First verify the exact uniform number of the reference player
    2. Search for the complete roster with all uniform numbers
    3. Find players with N-1 and N+1 uniform numbers
    4. Verify both players match any position/role requirements stated in the question
    5. Double-check you have the correct player name for each adjacent number

- **Calculations**:
  - ALWAYS use `calculator` for arithmetic - never calculate manually
  - Break complex calculations into verified steps
  - State your calculation expression clearly before executing

- **Academic Paper Research**:
    - NEVER guess or invent URLs, DOI values, arXiv IDs, paper titles,
        award numbers, grant numbers, dates, or author names.
        1. Use `fetch_webpage` to inspect the webpage.
        2. Extract the exact paper title and, when available, its actual URL.
        3. If the URL is unavailable, call `arxiv_search` using the exact
            paper title in quotation marks.
        4. Verify that the title and authors match before continuing.
        5. Obtain the canonical PDF URL from the search result.
        6. Use `download_and_read_pdf` to inspect the paper.
        7. For funding questions, search the complete PDF for:
            - acknowledgments
            - acknowledgement
            - supported
            - funding
            - grant
            - award
            - NASA
            - the author's full name and initials
        8. Answer only when the requested value appears in the paper itself.
    - NEVER construct an arXiv URL from a guessed identifier.
    - NEVER use `fetch_webpage` to retrieve arXiv content.
    - Use `arxiv_search` for arXiv discovery and
    `download_and_read_pdf` for the paper content.
    - A failed URL is not evidence that the paper cannot be found.
    Retry using the exact title, authors, and distinctive title phrases.


- **Web Research**:
  - Use `web_search` for current events, recent data, or when Wikipedia is insufficient
  - Use `arxiv_search` for academic papers only
  - Use `download_and_read_pdf` for PDF content retrieval
  - Use `fetch_webpage` for specific URLs (but NOT for Wikipedia or arXiv)

- **YouTube Videos**:
    - Use `get_youtube_transcript` to retrieve transcripts
    - If transciption is incomplete, try to look information about the video on the web with `web_search` or `fetch_webpage`.

- **LibreText**:
    - Use `web_search` to find relevant LibreText: query "site:chem.libretexts.org" + exact section title, author names, license type (e.g., "CK-12"), compilation date (e.g., "08/21/2023"), and distinctive phrases
    - When found the correct website, use `fetch_webpage` with max_chars=10000+ to retrieve full content
    - Search within the fetched content for the specific detail requested
    - If not found, try alternative phrasings, related sections, or different LibreText URLs with the same compilation date

- **Chess Position Analysis**:
    - Use `describe_image_file` to get complete board description
    - State the complete position explicitly before analyzing moves
    - For "guaranteed win" or "correct next move" questions, verify the move carefully
    - Consider multiple candidate moves systematically
    - If analysis is complex, acknowledge potential limitations of image-based chess analysis

- Execute Python code to solve problems or analyze data when necessary. Use `execute_python_code` for snippets and `execute_python_file` for full scripts.

BOTANICAL VS CULINARY DEFINITIONS:

When a question emphasizes "botanical" accuracy or mentions a "botany professor" or "botanist":
- **Botanical vegetables**: Plant parts that are NOT the fruit/seed-bearing structure
  * TRUE botanical vegetables: lettuce, celery, broccoli, fresh basil, sweet potatoes, spinach, kale
- **Botanical fruits** (NOT vegetables, even if commonly called vegetables):
  * Bell pepper, zucchini, green beans, corn, peas, tomatoes, cucumbers, pumpkins, eggplant, okra
- When in doubt about botanical classification, use `wikipedia_search` for the specific item + "botanical classification" or "fruit"
- Example: "bell pepper botanical classification" will confirm it's botanically a fruit

FOOD CATEGORIZATION:

- For "food (not including drinks)" or similar distinctions:
  * Drinks typically include: soda, juice, coffee, tea, milk, water, sports drinks
  * Food typically includes: ice cream, frozen yogurt, smoothies (unless context suggests otherwise)
  * When uncertain, state your interpretation explicitly before calculating
  * Verify category boundaries by checking if items are liquid beverages vs solid/semi-solid foods

VERIFICATION PROCESS:

1. After retrieving data, explicitly state what you found
2. Before calculating, list all values you'll use
3. Verify your calculation includes all required categories and excludes forbidden ones
4. Check your answer matches the requested format and precision
5. You can check for a list of vegetables on the wiki page "List of vegetables" to verify if a vegetable is included in the list.
6. If you are unsure about a food or drink item, you can verify it on the wiki page like "Ice cream".

HANDLING AMBIGUITY:

- If data is unclear, state your interpretation explicitly
- If multiple interpretations exist, mention them
- If data is missing, try alternative sources before giving up

MANDATORY RELIABILITY RULES:

- Never guess URLs, arXiv IDs, DOI values, award numbers, dates,
  uniform numbers, names, or other exact identifiers.

- For finite operation tables, commutativity, associativity, or exhaustive
  comparisons, MUST use `execute_python_code`. The row is the first operand
  and the column is the second operand.
  * For counter-examples: Find ALL pairs (x, y) where the property fails
  * Return ONLY the unique elements involved in those specific failing pairs
  * Example: if a*b ≠ b*a and b*e ≠ e*b fail, return {a, b, e} NOT all elements
  * Use set() to collect unique elements from failing pairs, then sort as requested
  * Do NOT return elements that are not involved in any counter-example

- For academic-paper questions:
  1. Locate the referenced webpage with `web_search`.
  2. Inspect it with `fetch_webpage`.
  3. Extract the exact paper title or verified link.
  4. Search the exact title with `arxiv_search`; never guess an arXiv ID.
  5. Verify title and authors.
  6. Use `download_and_read_pdf` for acknowledgments, grants, funding,
     awards, or facts found inside the paper.

- For LibreTexts questions, search the official domain using exact section
  titles, author names, license, compilation date, and distinctive phrases.
  Use fetch_webpage with max_chars=10000+ to get full content, not just snippets.
  Search within the fetched content for the specific detail. If the first URL
  doesn't have the exact compilation date, try alternative LibreText URLs.

- For historical sports questions, treat the requested date as mandatory.
  Verify the player's number on a dated roster. For “number before and
  after": (1) confirm the reference player's exact number N, (2) find the
  complete roster, (3) identify players with numbers N-1 and N+1, (4) verify
  both match any position requirements, (5) double-check names before answering.

- If a result is irrelevant or incomplete, refine the search at least once
  using exact titles, phrases, authors, dates, or official sources.

FINAL ANSWER FORMAT:

Finish with exactly: FINAL ANSWER: [answer]

CRITICAL FORMAT RULES:

1. **Numbers - NEVER include symbols or unit text**:
   - No thousands separators (not 89,706 but 89706)
   - No currency symbols (not $89706 but 89706)
   - No unit abbreviations (not 89706 USD but just 89706.00)
   - No percentage signs (not 50% but just 50)

2. **Understanding unit phrases**:
   - "Express in USD" = use USD as the measurement scale (DO NOT append "USD")
   - "Express in dollars" = same as above (DO NOT append "dollars")
   - "in kilometers" = use km scale (DO NOT append "km")
   - "as a percentage" = calculate percentage value (DO NOT append "%")

3. **When TO append unit text** (RARE):
   - ONLY when explicitly asked: "include the text USD after the number"
   - ONLY when explicitly asked: "append the currency code"
   - Otherwise, NEVER append unit text

4. **Decimal precision**:
   - Match EXACTLY the requested decimal places
   - "two decimal places" → 89706.00 (not 89706 or 89706.0)
   - If no precision specified, use reasonable precision for the context

5. **String answers**:
   - Use minimal words, be direct
   - No extra explanation after FINAL ANSWER
   - **Preserve EXACT capitalization from the source** (transcript, text, webpage)
   - If source says "Extremely", answer "Extremely" not "extremely"
   - If source says "Right", answer "Right" not "right"

6. **List answers**:
   - Comma-separated values
   - No extra commentary
   - Preserve exact capitalization from source when applicable

CORRECT EXAMPLES:
- "Express your answer in USD with two decimal places" → FINAL ANSWER: 89706.00
- "Give the answer in kilometers" → FINAL ANSWER: 42.5
- "What percentage?" → FINAL ANSWER: 23.4
- "Include USD after the number" → FINAL ANSWER: 89706.00 USD

INCORRECT EXAMPLES (DO NOT DO THIS):
- FINAL ANSWER: 89706.00 USD  ← WRONG if question says "in USD" (should be 89706.00)
- FINAL ANSWER: $89706.00  ← WRONG (never use $ symbol)
- FINAL ANSWER: 89,706.00  ← WRONG (no thousands separators)
- FINAL ANSWER: 23.4%  ← WRONG if question says "as a percentage" (should be 23.4)
"""

FINAL_REVIEW_SYSTEM_PROMPT = """
You are a final-answer verification specialist for GAIA benchmark tasks.

You receive:
- <original_request>: The user's question
- <current_attempt>: Tool calls and results from this attempt
- <candidate_answer>: The proposed final answer

Your task: Verify the answer is correct based ONLY on the current attempt's evidence.

IMPORTANT: Be lenient and pragmatic. Only flag critical errors that make the answer wrong.
Minor format issues are acceptable if the core answer value is correct.

VALIDATION PRIORITIES (in order):

1. **Answer Value Correctness** (MOST IMPORTANT):
   - Does the numeric/text value match the tool evidence?
   - Is the calculation result correct according to the calculator output?
   - If the value is right, proceed to next checks

2. **Data Selection**:
   - Were the required data points used?
   - Only flag if CRITICAL data was missed or wrong data was included
   - Don't flag if alternative valid approaches were used

3. **Tool Appropriateness**:
   - Only flag if the WRONG tool was used (e.g., Wikipedia for historical data)
   - Trust calculator tool results - don't second-guess the math

4. **Format Compliance** (CHECK CAREFULLY):
   - NO currency symbols ($, €, £)
   - NO unit text UNLESS explicitly requested ("include the text USD")
   - "Express in USD" means the value in USD scale, NOT appending "USD" text
   - NO thousands separators
   - Correct decimal precision as requested
   - BUT: if answer value is correct and only has minor format issues, be lenient

WHEN TO ACCEPT AN ANSWER:

Accept if ANY of these are true:
- The answer value is correct and format is reasonable (even if not perfect)
- The answer is mathematically correct per the calculator tool
- Minor format issues exist but the core answer is right and interpretable

WHEN TO REJECT AN ANSWER:

Only reject if:
- The numeric/text value is factually wrong
- Critical required data was not used in the calculation
- The wrong calculation was performed
- Major format violation (e.g., has "USD" text when question said "in USD")
- Wrong decimal precision (e.g., 0 decimals when 2 were requested)

DO NOT REJECT FOR:
- Minor decimal precision differences if reasonable
- Extra/missing trailing zeros (unless explicitly required)
- Different but equally valid interpretation of ambiguous questions

COMMON FORMAT ERRORS TO CATCH:
- "Express in USD" → answer has "USD" appended (REJECT - should be just the number)
- "in kilometers" → answer has "km" appended (REJECT - should be just the number)
- Answer has $ symbol (REJECT - never allowed)
- Answer has commas (REJECT - no thousands separators)

TRUST THE TOOLS:
- If calculator was used, trust its output - don't recalculate manually
- If the right data went into calculator, and result was transcribed correctly, accept it

OUTPUT REQUIREMENTS:

Return structured ReviewResult with:
- `data_selection_correct`: Were the right data points used? (Be lenient)
- `answer_correct`: Is the core value correct per the evidence? (Most important)
- `format_correct`: Is formatting correct per rules? (Check unit text carefully)
- `errors`: ONLY list critical errors that make the answer wrong
- `feedback`: If rejecting, provide specific correction (don't reject for trivial issues)

When in doubt about VALUE correctness, ACCEPT. But enforce format rules for units/symbols.
"""


tools = [
    calculator,
    wikipedia_search,
    wikipedia_section,
    arxiv_search,
    web_search,
    fetch_webpage,
    read_excel_file,
    get_youtube_transcript,
    execute_python_file,
    execute_python_code,
    transcript_audio_file,
    describe_image_file,
    download_and_read_pdf,
]


class MessagesStateWithRetry(MessagesState):
    """Extended state that includes retry counter to prevent infinite loops."""

    retry_count: int = 0
    max_retries: int = 3  # Maximum number of review retry attempts


class ReviewResult(BaseModel):
    data_selection_correct: bool = Field(
        description=(
            "Whether all required rows, columns, categories, and values "
            "were included and excluded correctly."
        )
    )

    answer_correct: bool = Field(
        description=(
            "Whether the candidate answer has the correct value according "
            "to the raw tool evidence."
        )
    )

    format_correct: bool = Field(
        description=("Whether the candidate follows the requested final-answer format.")
    )

    errors: list[str] = Field(
        description=(
            "All material errors. Do not mention missing explanations unless "
            "the original user explicitly requested an explanation."
        )
    )

    feedback: str = Field(
        description=(
            "Actionable correction instructions for the assistant. "
            "Empty when everything is correct."
        )
    )


rate_limiter = InMemoryRateLimiter(
    requests_per_second=1 / 30,
    check_every_n_seconds=0.2,
    max_bucket_size=1,
)

llm = ChatOpenAI(
    model="gpt-4o",
    temperature=0,
    rate_limiter=rate_limiter,
    max_retries=5,
    timeout=120,
)


llm_with_tools = llm.bind_tools(tools, parallel_tool_calls=False)

review_model = llm.with_structured_output(ReviewResult)


def assistant(state: MessagesStateWithRetry):
    return {
        "messages": [
            llm_with_tools.invoke([SystemMessage(SYSTEM_PROMPT)] + state["messages"])
        ]
    }


def build_review_input(state: MessagesStateWithRetry) -> str:
    messages = state["messages"]

    original_request = messages[0].content

    for candidate_index in range(len(messages) - 1, -1, -1):
        if (
            isinstance(messages[candidate_index], AIMessage)
            and messages[candidate_index].content
            and not messages[candidate_index].tool_calls
        ):
            break

    candidate_answer = messages[candidate_index].content

    # Last feedback index
    last_review_index = next(
        (
            index
            for index in range(candidate_index - 1, -1, -1)
            if (
                isinstance(messages[index], HumanMessage)
                and str(messages[index].content).startswith(
                    "FINAL ANSWER REVIEW FAILED"
                )
            )
        ),
        -1,
    )

    attempt_start = last_review_index + 1

    current_attempt = []

    for message in messages[attempt_start:candidate_index]:
        if isinstance(message, AIMessage):
            for t in message.tool_calls or []:
                current_attempt.append(
                    f"<tool_call>\nname: {t['name']}\narguments: {t.get('args', {})}\n</tool_call>"
                )

        elif isinstance(message, ToolMessage):
            current_attempt.append(
                f"<tool_result>\n"
                f"name: {message.name}\n"
                f"content:\n{message.content}\n"
                f"</tool_result>"
            )

    return f"""
<original_request>
{original_request}
</original_request>

<current_attempt>
{"\n\n---\n\n".join(current_attempt)}
</current_attempt>

<candidate_answer>
{candidate_answer}
</candidate_answer>
""".strip()


def review_passed(review: ReviewResult) -> bool:
    """Check if review passed with lenient criteria.

    Accept if:
    - Answer is correct (most important)
    - No critical errors
    - Format is reasonable (even if not perfect)
    """
    # Most important: is the answer value correct?
    if not review.answer_correct:
        return False

    # If answer is correct but has minor format issues, still accept
    if review.format_correct and review.data_selection_correct:
        return True

    # If answer is correct but format/data selection slightly off,
    # only reject if there are actual critical errors listed
    if not review.errors:
        # No errors listed means issues are minor, accept it
        return True

    # Check if errors are critical
    critical_keywords = [
        "wrong value",
        "incorrect calculation",
        "missing required",
        "excluded incorrectly",
    ]
    has_critical_error = any(
        any(keyword in error.lower() for keyword in critical_keywords)
        for error in review.errors
    )

    # If no critical errors and answer is correct, accept
    return not has_critical_error


def final_review(
    state: MessagesStateWithRetry,
) -> Command[Literal["assistant", "__end__"]]:
    review = review_model.invoke(
        [
            SystemMessage(content=FINAL_REVIEW_SYSTEM_PROMPT),
            HumanMessage(build_review_input(state)),
        ]
    )

    if review_passed(review):
        # Last assistant message is the final answer, so we can end the conversation
        return Command(goto=END)

    # Check if we've exceeded max retries
    current_retry = state.get("retry_count", 0)
    max_retries = state.get("max_retries", 3)

    if current_retry >= max_retries:
        print(
            f"Max retries ({max_retries}) reached. Accepting current answer despite errors."
        )
        # End the conversation with the current answer even if it has errors
        return Command(goto=END)

    errors = "\n".join(f"- {error}" for error in review.errors)

    reviewer_feedback = HumanMessage(
        content=(
            f"FINAL ANSWER REVIEW FAILED (Attempt {current_retry + 1}/{max_retries}).\n\n"
            f"Errors found:\n{errors}\n\n"
            f"Correction instructions:\n{review.feedback}\n\n"
            "Review the raw source data and only the current attempt. "
            "Include all required categories, exclude all disallowed "
            "categories, recalculate when necessary, and finish with the "
            "required FINAL ANSWER format."
        )
    )

    return Command(
        update={
            "messages": [reviewer_feedback],
            "retry_count": current_retry + 1,
        },
        goto="assistant",
    )


def route_after_assistant(
    state: MessagesStateWithRetry,
) -> Literal["tools", "final_review"]:
    last_message = state["messages"][-1]

    if isinstance(last_message, AIMessage) and last_message.tool_calls:
        return "tools"

    return "final_review"


builder = StateGraph(MessagesStateWithRetry)

# Nodes
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
builder.add_node("final_review", final_review)


# Edges
builder.add_edge(START, "assistant")
builder.add_conditional_edges("assistant", route_after_assistant)
builder.add_edge("tools", "assistant")


# Basic Gaia Agent
class BasicAgent:
    def __init__(self, max_retries: int = 3):
        # Compile
        self.graph = builder.compile()
        self.max_retries = max_retries

    def __call__(self, question: str) -> str:
        messages = self.graph.invoke(
            {"messages": question, "retry_count": 0, "max_retries": self.max_retries}
        )
        m = re.search(
            r".*FINAL ANSWER: (?P<answer>.*)", messages["messages"][-1].content
        )
        if m is None:
            return {"messages": messages["messages"], "answer": "Unresolved"}

        final_answer = m.group("answer")
        return {"messages": messages["messages"], "answer": final_answer}

    def draw_mermaid(self):
        return self.graph.get_graph(xray=True).draw_mermaid_png()


__all__ = ["BasicAgent"]