import os
import json
import time
import inspect
import yaml
import requests
import gradio as gr
import pandas as pd

from smolagents import CodeAgent, OpenAIServerModel, LiteLLMModel, Tool, HfApiModel, ToolCallingAgent

try:
    from litellm.exceptions import RateLimitError
except ImportError:
    RateLimitError = None

from tools.calculator import Calculator
from tools.final_answer import FinalAnswerTool
from tools.visit_webpage import VisitWebpageTool
from tools.web_search import DuckDuckGoSearchTool

from research_state import ResearchState
from verifier import verify_research
import litellm
from tools.youtube_transcript import YouTubeTranscriptTool
from tools.youtube_video_analysis import YouTubeVideoAnalysisTool

youtube_transcript = YouTubeTranscriptTool()

litellm._turn_on_debug()
# =========================================================
# Constants
# =========================================================

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

MODEL_ID = "gemini/gemini-2.5-flash"  #"" "Qwen/Qwen2.5-Coder-32B-Instruct" #"groq/groq/compound-mini" #"meta-llama/llama-prompt-guard-2-22m"

# Which agent to build: "tool_calling" (ToolCallingAgent) or "code" (CodeAgent).
# Both use prompts.yaml.
AGENT_TYPE = "tool_calling"

MAX_STEPS = 6
MAX_TOKENS = 900
TEMPERATURE = 0.2

# If this file exists next to app.py, questions are read from it instead of
# being fetched from the scoring server (avoids 429s from the shared Space IP).
QUESTIONS_FILE = "questions.json"

# Number of questions to skip from the start of the list (set to 0 to run all)
SKIP_FIRST_N = 1

# Pacing / rate-limit handling (model provider)
MIN_SECONDS_BETWEEN_QUESTIONS = 60
MAX_RATE_LIMIT_RETRIES = 5
RATE_LIMIT_BASE_DELAY = 30  # seconds; doubles each retry

# Throttling / retry for every individual LLM call (each agent step is one call)
LLM_MIN_SECONDS_BETWEEN_CALLS = 7   # ~8 calls/min; lower this if your quota allows
LLM_MAX_RETRIES = 5
LLM_RETRY_BASE_DELAY = 15           # seconds; doubles each retry, capped at 120

# Rate-limit handling for the scoring server (fetch questions / submit)
HTTP_MAX_RETRIES = 6
HTTP_BASE_DELAY = 10  # seconds; doubles each retry (10, 20, 40, 80, ...)
HTTP_MAX_DELAY = 120  # cap on any single wait


# =========================================================
# Research Context
# =========================================================

class ResearchContext:
    """
    Holds the state for a single research run.
    The context is shared by the research tools so that
    searches, webpages, claims, and verification results
    belong to the same research session.
    """

    def __init__(self):
        self.state = None

    def start(self, question: str):
        self.state = ResearchState(question=question)

    def ensure_started(self, question: str = ""):
        if self.state is None:
            self.start(question)

    def reset(self):
        self.state = None


research_context = ResearchContext()


# =========================================================
# Stateful Web Search
# =========================================================

class StatefulWebSearchTool(DuckDuckGoSearchTool):
    """
    Web search tool that also records discovered sources
    in the current ResearchState.
    """

    name = "web_search"

    description = (
        "Searches the web for relevant sources. "
        "Use this for source discovery. Search results are "
        "leads and should normally be followed by visit_webpage "
        "when factual evidence is needed."
    )

    def forward(self, query: str) -> str:
        research_context.ensure_started()

        results = super().forward(query)

        # Store the search itself as research activity.
        # The actual source evidence is recorded more precisely
        # when visit_webpage reads the source.
        return results


# =========================================================
# Stateful Webpage Reader
# =========================================================

class StatefulVisitWebpageTool(VisitWebpageTool):
    """
    Reads webpages and records the resulting source content
    in ResearchState.
    """

    name = "visit_webpage"

    description = (
        "Visits a webpage and reads its contents. "
        "Use this after web_search when you need actual "
        "source evidence rather than only a search snippet."
    )

    def forward(self, url: str) -> str:
        research_context.ensure_started()

        content = super().forward(url)

        # Record the source as evidence.
        #
        # At this stage we don't know which exact claim the
        # source supports. The verification step can associate
        # the collected sources with the claims supplied by
        # the agent.
        research_context.state.add_evidence(
            claim="__SOURCE__",
            source_url=url,
            excerpt=content[:4000],
            supports_claim=None,
        )

        return content


# =========================================================
# Verification Tool
# =========================================================

