# ==========================================================
# agent.py — Flat routing, single model path for all questions
# ==========================================================

import os
import re
import sys
import json
import tempfile
import subprocess
from pathlib import Path
from urllib.parse import urlparse

from dotenv import load_dotenv

from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import ToolNode, tools_condition

from langchain_groq import ChatGroq
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage

from langchain_community.document_loaders import WikipediaLoader
from langchain_community.tools import DuckDuckGoSearchResults

import requests
from bs4 import BeautifulSoup
import pandas as pd

try:
    from youtube_transcript_api import YouTubeTranscriptApi
except Exception:
    YouTubeTranscriptApi = None

load_dotenv()

# ==========================================================
# OBSERVABILITY GLOBALS
# ==========================================================

LAST_MODEL_USED = "N/A"
LAST_MODEL_FALLBACK = "No"
LAST_MODEL_ERROR = "None"

# ==========================================================
# TOOLS
# ==========================================================

@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for encyclopedic facts, biographies, history, and stable knowledge."""
    try:
        docs = WikipediaLoader(query=query, load_max_docs=2).load()
        if not docs:
            return "No Wikipedia results found. Try web_search instead."
        return "\n\n".join(d.page_content[:2000] for d in docs)
    except Exception as e:
        return f"Wikipedia unavailable: {type(e).__name__}. Try web_search instead."


@tool
def web_search(query: str) -> str:
    """Search the web and return compact title/url/snippet results."""
    try:
        result = DuckDuckGoSearchResults(
            max_results=5,
            output_format="list",
        ).run(query)
        return result if result else "No results found. Try a different query."
    except Exception as e:
        return f"Web search unavailable: {type(e).__name__}. Try wiki_search instead."


@tool
def fetch_page(url: str) -> str:
    """Fetch and extract readable text from a URL."""
    try:
        r = requests.get(url, timeout=12, headers={"User-Agent": "Mozilla/5.0"})
        r.raise_for_status()
        soup = BeautifulSoup(r.text, "html.parser")
        for tag in soup(["script", "style", "nav", "footer", "header"]):
            tag.decompose()
        text = soup.get_text(separator="\n", strip=True)
        return text[:8000] if text else "Page appears empty."
    except Exception as e:
        return f"Could not fetch page: {type(e).__name__}."


def _youtube_video_id(text: str) -> str | None:
    patterns = [
        r"(?:v=|youtu\.be/|shorts/|embed/)([A-Za-z0-9_-]{11})",
        r"^([A-Za-z0-9_-]{11})$",
    ]
    for pattern in patterns:
        match = re.search(pattern, text)
        if match:
            return match.group(1)
    return None


@tool
def youtube_transcript(video_url_or_id: str) -> str:
    """Get available YouTube captions/transcript for a video URL or video id."""
    if YouTubeTranscriptApi is None:
        return "YouTube transcript library is unavailable."
    video_id = _youtube_video_id(video_url_or_id)
    if not video_id:
        return "Could not identify a YouTube video id."
    try:
        rows = YouTubeTranscriptApi.get_transcript(video_id, languages=["en", "en-US", "en-GB"])
    except Exception:
        try:
            rows = YouTubeTranscriptApi.get_transcript(video_id)
        except Exception as e:
            return f"Transcript unavailable: {type(e).__name__}. Use web_search for quotes or descriptions."
    lines = []
    for row in rows[:220]:
        start = int(float(row.get("start", 0)))
        text = " ".join(str(row.get("text", "")).split())
        if text:
            lines.append(f"{start}s: {text}")
    return "\n".join(lines)[:7000] if lines else "Transcript is empty."


def _download_to_temp(source: str) -> tuple[Path | None, str]:
    if re.fullmatch(r"[0-9a-fA-F-]{36}", source.strip()):
        source = f"https://agents-course-unit4-scoring.hf.space/files/{source.strip()}"
    if not source.startswith(("http://", "https://")):
        return None, "Input must be a URL or benchmark task_id."

    try:
        r = requests.get(source, timeout=25, headers={"User-Agent": "Mozilla/5.0"})
        if r.status_code == 404:
            return None, "No attached file found for this task."
        r.raise_for_status()
    except Exception as e:
        return None, f"Download failed: {type(e).__name__}."

    parsed = urlparse(source)
    suffix = Path(parsed.path).suffix
    content_type = r.headers.get("content-type", "").lower()
    if not suffix:
        if "spreadsheet" in content_type or "excel" in content_type:
            suffix = ".xlsx"
        elif "csv" in content_type:
            suffix = ".csv"
        elif "pdf" in content_type:
            suffix = ".pdf"
        elif "python" in content_type or "text" in content_type:
            suffix = ".txt"
        elif "image" in content_type:
            suffix = "." + content_type.split("/")[-1].split(";")[0]
        else:
            cd = r.headers.get("content-disposition", "")
            match = re.search(r'filename="?([^";]+)', cd)
            suffix = Path(match.group(1)).suffix if match else ".bin"

    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as f:
        f.write(r.content)
        return Path(f.name), f"Downloaded {len(r.content)} bytes as {suffix or '.bin'}."


def _read_text_file(path: Path) -> str:
    for enc in ("utf-8", "latin-1"):
        try:
            return path.read_text(encoding=enc)[:9000]
        except Exception:
            continue
    return "Could not decode text file."


@tool
def inspect_file(source: str) -> str:
    """Download and inspect an attached benchmark file or URL. Input can be a task_id or URL."""
    path, status = _download_to_temp(source)
    if path is None:
        return status

    try:
        suffix = path.suffix.lower()

        if suffix in {".py", ".txt", ".md", ".json", ".html", ".csv"}:
            if suffix == ".csv":
                df = pd.read_csv(path)
                return _summarize_dataframe(df, "csv")
            text = _read_text_file(path)
            if suffix == ".py":
                run = subprocess.run(
                    [sys.executable, str(path)],
                    capture_output=True,
                    text=True,
                    timeout=20,
                )
                stdout = run.stdout.strip()
                stderr = run.stderr.strip()
                return f"{status}\nPython stdout:\n{stdout[:3000] or '(empty)'}\nPython stderr:\n{stderr[:1000] or '(empty)'}\n\nCode preview:\n{text[:4000]}"
            return f"{status}\n{text}"

        if suffix in {".xlsx", ".xls"}:
            xl = pd.ExcelFile(path)
            parts = [status, f"Workbook sheets: {', '.join(xl.sheet_names)}"]
            for sheet in xl.sheet_names[:4]:
                df = xl.parse(sheet)
                parts.append(_summarize_dataframe(df, sheet))
            return "\n\n".join(parts)[:9000]

        if suffix == ".pdf":
            try:
                from pypdf import PdfReader
                reader = PdfReader(str(path))
                text = "\n".join((p.extract_text() or "") for p in reader.pages[:8])
                return f"{status}\nPDF pages: {len(reader.pages)}\n{text[:8500] or 'No extractable text.'}"
            except Exception as e:
                return f"{status}\nPDF extraction unavailable: {type(e).__name__}."

        if suffix in {".png", ".jpg", ".jpeg", ".webp", ".gif"}:
            return f"{status}\nImage file detected. If the question requires visual reasoning, use any visible description in the question and web_search; this runtime has no vision model."

        if suffix in {".mp3", ".wav", ".m4a", ".ogg", ".flac"}:
            return f"{status}\nAudio file detected. This runtime has no local speech-to-text model; use web_search if the recording is from public material."

        return f"{status}\nUnsupported file type: {suffix or 'unknown'}."
    except Exception as e:
        return f"{status}\nInspection failed: {type(e).__name__}: {e}"
    finally:
        try:
            path.unlink(missing_ok=True)
        except Exception:
            pass


def _summarize_dataframe(df: pd.DataFrame, name: str) -> str:
    rows, cols = df.shape
    df = df.dropna(how="all")
    preview = df.head(12).to_string(index=False)
    numeric = df.select_dtypes(include="number")
    sums = numeric.sum(numeric_only=True).to_dict()
    sums_text = json.dumps({str(k): round(float(v), 4) for k, v in sums.items()}, ensure_ascii=True)
    cols_text = ", ".join(map(str, df.columns))
    food_hint = ""
    lower_cols = {str(c).lower(): c for c in df.columns}
    category_col = next((c for key, c in lower_cols.items() if any(x in key for x in ["category", "type", "item type"])), None)
    sales_col = next((c for key, c in lower_cols.items() if any(x in key for x in ["sales", "revenue", "amount", "total"])), None)
    if category_col is not None and sales_col is not None:
        try:
            grouped = df.groupby(category_col)[sales_col].sum(numeric_only=True).to_dict()
            food_hint = "\nGrouped sums: " + json.dumps({str(k): round(float(v), 2) for k, v in grouped.items()}, ensure_ascii=True)
        except Exception:
            pass
    return f"Sheet/table {name}: {rows} rows x {cols} cols\nColumns: {cols_text}\nNumeric sums: {sums_text}{food_hint}\nPreview:\n{preview}"


@tool
def run_python(code: str) -> str:
    """Execute Python code and return stdout. Use for calculations, counting, data processing."""
    fname = None
    try:
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write(code)
            fname = f.name
        p = subprocess.run(
            [sys.executable, fname],
            capture_output=True,
            text=True,
            timeout=20
        )
        out = p.stdout.strip()
        err = p.stderr.strip()
        if out:
            return out
        if err:
            return f"Error: {err}"
        return "Script ran but produced no output."
    except subprocess.TimeoutExpired:
        return "Script timed out after 20 seconds."
    except Exception as e:
        return f"Execution failed: {type(e).__name__}."
    finally:
        if fname:
            try:
                Path(fname).unlink(missing_ok=True)
            except Exception:
                pass


@tool
def reverse_text(text: str) -> str:
    """Reverse a string character by character."""
    return text[::-1]


TOOLS = [wiki_search, web_search, fetch_page, youtube_transcript, inspect_file, run_python, reverse_text]

# ==========================================================
# MODELS — primary + ordered fallback chain
# ==========================================================

def _llm(name: str) -> ChatGroq:
    return ChatGroq(
        model=name,
        api_key=os.getenv("GROQ_API_KEY"),
        temperature=0,
        max_tokens=768,
        timeout=45,
        max_retries=1,
    )


# All questions use the same primary model.
# Fallback chain kicks in only on errors (rate limits, timeouts, etc.)
MODEL_PRIMARY  = _llm("qwen/qwen3-32b")
MODEL_FALLBACK = _llm("llama-3.3-70b-versatile")
MODEL_LAST     = _llm("llama-3.1-8b-instant")

FALLBACK_CHAIN = [MODEL_PRIMARY, MODEL_FALLBACK, MODEL_LAST]

# ==========================================================
# SYSTEM PROMPT — single prompt for all question types
# ==========================================================

SYSTEM_PROMPT = """You are a precise benchmark task solver.

