import os
import re
import gradio as gr
import requests
import pandas as pd

# smolagents: HF가 만든 에이전트 프레임워크. CodeAgent는 LLM이 매 스텝마다 파이썬
# 코드를 생성·실행해 도구를 호출하는 ReAct 변형이다.
from smolagents import CodeAgent, InferenceClientModel

# 도구는 tools/ 패키지에 분리되어 있다. 각 파일이 하나의 @tool 함수를 담당.
from tools import (
    web_search,
    visit_webpage,
    wikipedia_search,
    youtube_info,
    exec_python_code,
    get_attached_file,
    prefetch_question_index,
    set_question_index,
    set_current_task,
)
# GAIA exact-match 채점에 맞춘 시스템 프롬프트 가이드라인.
from prompts import GAIA_ANSWER_GUIDELINES
# 멀티홉 질문 사전 분해(query decomposition).
from decomposer import decompose_question
# 답변 캐싱(재실행 시 처리한 문제 스킵, 한 문제 실패의 cascade 방지).
from answer_cache import load_cache, save_answer, is_retryable_answer
# 답변 포맷 후처리(exact-match 채점 보정).
from formatter import coerce_answer, final_format_pass

# (Keep Constants as is)
# --- Constants ---
# 채점 서버 베이스 URL. /questions 로 문제를 받고, /submit 으로 답을 제출한다.
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"


# --- Basic Agent Definition ---
# ----- THIS IS WHERE YOU CAN BUILD WHAT YOU WANT ------
class BasicAgent:
    """GAIA Level 1 문제를 푸는 에이전트.

    실제 추론은 smolagents.CodeAgent에 위임한다. CodeAgent는 매 스텝마다
    LLM(InferenceClientModel)에 컨텍스트를 보내 파이썬 코드를 받아오고,
    그 코드를 안전한 샌드박스에서 실행해 도구 결과를 다시 LLM에게 전달한다.
    최종적으로 LLM이 final_answer(...)를 호출하면 그 값이 self.agent.run의 반환값이 된다.
    """

    def __init__(self):
        print("BasicAgent initialized.")
        # 모델: Qwen2.5-72B-Instruct (오픈웨이트, 32k ctx).
        # provider="hf-inference"로 명시 — HF 네이티브 serverless 라인이라 무료 풀.
        # 시도했던 다른 모델들의 결과:
        # - DeepSeek-V3 + provider="auto" → Together로 라우팅 → 503/402 (크레딧 소진)
        # - Llama-3.3-70B-Instruct + provider="hf-inference" → 400 Bad request
        #   (hf-inference에 호스팅 안 됨, paid provider 전용)
        # Qwen2.5-72B는 hf-inference에서 호스팅이 확인된 모델 중 추론력 가장 강함.
        # 큐 대기 가끔 있어도 키 정책 + 무료 제약에서는 최선의 선택.
        # 코더 모델로 바꾸지 말 것: 매 스텝 마크다운 잔재(```, </code])를 흘려
        # smolagents 코드 파서가 깨진다(이전 32B 코더 시도에서 확인됨).
        self.model = InferenceClientModel(
            model_id="Qwen/Qwen2.5-72B-Instruct",
            provider="auto",
            max_tokens=2048,  # 한 스텝당 LLM 응답 토큰 한도
        )

        # /questions 한 번 prefetch 해서 {질문본문: task_id} 인덱스 빌드.
        # tools.attachments 모듈 전역에 등록 → __call__ 진입 시 set_current_task()가 사용.
        idx = prefetch_question_index()
        set_question_index(idx)
        print(f"Prefetched question index: {len(idx)} entries")

        # 도구 6종: web_search, visit_webpage, wikipedia_search, youtube_info,
        # exec_python_code, get_attached_file.
        # max_steps=12: 8스텝에선 검색 실패로 다른 쿼리를 시도하다 한도에 걸리는 일이 잦았다.
        # additional_authorized_imports: 샌드박스에서 표 처리/계산이 필요할 때 import 허용.
        self.agent = CodeAgent(
            tools=[
                web_search,
                visit_webpage,
                wikipedia_search,
                youtube_info,
                exec_python_code,
                get_attached_file,
            ],
            model=self.model,
            max_steps=12,
            additional_authorized_imports=[
                "pandas", "openpyxl", "json", "re", "math", "statistics", "itertools",
                "datetime", "collections", "urllib.parse",
            ],
        )
        # CodeAgent의 기본 시스템 프롬프트 뒤에 GAIA용 채점 규칙을 덧붙인다.
        try:
            current_sp = self.agent.prompt_templates.get("system_prompt", "")
            self.agent.prompt_templates["system_prompt"] = (
                current_sp + "\n\n" + GAIA_ANSWER_GUIDELINES
            )
        except Exception as e:
            print(f"Warning: could not patch system prompt: {e}")

    def __call__(self, question: str) -> str:
        # 시그니처는 (self, question: str) -> str로 고정. run_and_submit_all이
        # `agent(question_text)` 형태로 호출하므로 인자 추가 금지.
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        # 현재 문제의 task_id를 tools.attachments 전역에 세팅 → get_attached_file() 가
        # 인자 없이 동작. 매칭 실패 시 None(첨부 없는 문제처럼 처리됨).
        tid = set_current_task(question)
        if tid:
            print(f"  → matched task_id: {tid}")
        else:
            print("  → no matched task_id (question not in cache)")

        # 멀티홉 질문은 1콜로 plan을 뽑아 prompt에 prepend 한다. 본 루프(12스텝)가
        # 첫 스텝부터 곧장 도구 호출로 들어가도록 유도. 단일 lookup이면 None이
        # 반환되어 원본 질문 그대로 진행. 분해 실패도 None → degrade 안전.
        plan = decompose_question(question)
        if plan:
            print(f"  → decomposition plan:\n{plan}")
            prompt_question = (
                f"{question}\n\n"
                f"--- Suggested decomposition plan (guidance — deviate as tool results show) ---\n"
                f"{plan}\n"
                f"--- end plan ---\n"
                f"The final answer must address the ORIGINAL question above, not the plan."
            )
        else:
            prompt_question = question

        try:
            raw = self.agent.run(prompt_question)
            answer = str(raw).strip()
            # 1) "FINAL ANSWER:" / "FINAL ANSWER -" 같은 prefix 제거(case-insensitive).
            answer = re.sub(
                r"^\s*FINAL\s*ANSWER\s*[:\-]?\s*",
                "",
                answer,
                flags=re.IGNORECASE,
            ).strip()
            # 2) 양끝을 둘러싼 따옴표 제거. (LLM이 종종 "Answer" 형태로 따옴표를 붙인다.)
            if len(answer) >= 2 and (
                (answer[0] == '"' and answer[-1] == '"')
                or (answer[0] == "'" and answer[-1] == "'")
            ):
                answer = answer[1:-1].strip()
            # 3) Final-answer formatter pass — 별도 LLM 호출로 GAIA 포맷 강제.
            #    내용 맞고 형식 위반인 B 카테고리 회복용. 호출 실패 시 raw 유지(graceful degrade).
            answer = final_format_pass(question, answer)
            # 4) 결정적 regex 후처리(yes/no, 숫자, 통화). final_format_pass가 놓친 패턴 안전망.
            answer = coerce_answer(question, answer)
            print(f"Agent returning answer: {answer}")
            return answer
        except Exception as e:
            # 한 문제에서 raise되면 전체 채점이 멈추므로 여기서 흡수하고
            # AGENT_ERROR 문자열을 답으로 제출한다(어차피 오답 처리됨).
            # 제출 문자열은 타입만 노출(상세는 로그에만) — 예외 메시지 유출 완화.
            import traceback
            err_type = type(e).__name__
            print(f"Agent error ({err_type}): {e}")
            print(traceback.format_exc()[-600:])
            return f"AGENT_ERROR: {err_type}"