class VerifyResearchTool(Tool):
    """
    Runs the deterministic verification layer over the
    current research state.
    """

    name = "verify_research"

    description = (
        "Verifies the claims in the current research session "
        "against the sources that have been collected. "
        "Use this before producing the final answer."
    )

    inputs = {
        "claims": {
            "type": "string",
            "description": (
                "The factual claims that need verification, "
                "separated by newline."
            ),
        }
    }

    output_type = "string"

    def forward(self, claims: str) -> str:
        research_context.ensure_started()

        state = research_context.state

        # Normalize claims supplied by the agent.
        claim_list = [
            claim.strip()
            for claim in claims.splitlines()
            if claim.strip()
        ]

        # Register the claims.
        for claim in claim_list:
            state.add_claim(claim)

        # Existing evidence was initially recorded against
        # __SOURCE__. Associate the collected sources with
        # the claims being verified.
        #
        # This is intentionally conservative: the verifier
        # knows that these sources were collected, but the
        # model still needs to determine whether they actually
        # support or contradict each claim.
        source_evidence = [
            evidence
            for evidence in state.evidence
            if evidence.claim == "__SOURCE__"
        ]

        # Rebuild claim-level evidence.
        #
        # We do not automatically mark sources as supporting.
        # They remain uncertain until the agent has inspected
        # and reasoned about the evidence.
        for claim in claim_list:
            existing = state.evidence_for_claim(claim)

            if not existing:
                for source in source_evidence:
                    state.add_evidence(
                        claim=claim,
                        source_url=source.source_url,
                        excerpt=source.excerpt,
                        supports_claim=None,
                    )

        verification = verify_research(state)

        return _format_verification_result(verification)


def _format_verification_result(result: dict) -> str:
    lines = []

    lines.append("## Verification Results")
    lines.append("")

    for item in result["claims"]:
        lines.append(
            f"### {item['status']}: {item['claim']}"
        )
        lines.append(item["reason"])

        if item["sources"]:
            lines.append("Sources:")
            for source in item["sources"]:
                lines.append(f"- {source}")

        lines.append("")

    if result["verified"]:
        lines.append("Overall result: SUPPORTED")
    else:
        lines.append(
            "Overall result: NOT FULLY VERIFIED"
        )

    return "\n".join(lines)


# =========================================================
# Rate limit / quota handling (model provider)
# =========================================================

class DailyQuotaExhaustedError(Exception):
    """Raised when the model provider reports a daily/per-day quota
    exhaustion, as opposed to a transient per-minute rate limit."""
    pass


def _is_rate_limit_error(exc: Exception) -> bool:
    """
    Detect rate-limit / quota errors across providers, since exact
    exception types can vary by SDK version. Falls back to matching
    on common message fragments.
    """
    if RateLimitError is not None and isinstance(exc, RateLimitError):
        return True

    message = str(exc).lower()
    signals = [
        "rate limit",
        "rate_limit",
        "429",
        "quota",
        "resource_exhausted",
        "too many requests",
    ]
    return any(signal in message for signal in signals)


def _is_daily_quota_exhausted(exc: Exception) -> bool:
    """
    Heuristically detect a *daily/per-day* quota exhaustion (as opposed
    to a transient per-minute rate limit). These typically won't resolve
    with a short backoff, so retrying is pointless and the whole run
    should stop early rather than burn time and further requests.
    Gemini/Google API quota errors usually include phrases like
    "PerDay", "per day", "daily limit", or a quota_id containing
    "PerDay" when the daily free-tier quota is hit.
    """
    message = str(exc).lower()
    signals = [
        "per day",
        "perday",
        "daily limit",
        "daily quota",
        "generate_content_free_tier_requests_per_day",
    ]
    return any(signal in message for signal in signals)


def _run_agent_with_backoff(agent_runner, question_text,
                             max_retries=MAX_RATE_LIMIT_RETRIES,
                             base_delay=RATE_LIMIT_BASE_DELAY):
    """
    Runs the agent on a single question, retrying with exponential
    backoff if a transient rate-limit error is hit. If the error looks
    like a daily/per-day quota exhaustion, raises DailyQuotaExhaustedError
    immediately instead of retrying, since backoff won't help.
    Non-rate-limit errors are raised immediately (handled by the caller).
    """
    for attempt in range(max_retries + 1):
        try:
            return agent_runner(question_text)
        except Exception as e:
            if _is_daily_quota_exhausted(e):
                raise DailyQuotaExhaustedError(str(e)) from e

            if _is_rate_limit_error(e) and attempt < max_retries:
                delay = base_delay * (2 ** attempt)  # 30, 60, 120, 240, 480...
                print(
                    f"Rate limit/quota hit (attempt {attempt + 1}/{max_retries}). "
                    f"Backing off for {delay}s..."
                )
                time.sleep(delay)
                continue
            raise


