import os
import json
from fastapi import FastAPI, Request, HTTPException, Response
from fastapi.responses import PlainTextResponse
from huggingface_hub import InferenceClient
import io
from PIL import Image

app = FastAPI(title="VK AI Webhook Bot")

# Переменные окружения (добавьте их в Settings → Variables в Space)
VK_TOKEN = os.getenv("VK_TOKEN")               # access_token группы
VK_SECRET = os.getenv("VK_SECRET")             # секретное слово (придумайте сами)
HF_TOKEN = os.getenv("HF_TOKEN")               # ваш HF токен
CONFIRMATION_STRING = os.getenv("VK_CONFIRM")  # строка подтверждения из VK

client = InferenceClient(token=HF_TOKEN)

LLM_MODEL = "Qwen/Qwen2.5-7B-Instruct"
IMG_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"

@app.get("/")
async def root():
    return {"status": "VK Webhook Bot is running"}

@app.post("/webhook")
async def vk_webhook(request: Request):
    try:
        body = await request.json()
    except:
        raise HTTPException(400, "Invalid JSON")

    # Проверка секретного ключа (обязательно!)
    if body.get("secret") != VK_SECRET:
        raise HTTPException(403, "Invalid secret")

    event_type = body.get("type")

    if event_type == "confirmation":
        # VK при настройке шлёт confirmation → возвращаем строку подтверждения
        if CONFIRMATION_STRING:
            return PlainTextResponse(CONFIRMATION_STRING)
        else:
            raise HTTPException(500, "Confirmation string not set")

    if event_type == "message_new":
        msg = body["object"]["message"]
        peer_id = msg["peer_id"]
        text = msg["text"].strip()
        from_id = msg["from_id"]
        msg_id = msg.get("id")  # для reply_to

        low_text = text.lower()

        # Пример обработки команд
        if low_text.startswith("!картинка") or "нарисуй" in low_text:
            prompt = text.replace("!картинка", "").replace("нарисуй", "").strip()
            if not prompt:
                # Отправляем ответ (через vk_api или requests — см. ниже)
                await send_vk_message(peer_id, "Укажите, что нарисовать!", msg_id)
                return {"ok": True}

            await send_vk_message(peer_id, f"Генерирую картинку: «{prompt}»...", msg_id)

            try:
                image = client.text_to_image(
                    prompt,
                    model=IMG_MODEL,
                    num_inference_steps=25,
                    guidance_scale=7.0
                )
                buf = io.BytesIO()
                image.save(buf, format="PNG")
                buf.seek(0)

                # Здесь нужно загрузить фото в ВК и получить attachment
                # Для этого используйте vk_api или requests + upload сервер
                # Простой вариант — пока просто логируем
                print("Картинка сгенерирована")
                await send_vk_message(peer_id, "Картинка готова! (загрузка пока не реализована)", msg_id)

            except Exception as e:
                await send_vk_message(peer_id, f"Ошибка генерации: {str(e)[:100]}", msg_id)

        else:
            # Обычный чат с AI
            try:
                response = client.chat.completions.create(
                    model=LLM_MODEL,
                    messages=[{"role": "user", "content": text}],
                    max_tokens=400
                )
                answer = response.choices[0].message.content.strip()
                await send_vk_message(peer_id, answer, msg_id)
            except Exception as e:
                await send_vk_message(peer_id, "Ошибка ИИ, попробуйте позже", msg_id)

    return {"ok": True}  # VK требует "ok" или 200 OK

async def send_vk_message(peer_id: int, text: str, reply_to: int = None):
    """Отправка сообщения через VK API (простой вариант через requests)"""
    import requests

    params = {
        "peer_id": peer_id,
        "message": text,
        "random_id": 0,
        "access_token": VK_TOKEN,
        "v": "5.199"
    }
    if reply_to:
        params["reply_to"] = reply_to

    try:
        requests.post("https://api.vk.com/method/messages.send", params=params)
    except:
        print("Ошибка отправки сообщения в VK")