import dotenv
import importlib.resources
import json
import os
from typing import Any

import gradio as gr
import requests
import pandas as pd
from pathlib import Path
from langchain_core.messages import HumanMessage
from langgraph.graph import MessagesState
from langgraph.graph.graph import CompiledGraph

from agent_factory import AgentFactory

dotenv.load_dotenv()

HF_ACCESS_TOKEN = os.getenv("HF_ACCESS_TOKEN")
HF_USERNAME = os.getenv("HF_USERNAME")

# (Keep Constants as is)
# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

DATA_PATH = Path(str(importlib.resources.files("data")))
QUESTIONS_FILE_PATH = DATA_PATH / "questions.jsonl"
AGENT_ANSWERS_FILE_PATH = DATA_PATH / "agent-answers.jsonl"


# --- Basic Agent Definition ---
# ----- THIS IS WHERE YOU CAN BUILD WHAT YOU WANT ------
class BasicAgent:

    __agent_factory: AgentFactory
    __agent: CompiledGraph

    def __init__(self):
        self.__agent_factory = AgentFactory()
        self.__agent = self.__agent_factory.get()
        print("BasicAgent initialized.")

    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")

        initial_state = MessagesState(
            messages=[
                self.__agent_factory.system_prompt,
                HumanMessage(content=question)
            ]
        )

        final_state = self.__agent.invoke(input=initial_state)

        answer = final_state["messages"][-1].content

        print(f"Agent returning answer: {answer}")
        return answer


def retrieve_downloaded_questions() -> list[dict[str, Any]]:
    with open(QUESTIONS_FILE_PATH, mode="r") as f:
        return [json.loads(line) for line in f]

def download_questions_and_files() -> list[dict[str, Any]]:
    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    files_base_url = f"{api_url}/files"

    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 [{
                "error": "Fetched questions list is empty or invalid format."
            }]
        print(f"Fetched {len(questions_data)} questions.")
    except requests.exceptions.RequestException as e:
        print(f"Error fetching questions: {e}")
        return [{
            "error": f"Error fetching questions: {e}"
        }]
    except requests.exceptions.JSONDecodeError as e:
        print(f"Error decoding JSON response from questions endpoint: {e}")
        print(f"Response text: {response.text[:500]}")
        return [{
            "error": f"Error decoding server response for questions: {e}"
        }]
    except Exception as e:
        print(f"An unexpected error occurred fetching questions: {e}")
        return [{
            "error": f"An unexpected error occurred fetching questions: {e}"
        }]

    # Save input questions and related files into the data subdirectory
    try:
        with open(QUESTIONS_FILE_PATH, mode="w") as f:
            for cur_question in questions_data:
                json.dump(cur_question, f)
                f.write("\n")

                file_name = cur_question["file_name"]
                if len(file_name) > 0:
                    file_url = f"{files_base_url}/{cur_question["task_id"]}"
                    response = requests.get(file_url)
                    out_file_path = DATA_PATH / file_name
                    with open(out_file_path, 'wb') as file:
                        file.write(response.content)
    except requests.exceptions.RequestException as e:
        print(f"Error fetching question-related file: {e}")
        return [{
            "error": f"Error fetching question-related file: {e}"
        }]
    except Exception as e:
        print(f"An unexpected error occurred fetching question-related file: {e}")
        return [{
            "error": f"An unexpected error occurred fetching question-related file: {e}"
        }]

    return questions_data


def create_answers_file_if_not_exists() -> None:
    if not os.path.exists(AGENT_ANSWERS_FILE_PATH):
        with open(AGENT_ANSWERS_FILE_PATH, 'w'):
            pass


def get_answers_payload() -> list[dict[str, Any]]:
    with open(AGENT_ANSWERS_FILE_PATH, mode="r") as f:
        answers_payload = [json.loads(line) for line in f]
    return answers_payload


def get_task_ids_to_process() -> list[str]:
    with open(QUESTIONS_FILE_PATH, mode="r") as f:
        all_tasks = set([json.loads(line)["task_id"] for line in f])
    answers = get_answers_payload()
    answered_tasks = set([answer["task_id"] for answer in answers])
    tasks_to_answer = all_tasks - answered_tasks
    return list(tasks_to_answer)


def run_and_submit_all() -> tuple[str, pd.DataFrame | 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") # Get the SPACE_ID for sending link to the code

    username= f"{HF_USERNAME}"
    print(f"User: {username}")

    api_url = DEFAULT_API_URL
    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
    # In the case of an app running as a hugging Face space, this link points
    # towards your codebase ( useful for others so please keep it public)
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(agent_code)

    # 2. Fetch Questions and related files (they get saved into the data directory)
    if os.path.exists(QUESTIONS_FILE_PATH):
        questions_data = retrieve_downloaded_questions()
    else:
        questions_data = download_questions_and_files()

    # 3. Run your Agent and save agent's answers for later review
    create_answers_file_if_not_exists()
    task_ids_to_process = get_task_ids_to_process()
    results_log = []
    print(f"Running agent on {len(questions_data)} questions...")
    with open(AGENT_ANSWERS_FILE_PATH, mode="a") as f:
        for item in questions_data:
            task_id = item.get("task_id")
            if task_id not in task_ids_to_process:
                print(f"Skipping already answered question: {item}")
                continue
            question_text = json.dumps(item)
            if not task_id or question_text is None:
                print(f"Skipping item with missing task_id or question: {item}")
                continue
            try:
                answer_to_submit = agent(question_text)
                answer_payload = {"task_id": task_id, "submitted_answer": answer_to_submit}
                json.dump(answer_payload, f)
                f.write("\n")
                f.flush()
                results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": answer_to_submit})
            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}"})

    answers_payload = get_answers_payload()
    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)

    if len(answers_payload) < len(questions_data):
        msg = "Still need to process all the questions. Rerun until all questions are answered."
        print(msg)
        return msg, 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}")
    headers = {
        "Authorization": f"Bearer {HF_ACCESS_TOKEN}",
        "Content-Type": "application/json"
    }
    try:
        response = requests.post(
            submit_url,
            json=submission_data,
            headers=headers,
            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.  Read the `README.md` file for configuration.
        2.  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).
        """
    )

    # 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)