## Goal
Produce the exact correct answer — nothing more, nothing less.

## Tool use
- Use wiki_search for historical facts, biographies, science, geography.
- Use web_search for recent events, specific articles, prices, or anything time-sensitive.
- Use fetch_page when a URL is provided or a search result points to a relevant page.
- Use youtube_transcript first for YouTube questions.
- Use inspect_file whenever the question says attached file, attached image, spreadsheet, audio, Python code, or provides a task file URL.
- Use run_python for any arithmetic, counting, sorting, or data transformation.
- Use reverse_text only when asked to reverse a string.
- Do not inspect a task file unless the question mentions an attachment, image, audio, spreadsheet, code file, or file URL.
- For YouTube questions, if transcript is unavailable, search the exact video id plus the specific requested phrase/object.
- Prefer tools over guessing. Use compact searches. Stop as soon as you have a confident answer.

## Answer format rules
1. Output the raw value only — no explanation, no preamble.
2. If asked for a first name, output only the first name.
3. If asked for a surname, output only the surname.
4. Numbers: digits only unless units were explicitly requested.
5. Lists: comma-separated on one line.
6. If you cannot find the answer after searching, output: N/A

## Required final line
Always end your response with exactly:
FINAL ANSWER: <your answer>
"""

# ==========================================================
# INVOKE — with fallback chain
# ==========================================================

def invoke(messages: list) -> object:
    global LAST_MODEL_USED, LAST_MODEL_FALLBACK, LAST_MODEL_ERROR

    LAST_MODEL_FALLBACK = "No"
    LAST_MODEL_ERROR = "None"

    seen: set[str] = set()
    first = True

    for model in FALLBACK_CHAIN:
        key = model.model_name
        if key in seen:
            continue
        seen.add(key)

        try:
            LAST_MODEL_USED = key
            if not first:
                LAST_MODEL_FALLBACK = "Yes"
            first = False
            return model.bind_tools(TOOLS).invoke(messages)
        except Exception as e:
            LAST_MODEL_ERROR = str(e)
            continue

    raise RuntimeError(f"All models failed. Last error: {LAST_MODEL_ERROR}")

# ==========================================================
# GRAPH NODES
# ==========================================================

def assistant(state: MessagesState) -> dict:
    messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
    result = invoke(messages)
    return {"messages": [result]}


def build_graph():
    g = StateGraph(MessagesState)
    g.add_node("assistant", assistant)
    g.add_node("tools", ToolNode(TOOLS))
    g.add_edge(START, "assistant")
    g.add_conditional_edges("assistant", tools_condition)
    g.add_edge("tools", "assistant")
    return g.compile()

# ==========================================================
# ANSWER EXTRACTION + CLEANING
# ==========================================================

def _clean_answer(raw: str) -> str:
    """Normalise the extracted answer string."""
    answer = raw.strip()
    # Strip trailing punctuation that the model sometimes adds
    answer = answer.rstrip(".,;:")
    # Collapse internal whitespace / newlines
    answer = " ".join(answer.split())
    # Remove common LLM filler prefixes the regex sometimes captures
    for prefix in (
        "the answer is",
        "answer is",
        "answer:",
        "it is",
        "it's",
        "that is",
        "this is",
    ):
        if answer.lower().startswith(prefix):
            answer = answer[len(prefix):].strip()
    return answer


def extract_final_answer(output: dict) -> str:
    msgs = output.get("messages", [])
    for m in reversed(msgs):
        txt = getattr(m, "content", "")
        if isinstance(txt, str):
            match = re.search(r"FINAL ANSWER:\s*(.+)", txt, re.I | re.S)
            if match:
                raw = match.group(1).strip()[:300]
                return _clean_answer(raw)
    return "N/A"


def extract_tools_used(output: dict) -> list[str]:
    msgs = output.get("messages", [])
    seen: set[str] = set()
    tools: list[str] = []

    for m in msgs:
        name = getattr(m, "name", None)
        if name and name not in seen:
            tools.append(name)
            seen.add(name)
        for tc in getattr(m, "tool_calls", []) or []:
            n = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", None)
            if n and n not in seen:
                tools.append(n)
                seen.add(n)

    return tools


def get_last_trace() -> dict:
    return {
        "model": LAST_MODEL_USED,
        "fallback": LAST_MODEL_FALLBACK,
        "model_error": LAST_MODEL_ERROR,
    }