import os
import gradio as gr
import requests
import pandas as pd
from pathlib import Path

from smolagents import (
    CodeAgent,
    InferenceClientModel,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    PythonInterpreterTool,
    tool
)

from pypdf import PdfReader

# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
MODEL_ID = "Qwen/Qwen2.5-32B-Instruct"

class BasicAgent:
    def __init__(self):        
        self.model = InferenceClientModel(
            model_id=MODEL_ID,
            token=os.getenv("HF_TOKEN"),
            temperature=0.05,
            max_tokens=1024,
        )

        tools = [
            DuckDuckGoSearchTool(max_results=10),
            VisitWebpageTool(),
            PythonInterpreterTool(),
        ]

        @tool
        def download_file(url: str) -> str:
            """
            Скачивает файл по URL.

            Args:
                url (str): URL файла

            Returns:
                str: Путь или ошибка
            """
            try:
                downloads = Path("./downloads")
                downloads.mkdir(exist_ok=True)
                fname = url.split("/")[-1].split("?")[0] or "file"
                path = downloads / fname
                
                r = requests.get(url, stream=True, timeout=45)
                r.raise_for_status()
                with open(path, "wb") as f:
                    for chunk in r.iter_content(8192):
                        f.write(chunk)
                return f"Скачано: {path.absolute()}. Анализируй."
            except Exception as e:
                return f"Ошибка скачивания: {str(e)}"

        @tool
        def read_pdf(path: str) -> str:
            """
            Читает PDF.

            Args:
                path (str): Путь

            Returns:
                str: Текст (до 4000 символов)
            """
            try:
                reader = PdfReader(path)
                text = "\n".join(page.extract_text() or "" for page in reader.pages)
                return text[:4000]
            except Exception as e:
                return f"Ошибка PDF: {str(e)}"

        @tool
        def read_excel(path: str, sheet: str = None) -> str:
            """
            Читает Excel.

            Args:
                path (str): Путь
                sheet (str, optional): Лист

            Returns:
                str: Таблица или ошибка
            """
            try:
                df = pd.read_excel(path, sheet_name=sheet)
                return df.to_string(max_rows=20, max_cols=10)
            except Exception as e:
                return f"Ошибка Excel: {str(e)}"

        tools.extend([download_file, read_pdf, read_excel])

        self.agent = CodeAgent(
            tools=tools,
            model=self.model,
            add_base_tools=True,
            max_steps=18,
        )
        print("Агент готов!")

    def __call__(self, question: str) -> str:
        print(f"Вопрос: {question[:120]}...")

        if len(question) > 2500:
            question = question[:2500] + "\n[Обрезано из-за длины.]"

        q = question.lower()
        if any(k in q for k in [".mp3", "audio", "recording", "voice", "youtube.com", "video", "attached", "file", "excel", "pdf", "image"]):
            question += "\nЕсли есть URL — скачай файл/видео и читай/ищи транскрипт/описание. Не придумывай."

        try:
            result = self.agent.run(question)
            answer = str(result).strip()

            prefixes = [
                "Final Answer", "Final answer", "Answer:", "The answer is",
                "So the final answer is", "```", "boxed{", "}", "[/INST]", "</s>",
                "Thought:", "Observation:", "Action:"
            ]
            for p in prefixes:
                if p.lower() in answer.lower():
                    answer = answer.split(p, 1)[-1].strip(": []{}\n`")
                    break

            if answer.startswith("[") and answer.endswith("]"):
                answer = answer[1:-1].strip()

            answer = answer.strip()

            if len(answer) > 300 or "придум" in answer.lower():
                answer = answer[:150] + "..." if len(answer) > 150 else answer

            print(f"Ответ: {answer[:150]}...")
            return answer or "Нет ответа"

        except Exception as e:
            err = f"Ошибка: {str(e)[:200]}"
            print(err)
            return err

# --- run_and_submit_all (без изменений) ---
def run_and_submit_all(profile: gr.OAuthProfile | None):
    space_id = os.getenv("SPACE_ID")
    if profile:
        username = profile.username
        print(f"Вход: {username}")
    else:
        return "Войдите в HF", None

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

    try:
        agent = BasicAgent()
    except Exception as e:
        return f"Ошибка агента: {e}", None

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

    try:
        resp = requests.get(questions_url, timeout=15)
        resp.raise_for_status()
        questions = resp.json()
        if not questions:
            return "Вопросов нет", None
        print(f"Вопросов: {len(questions)}")
    except Exception as e:
        return f"Ошибка вопросов: {e}", None

    results = []
    payload = []
    for item in questions:
        tid = item.get("task_id")
        q = item.get("question")
        if not tid or not q:
            continue
        try:
            ans = agent(q)
            payload.append({"task_id": tid, "submitted_answer": ans})
            results.append({"Task ID": tid, "Question": q, "Answer": ans})
        except Exception as e:
            results.append({"Task ID": tid, "Question": q, "Answer": f"ERROR: {e}"})

    if not payload:
        return "Нет ответов", pd.DataFrame(results)

    data = {"username": username.strip(), "agent_code": agent_code, "answers": payload}

    try:
        resp = requests.post(submit_url, json=data, timeout=60)
        resp.raise_for_status()
        res = resp.json()
        status = (
            f"Успех!\n"
            f"Пользователь: {res.get('username')}\n"
            f"Балл: {res.get('score', 'N/A')}% "
            f"({res.get('correct_count', '?')}/{res.get('total_attempted', '?')})\n"
            f"{res.get('message', '')}"
        )
        return status, pd.DataFrame(results)
    except Exception as e:
        return f"Ошибка отправки: {e}", pd.DataFrame(results)

# --- Gradio ---
with gr.Blocks() as demo:
    gr.Markdown("# Агент для финального задания")
    gr.Markdown("""
1. Клонируй и дорабатывай.
2. Войди через кнопку.
3. Нажми кнопку — увидишь score.
""")
    gr.LoginButton()
    btn = gr.Button("Запустить оценку и отправить")
    status = gr.Textbox(label="Результат", lines=6)
    table = gr.DataFrame(label="Ответы", wrap=True)
    btn.click(run_and_submit_all, outputs=[status, table])

if __name__ == "__main__":
    print("Запуск...")
    demo.launch(debug=True, share=False)