# =========================================================
# Rate limit handling (scoring server HTTP requests)
# =========================================================

def _request_with_backoff(method, url,
                          max_retries=HTTP_MAX_RETRIES,
                          base_delay=HTTP_BASE_DELAY,
                          max_delay=HTTP_MAX_DELAY,
                          **kwargs):
    """
    Performs an HTTP request, retrying on 429 (Too Many Requests) and
    transient 502/503/504 responses. Honours the server's Retry-After
    header when present, otherwise uses exponential backoff.
    Raises requests.exceptions.HTTPError (via raise_for_status) if the
    request still fails after all retries or fails for another reason.
    """
    retry_statuses = {429, 502, 503, 504}

    for attempt in range(max_retries + 1):
        response = requests.request(method, url, **kwargs)

        if response.status_code in retry_statuses and attempt < max_retries:
            retry_after = response.headers.get("Retry-After")
            try:
                delay = float(retry_after)
            except (TypeError, ValueError):
                delay = base_delay * (2 ** attempt)
            delay = min(delay, max_delay)

            print(
                f"{method} {url} returned {response.status_code} "
                f"(attempt {attempt + 1}/{max_retries}). "
                f"Retrying in {delay:.0f}s..."
            )
            time.sleep(delay)
            continue

        response.raise_for_status()
        return response


def _load_questions(questions_url):
    """
    Loads questions from a local QUESTIONS_FILE if present, otherwise
    fetches them from the scoring server (with backoff).
    """
    if os.path.exists(QUESTIONS_FILE):
        print(f"Loading questions from local file: {QUESTIONS_FILE}")
        with open(QUESTIONS_FILE, "r", encoding="utf-8") as f:
            return json.load(f)

    response = _request_with_backoff("GET", questions_url, timeout=15)
    return response.json()


# =========================================================
# Throttled + retrying LLM calls
# =========================================================
# smolagents catches model errors internally and turns them into step
# errors / a garbage "final answer", so rate limits never reach the
# question-level backoff. Handling them at the litellm.completion level
# makes every agent step wait and retry instead.

_original_litellm_completion = litellm.completion
_last_llm_call_time = 0.0


def _throttled_completion(*args, **kwargs):
    global _last_llm_call_time

    for attempt in range(LLM_MAX_RETRIES + 1):
        wait = LLM_MIN_SECONDS_BETWEEN_CALLS - (time.monotonic() - _last_llm_call_time)
        if wait > 0:
            time.sleep(wait)
        _last_llm_call_time = time.monotonic()

        try:
            return _original_litellm_completion(*args, **kwargs)
        except Exception as e:
            if _is_daily_quota_exhausted(e):
                raise

            if _is_rate_limit_error(e) and attempt < LLM_MAX_RETRIES:
                delay = min(LLM_RETRY_BASE_DELAY * (2 ** attempt), 120)
                print(
                    f"LLM rate limit hit (attempt {attempt + 1}/{LLM_MAX_RETRIES}). "
                    f"Waiting {delay}s..."
                )
                time.sleep(delay)
                continue
            raise


litellm.completion = _throttled_completion


# =========================================================
# Load Prompts
# =========================================================

with open("prompts.yaml", "r", encoding="utf-8") as stream:
    prompt_templates = yaml.safe_load(stream)


# =========================================================
# Tools
# =========================================================

web_search = StatefulWebSearchTool(
    max_results=10
)

visit_webpage = StatefulVisitWebpageTool()

calculator = Calculator()

verify_research_tool = VerifyResearchTool()

final_answer = FinalAnswerTool()

youtube_video_analysis = YouTubeVideoAnalysisTool(
    model="gemini-2.5-flash"
)

tools = [
    web_search,
    visit_webpage,
    calculator,
    verify_research_tool,
    youtube_transcript,
    final_answer,
    youtube_video_analysis,
]


# =========================================================
# Model
# =========================================================
"""
model = HfApiModel(
    model_id=MODEL_ID,
    provider="auto",     # or "together", "nebius", "fireworks-ai", etc.
    token=os.environ["GEMINI_API_KEY"],
    max_tokens=MAX_TOKENS,
    temperature=TEMPERATURE,
    timeout=120,
)
"""

model = LiteLLMModel(
    model_id=MODEL_ID,
    api_key=os.environ["GEMINI_API_KEY"],
    max_tokens=MAX_TOKENS,
    temperature=TEMPERATURE,
    #tool_config={"function_calling_config": {"mode": "NONE"}},
)


# =========================================================
# Agent
# =========================================================

