import gradio as gr
import requests
import re

PROMPT_SYSTEM = (
    "Ты — виртуальный астролог, отвечай интересно, понятно и позитивно. "
    "Умей давать грамотные гороскопы на день, неделю, месяц, год для каждого знака зодиака, "
    "расскажи о кармических задачах, характеристиках, давай советы для самореализации и отношений. "
    "Расскажи о совместимости двух знаков (например, 'Совместимость Дева и Лев'). "
    "Если спрашивают гороскоп — отвечай только про тот срок (день/неделя/месяц/год), который просят. "
    "Если спрашивают про совместимость — только совместимость этих знаков. "
    "Не давай медицинских или финансовых советов."
)

# --- Распознавание даты по ключевым словам
def detect_timeframe(text):
    if re.search(r"\b(сегодня|на день|на сегодня|день)\b", text, re.IGNORECASE):
        return "день"
    elif re.search(r"\b(на неделю|неделя|неделю)\b", text, re.IGNORECASE):
        return "неделя"
    elif re.search(r"\b(на месяц|месяц|месяца)\b", text, re.IGNORECASE):
        return "месяц"
    elif re.search(r"\b(на год|год|году)\b", text, re.IGNORECASE):
        return "год"
    else:
        return None

# --- Вспомогательная функция извлечения знаков зодиака из текста
ZODIACS = [
    "овен", "телец", "близнец", "рак", "лев", "дева", "весы", "скорпион",
    "стрелец", "козерог", "водолей", "рыбы", "овна", "тельца", "близнецов",
    "рака", "льва", "девы", "весов", "скорпиона", "стрельца", "козерога",
    "водолея", "рыб"
]
def extract_zodiacs(text):
    found = []
    for z in ZODIACS:
        if re.search(rf"\b{z}\b", text, re.IGNORECASE):
            found.append(z.capitalize())
    # Убираем дубликаты
    return list(set(found))

# --- Основная функция общения с HuggingFace Inference API
def astro_assistant(message, history=[]):
    prompt = PROMPT_SYSTEM
    time_type = detect_timeframe(message)
    signs = extract_zodiacs(message)
    if "совместим" in message.lower() or "пара" in message.lower():
        topic = "совместимость"
    elif "карми" in message.lower():
        topic = "карма"
    elif time_type:
        topic = f"гороскоп на {time_type}"
    else:
        topic = "общее астрологическое описание"
    
    chat_history = []
    # Добавим историю беседы (история max 3 сообщения, не обязательно – можно убрать)
    for prev_q, prev_a in history[-3:]:
        chat_history.append({"role": "user", "content": prev_q})
        chat_history.append({"role": "assistant", "content": prev_a})
    # Добавим системный промпт, затем историю, затем новое сообщение пользователя
    messages = (
        [{"role": "system", "content": prompt}]
        + chat_history
        + [{"role": "user", "content": message}]
    )
    payload = {
        "inputs": messages,
        "parameters": {
            "max_new_tokens": 320,
            "top_p": 0.95,
            "temperature": 0.8,
        }
    }
    try:
        response = requests.post(
            "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta",
            headers={},
            json=payload,
            timeout=60
        )
        if response.status_code != 200:
            return f"Ошибка HF API: {response.status_code}\n{response.text}", history
        data = response.json()
        reply = data.get("generated_text","Ответ не получен,попробуйте иначе переформулировать вопрос. ")
        history = history + [[message, reply]]
        return reply, history
    except Exception as e:
        return f"Ошибка: {e}", history

# --- Примеры для подсказок пользователю
example_questions = [
    "Гороскоп на день для Льва",
    "Совместимость Рыбы и Скорпион",
    "Гороскоп на неделю для Тельца",
    "Кармические задачи для Девы",
    "Гороскоп на год для всех знаков",
    "Совет для Овна в любви",
    "Общие черты Льва"
]

# --- Интерфейс Gradio (чата)
with gr.Blocks(title="AI Астролог-ассистент") as demo:
    gr.Markdown(
        """
        # 🌟 AI Астролог-ассистент
        Задавайте любые вопросы: <br>
        • **Гороскоп на день/неделю/месяц/год для любого знака**<br>
        • **Совместимость пары**<br>
        • **Кармические задачи**<br>
        • **Астрологические советы для развития**<br>
        ---
        """
    )
    chatbot = gr.Chatbot(label="Диалог с AI-астрологом", height=400)
    with gr.Row():
        msg = gr.Textbox(label="Ваш вопрос", placeholder="Например: Кармические задачи для Весов")
    with gr.Row():
        submit = gr.Button("Отправить")
        clear = gr.Button("Очистить чат")

    examples = gr.Examples(
        examples=example_questions,
        inputs=[msg]
    )

    history_state = gr.State([])

    def handle_msg(user_message, chat_history):
        reply, updated_history = astro_assistant(user_message, chat_history)
        return updated_history + [[user_message, reply]], updated_history

    submit.click(handle_msg, inputs=[msg, history_state], outputs=[chatbot, history_state])
    msg.submit(handle_msg, inputs=[msg, history_state], outputs=[chatbot, history_state])
    clear.click(lambda: ([], []), outputs=[chatbot, history_state])

demo.launch()