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

from huggingface_hub import InferenceClient
from pypdf import PdfReader

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-14B-Instruct")
HF_TOKEN = os.getenv("HF_TOKEN")


def clean_answer(text: str) -> str:
    if not text:
        return ""

    text = text.strip()

    # remove markdown fences
    text = re.sub(r"^```.*?\n", "", text, flags=re.DOTALL)
    text = text.replace("```", "").strip()

    # common prefixes
    text = re.sub(r"(?i)^final answer\s*:\s*", "", text).strip()
    text = re.sub(r"(?i)^answer\s*:\s*", "", text).strip()
    text = re.sub(r"(?i)^submitted_answer\s*:\s*", "", text).strip()

    # if model gave multiple lines, keep the first meaningful one
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    if lines:
        text = lines[0]

    # trim wrapping quotes
    text = text.strip().strip('"').strip("'").strip()

    return text


def try_extract_text_from_pdf(content: bytes) -> str:
    try:
        reader = PdfReader(io.BytesIO(content))
        pages = []
        for page in reader.pages[:10]:
            page_text = page.extract_text() or ""
            if page_text.strip():
                pages.append(page_text)
        return "\n".join(pages)[:12000]
    except Exception:
        return ""


def try_extract_text_from_bytes(content: bytes) -> str:
    for enc in ["utf-8", "latin-1"]:
        try:
            text = content.decode(enc, errors="ignore").strip()
            if text:
                return text[:12000]
        except Exception:
            pass
    return ""


def fetch_task_file_text(task_id: str) -> str:
    file_url = f"{DEFAULT_API_URL}/files/{task_id}"
    try:
        r = requests.get(file_url, timeout=30)
        if r.status_code != 200:
            return ""

        content_type = (r.headers.get("content-type") or "").lower()
        content = r.content

        if "pdf" in content_type:
            pdf_text = try_extract_text_from_pdf(content)
            if pdf_text:
                return pdf_text

        if any(x in content_type for x in ["text", "json", "csv", "xml", "html"]):
            return try_extract_text_from_bytes(content)

        # fallback: try text anyway
        return try_extract_text_from_bytes(content)

    except Exception:
        return ""


class BasicAgent:
    def __init__(self):
        if not HF_TOKEN:
            raise ValueError("Missing HF_TOKEN secret in your Space settings.")
        self.client = InferenceClient(token=HF_TOKEN)
        print(f"BasicAgent initialized with model: {MODEL_ID}")

    def __call__(self, question: str, file_text: str = "") -> str:
        system_prompt = (
            "You solve benchmark questions. "
            "Return only the exact final answer. "
            "Do not explain. "
            "Do not use markdown. "
            "Do not say FINAL ANSWER. "
            "If the answer is a number, date, name, or short phrase, return exactly that."
        )

        user_prompt = f"Question:\n{question}\n"
        if file_text.strip():
            user_prompt += f"\nAttached file content:\n{file_text}\n"

        completion = self.client.chat.completions.create(
            model=MODEL_ID,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.1,
            max_tokens=120,
        )

        raw = completion.choices[0].message.content
        answer = clean_answer(raw)
        print(f"RAW MODEL OUTPUT: {raw}")
        print(f"CLEANED ANSWER: {answer}")
        return answer


def run_random_test():
    random_url = f"{DEFAULT_API_URL}/random-question"

    try:
        agent = BasicAgent()
    except Exception as e:
        return f"Agent init error: {e}", None

    try:
        r = requests.get(random_url, timeout=20)
        r.raise_for_status()
        item = r.json()
    except Exception as e:
        return f"Could not fetch random question: {e}", None

    task_id = item.get("task_id", "")
    question = item.get("question", "")
    file_text = fetch_task_file_text(task_id) if task_id else ""

    try:
        answer = agent(question, file_text=file_text)
    except Exception as e:
        return f"Agent failed on random test: {e}", None

    preview = pd.DataFrame([
        {
            "Task ID": task_id,
            "Question": question,
            "Attached File Text Found": "yes" if file_text else "no",
            "Submitted Answer": answer,
        }
    ])

    return "Random test completed. Check whether the answer is short and clean.", preview


def run_and_submit_all(profile: gr.OAuthProfile | None):
    space_id = os.getenv("SPACE_ID")

    if profile:
        username = f"{profile.username}"
    else:
        return "Please login to Hugging Face first.", None

    if not space_id:
        return "SPACE_ID environment variable missing.", None

    questions_url = f"{DEFAULT_API_URL}/questions"
    submit_url = f"{DEFAULT_API_URL}/submit"

    try:
        agent = BasicAgent()
    except Exception as e:
        return f"Error initializing agent: {e}", None

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    try:
        response = requests.get(questions_url, timeout=20)
        response.raise_for_status()
        questions_data = response.json()
    except Exception as e:
        return f"Error fetching questions: {e}", None

    results_log = []
    answers_payload = []

    for item in questions_data:
        task_id = item.get("task_id")
        question_text = item.get("question", "")

        if not task_id or not question_text:
            continue

        try:
            file_text = fetch_task_file_text(task_id)
            submitted_answer = agent(question_text, file_text=file_text)

            answers_payload.append(
                {"task_id": task_id, "submitted_answer": submitted_answer}
            )

            results_log.append(
                {
                    "Task ID": task_id,
                    "Question": question_text,
                    "Attached File Text Found": "yes" if file_text else "no",
                    "Submitted Answer": submitted_answer,
                }
            )
        except Exception as e:
            results_log.append(
                {
                    "Task ID": task_id,
                    "Question": question_text,
                    "Attached File Text Found": "unknown",
                    "Submitted Answer": f"AGENT ERROR: {e}",
                }
            )

    if not answers_payload:
        return "No answers were produced.", pd.DataFrame(results_log)

    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload,
    }

    try:
        response = requests.post(submit_url, json=submission_data, timeout=120)
        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.')}"
        )

        return final_status, pd.DataFrame(results_log)

    except requests.exceptions.HTTPError as e:
        detail = f"Server responded with status {e.response.status_code}."
        try:
            detail_json = e.response.json()
            detail += f" Detail: {detail_json.get('detail', e.response.text)}"
        except Exception:
            detail += f" Response: {e.response.text[:500]}"
        return f"Submission failed: {detail}", pd.DataFrame(results_log)

    except Exception as e:
        return f"Submission failed: {e}", pd.DataFrame(results_log)


with gr.Blocks() as demo:
    gr.Markdown("# Unit 4 Cheap Baseline Agent")
    gr.Markdown(
        """
        1. Add your HF_TOKEN secret in Space settings.
        2. Login with Hugging Face below.
        3. Click 'Run One Cheap Test' first.
        4. If the answer looks clean, click 'Run Full Evaluation and Submit'.

        Notes:
        - This version is optimized for simplicity and low cost.
        - It tries to read attached text/PDF files.
        - It returns short exact answers for exact-match scoring.
        """
    )

    gr.LoginButton()

    test_button = gr.Button("Run One Cheap Test")
    run_button = gr.Button("Run Full Evaluation and Submit")

    status_output = gr.Textbox(label="Status", lines=6, interactive=False)
    results_table = gr.DataFrame(label="Agent Output", wrap=True)

    test_button.click(
        fn=run_random_test,
        outputs=[status_output, results_table],
    )

    run_button.click(
        fn=run_and_submit_all,
        outputs=[status_output, results_table],
    )

if __name__ == "__main__":
    demo.launch(debug=True, share=False)