import streamlit as st
import time
from transformers import pipeline, PipelineException

# --- Конфігурація ---
MODEL_NAME = "facebook/bart-large-mnli" # Ця модель добре працює для zero-shot
CONFIDENCE_THRESHOLD = 0.5 # Мінімальний рівень впевненості для визначення наміру

# --- Завантаження моделі ---
try:
    # Використовуємо zero-shot classification для визначення наміру
    intent_classifier = pipeline("zero-shot-classification", model=MODEL_NAME)
    st.sidebar.success(f"Модель '{MODEL_NAME}' завантажена.")
except Exception as e:
    st.error(f"Помилка завантаження моделі '{MODEL_NAME}': {e}")
    st.warning("Застосунок не може працювати без моделі.")
    # Зупиняємо виконання застосунку, якщо модель не завантажилась
    st.stop()

# --- Словник стратегій ---
# Ключі - це потенційні наміри (мітки для zero-shot)
методи_для_намірів = {
    "пошук інформації": {
        "розбиття": lambda задача: [задача], # Для простоти, поки не розбиваємо
        "метод": lambda підзадача: "Використовуйте пошукову систему (Google, DuckDuckGo)."
    },
    "написання тексту": {
        "розбиття": lambda задача: [задача], # Для простоти, поки не розбиваємо
        "метод": lambda підзадача: "Використовуйте текстовий редактор або онлайн-документ."
    },
    # Можна додати більше намірів тут
    "переклад тексту": {
        "розбиття": lambda задача: [задача],
        "метод": lambda підзадача: "Використовуйте сервіс онлайн-перекладу (Google Translate, DeepL)."
    }
}

# --- Функції ---
def визначити_намір(задача):
    """Визначає намір користувача за допомогою моделі zero-shot."""
    наміри = list(методи_для_намірів.keys())
    try:
        # Класифікуємо задачу за списком можливих намірів
        результат = intent_classifier(задача, candidate_labels=наміри)

        print(f"Результат класифікації: {результат}") # Для налагодження

        найкращий_намір = результат['labels'][0]
        найвища_впевненість = результат['scores'][0]

        # Перевіряємо, чи впевненість вища за поріг
        if найвища_впевненість >= CONFIDENCE_THRESHOLD:
            return найкращий_намір
        else:
            return "невизначено" # Якщо впевненість занизька

    except PipelineException as e:
         st.error(f"Помилка під час визначення наміру: {e}")
         return "помилка класифікації"
    except Exception as e:
        st.error(f"Неочікувана помилка під час визначення наміру: {e}")
        return "помилка класифікації"


def виконати_підзадачу(підзадача, метод):
    """Імітує виконання підзадачі."""
    with st.spinner(f"  Виконую підзадачу: '{підзадача}'..."):
        st.write(f"  *Метод:* {метод}")
        time.sleep(1.5) # Імітація роботи
    st.success(f"  Підзадачу '{підзадача}' виконано!")

# --- Інтерфейс Streamlit ---
st.title("ZerOnyX: Простий агент для задач")
st.caption("Введіть задачу, і я спробую визначити, що робити.")

задача_користувача = st.text_input("Опишіть вашу задачу:", placeholder="Наприклад: 'знайди інформацію про квантові комп'ютери' або 'напиши короткий опис проєкту'")

if задача_користувача:
    # Визначаємо намір
    намір = визначити_намір(задача_користувача)

    if намір != "помилка класифікації":
        st.subheader(f"Визначений намір: {намір.capitalize()}")

        if намір in методи_для_намірів:
            стратегія = методи_для_намірів[намір]
            підзадачі = стратегія["розбиття"](задача_користувача) # Зараз просто повертає початкову задачу

            st.subheader("План виконання:")
            if len(підзадачі) > 1:
                 st.write("Задача розбита на підзадачі:")
            for i, підзадача in enumerate(підзадачі):
                st.markdown(f"**Крок {i+1}:** {підзадача}")
                метод = стратегія["метод"](підзадача)
                виконати_підзадачу(підзадача, метод)
            st.balloons() # Святкуємо успішне виконання :)
        elif намір == "невизначено":
             st.warning(f"Не вдалося достатньо впевнено визначити намір для задачі: '{задача_користувача}'. Спробуйте переформулювати.")
             st.info(f"Я краще розумію задачі, пов'язані з: {', '.join(методи_для_намірів.keys())}.")
        # else: # Цей випадок тепер обробляється через "невизначено" або "помилка класифікації"
        #     st.warning(f"Намір '{намір}' не вдалося обробити.")