_agent_kwargs = dict(
    model=model,
    tools=tools,
    max_steps=MAX_STEPS,
    verbosity_level=1,
    planning_interval=None,
    name="web_research_agent",
    description=(
        "A web research and fact-verification agent that "
        "searches for evidence, reads sources, performs "
        "calculations, checks claims, detects uncertainty "
        "and produces cited answers."
    ),
    prompt_templates=prompt_templates,
)

if AGENT_TYPE == "code":
    agent = CodeAgent(grammar=None, **_agent_kwargs)
elif AGENT_TYPE == "tool_calling":
    agent = ToolCallingAgent(**_agent_kwargs)
else:
    raise ValueError(f"Unknown AGENT_TYPE: {AGENT_TYPE!r} (use 'code' or 'tool_calling')")

print(f"Using agent type: {AGENT_TYPE}")


class ResearchAgentWrapper:
    """
    Thin wrapper so the research agent can be called like
    BasicAgent: agent(question) -> answer string.
    Resets the shared research context for every new question
    so state doesn't leak between questions.
    """

    def __init__(self, agent: CodeAgent):
        self.agent = agent
        print("ResearchAgentWrapper initialized.")

    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        research_context.reset()
        research_context.start(question)
        answer = self.agent.run(question)

        # smolagents swallows model errors (e.g. rate limits) and returns them
        # as the "answer". Re-raise so the question-level backoff can retry it
        # and the error text is never submitted as a real answer.
        if str(answer).startswith("Error in generating final LLM output"):
            raise RuntimeError(str(answer))

        print(f"Agent returning answer (first 50 chars): {str(answer)[:50]}...")
        return str(answer)


# =========================================================
# Evaluation Runner
# =========================================================

def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the research agent 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_runner = ResearchAgentWrapper(agent)
    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 (with backoff on 429 / transient server errors)
    print(f"Fetching questions from: {questions_url}")
    try:
        questions_data = _load_questions(questions_url)
        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}")
        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 your Agent
    results_log = []
    answers_payload = []
    quota_exhausted = False

    # Skip the first N questions
    questions_data = questions_data[SKIP_FIRST_N:]
    print(
        f"Skipping first {SKIP_FIRST_N} question(s). "
        f"Running agent on {len(questions_data)} questions..."
    )

    for idx, item in enumerate(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

        start_time = time.monotonic()

        try:
            submitted_answer = _run_agent_with_backoff(agent_runner, question_text)
            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 DailyQuotaExhaustedError as e:
            print(f"Daily quota exhausted, stopping run early: {e}")
            results_log.append({
                "Task ID": task_id,
                "Question": question_text,
                "Submitted Answer": f"SKIPPED (daily quota exhausted): {e}",
            })
            quota_exhausted = True
            break
        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}"})

        # Pace requests: ensure at least MIN_SECONDS_BETWEEN_QUESTIONS
        # seconds between question starts, accounting for however long
        # this question took (including any backoff sleeps it already did).
        if idx < len(questions_data) - 1:
            elapsed = time.monotonic() - start_time
            remaining = MIN_SECONDS_BETWEEN_QUESTIONS - elapsed
            if remaining > 0:
                print(f"Waiting {remaining:.1f}s before the next question (pacing)...")
                time.sleep(remaining)

    if quota_exhausted:
        print(
            f"Stopped early due to daily quota exhaustion. "
            f"{len(answers_payload)} answers collected before stopping."
        )

    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. Prepare Submission
    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)

    # 5. Submit (with backoff on 429 / transient server errors)
    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
    try:
        response = _request_with_backoff("POST", submit_url, json=submission_data, timeout=60)
        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.')}"
        )
        if quota_exhausted:
            final_status += (
                f"\n\nNote: run stopped early after {len(answers_payload)} "
                f"questions due to daily quota exhaustion."
            )
        print("Submission successful.")
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
    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)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.Timeout:
        status_message = "Submission Failed: The request timed out."
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.RequestException as e:
        status_message = f"Submission Failed: Network error - {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except Exception as e:
        status_message = f"An unexpected error occurred during submission: {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df


# =========================================================
# Gradio Interface
# =========================================================

with gr.Blocks() as demo:
    gr.Markdown("# Web Research 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 the research agent, submit answers, and see the score.
        ---
        **Disclaimers:**
        Once clicking the submit button, it can take some time (the agent must search, read sources, and verify claims for every question).
        """
    )

    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 environment variable 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}")
        print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
    else:
        print("ℹ️  SPACE_ID environment variable not found (running locally?). Repo URL cannot be determined.")

    print("-" * (60 + len(" App Starting ")) + "\n")

    print("Launching Gradio Interface for Web Research Agent Evaluation...")
    demo.launch(debug=True, share=False)