def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the BasicAgent on them, submits all answers,
    and displays the results.
    """
    # --- Determine HF Space Runtime URL and Repo URL ---
    space_id = os.getenv("SPACE_ID")  # Space 배포 시 자동 설정; 로컬에서는 보통 없음

    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 ( modify this part to create your agent)
    try:
        agent = BasicAgent()
    except Exception as e:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None
    # SPACE_ID 없으면 /spaces/None/... 로 깨지지 않도록 고정 문서 URL 사용.
    if space_id:
        agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    else:
        agent_code = "https://huggingface.co/docs/hub/spaces"
        print("SPACE_ID unset — using docs URL for agent_code (set when deploying to HF Spaces).")
    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 your Agent
    results_log = []
    answers_payload = []
    # 캐시는 .cache/answers.json. 한 번 답을 받은 task_id는 재실행 시 LLM 호출
    # 없이 그대로 재사용 — 전체 채점 재시도 비용 절감 + 한 문제 실패 cascade 방지.
    cache = load_cache()
    print(f"Running agent on {len(questions_data)} questions... (cache: {len(cache)} entries)")
    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

        cached = cache.get(task_id)
        if cached and isinstance(cached, dict) and "answer" in cached:
            submitted_answer = cached["answer"]
            if is_retryable_answer(submitted_answer):
                print(
                    f"  [cache stale] task_id={task_id}: retrying "
                    f"instead of reusing {submitted_answer!r}"
                )
            else:
                print(f"  [cache hit] task_id={task_id}: {submitted_answer[:80]}")
                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})
                continue
        try:
            submitted_answer = agent(question_text)
            # AGENT_ERROR 결과는 save_answer 내부에서 캐시 안 함(다음 실행 때 재시도).
            save_answer(task_id, question_text, submitted_answer)
            if is_retryable_answer(submitted_answer):
                print(f"  [skip retryable answer] task_id={task_id}: {submitted_answer!r}")
                results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
                continue
            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:
            err_type = type(e).__name__
            print(f"Error running agent on task {task_id} ({err_type}): {e}")
            results_log.append(
                {
                    "Task ID": task_id,
                    "Question": question_text,
                    "Submitted Answer": f"AGENT_ERROR: {err_type}",
                }
            )

    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
    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.")
        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


# --- Build Gradio Interface using Blocks ---
with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner")
    gr.Markdown(
        """
        **Instructions:**

        1.  Please clone this space, then modify the code to define your agent's logic, the tools, the necessary packages, etc ...
        2.  Log in to your Hugging Face account using the button below. This uses your HF username for submission.
        3.  Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.

        ---
        **Disclaimers:**
        Once clicking on the "submit button, it can take quite some time ( this is the time for the agent to go through all the questions).
        This space provides a basic setup and is intentionally sub-optimal to encourage you to develop your own, more robust solution. For instance for the delay process of the submit button, a solution could be to cache the answers and submit in a seperate action or even to answer the questions in async.
        """
    )

    gr.LoginButton()

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

    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    # Removed max_rows=10 from DataFrame constructor
    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)
    # Check for SPACE_HOST and SPACE_ID at startup for information
    space_host_startup = os.getenv("SPACE_HOST")
    space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup

    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 repo URLs if SPACE_ID is found
        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 Basic Agent Evaluation...")
    demo.launch(debug=True, share=False)