import asyncio
import aiohttp
import os
import sys
import logging
import base64
from datetime import datetime
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import (
    Application,
    CommandHandler,
    MessageHandler,
    CallbackQueryHandler,
    ConversationHandler,
    filters,
    ContextTypes
)
from telegram.request import HTTPXRequest
import httpx

from config import TELEGRAM_BOT_TOKEN, PIApi_API_KEY, PIApi_BASE_URL, PROXY_URL, PROXY_TYPE
from database import init_db, save_user, save_image_record, update_image_record, get_user_images
import aiosqlite

# Настройка логирования
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

# Состояния для ConversationHandler
AWAITING_PROMPT = 1

# Создаем директории
UPLOAD_DIR = "uploads"
PROCESSED_DIR = "processed"
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(PROCESSED_DIR, exist_ok=True)

# Глобальная сессия
session = None


async def get_session():
    """Создание aiohttp сессии"""
    global session
    if session is None:
        connector = None
        if PROXY_URL:
            # Создаем коннектор с прокси для aiohttp
            from aiohttp_socks import ProxyConnector
            try:
                connector = ProxyConnector.from_url(PROXY_URL)
                logger.info(f"🔌 Создан SOCKS коннектор для aiohttp")
            except:
                connector = aiohttp.TCPConnector()
        else:
            connector = aiohttp.TCPConnector()

        session = aiohttp.ClientSession(connector=connector)
    return session


async def close_session():
    """Закрытие aiohttp сессии"""
    global session
    if session:
        await session.close()
        session = None


def get_telegram_request():
    """Правильная настройка прокси для Telegram"""
    if PROXY_URL:
        logger.info(f"🔌 Настройка прокси для Telegram: {PROXY_URL}")
        try:
            # Используем монтированный транспорт
            transport = httpx.AsyncHTTPTransport(proxy=PROXY_URL)
            request = HTTPXRequest(transport=transport)
            logger.info("✅ Прокси настроен успешно")
            return request
        except Exception as e:
            logger.error(f"Ошибка настройки прокси: {e}")
            # Возвращаем обычный запрос без прокси
            return HTTPXRequest()
    else:
        logger.info("🔌 Прокси не используется")
        return HTTPXRequest()


async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Обработчик команды /start"""
    logger.info(f"✅ Получена команда /start от {update.effective_user.id}")
    user = update.effective_user

    await save_user(
        user_id=user.id,
        username=user.username,
        first_name=user.first_name,
        last_name=user.last_name
    )

    welcome_text = (
        f"👋 Привет, {user.first_name}!\n\n"
        f"Я бот для обработки фотографий через нейросеть.\n\n"
        f"📸 Просто отправь мне фото и напиши, как его изменить!\n\n"
        f"🔧 Команды:\n"
        f"/start - Приветствие\n"
        f"/test - Проверка работы\n"
        f"/help - Помощь\n"
        f"/history - История обработок"
    )
    await update.message.reply_text(welcome_text)


async def test_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Тестовая команда"""
    logger.info(f"✅ Получена команда /test от {update.effective_user.id}")
    await update.message.reply_text("✅ Бот работает! Команды обрабатываются нормально.")


async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Обработчик команды /help"""
    logger.info(f"✅ Получена команда /help от {update.effective_user.id}")
    help_text = (
        "📖 **Инструкция:**\n\n"
        "1️⃣ Отправь фото\n"
        "2️⃣ Напиши текстовое описание изменений\n"
        "3️⃣ Получи результат!\n\n"
        "**Примеры промптов:**\n"
        "- 'сделай фон пляжем'\n"
        "- 'добавь солнечные очки'\n"
        "- 'преврати в рисунок акварелью'\n\n"
        "📊 /history - История\n"
        "🔧 /test - Проверка"
    )
    await update.message.reply_text(help_text, parse_mode='Markdown')


async def handle_photo(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Обработчик получения фото"""
    logger.info(f"✅ Получено фото от {update.effective_user.id}")
    user = update.effective_user

    photo_file = await update.message.photo[-1].get_file()
    file_id = photo_file.file_id

    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    local_filename = f"user_{user.id}_{timestamp}.jpg"
    local_path = os.path.join(UPLOAD_DIR, local_filename)

    await photo_file.download_to_drive(local_path)

    image_id = await save_image_record(
        user_id=user.id,
        original_file_id=file_id,
        original_url=local_path
    )

    context.user_data['current_image_id'] = image_id
    context.user_data['current_file_path'] = local_path

    await update.message.reply_text(
        "📸 Фото получено!\n\n"
        "Теперь напиши, как изменить это фото.\n"
        "❌ Для отмены отправь /cancel"
    )

    return AWAITING_PROMPT


async def cancel(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Отмена операции"""
    logger.info(f"Отмена операции от {update.effective_user.id}")
    if 'current_file_path' in context.user_data:
        file_path = context.user_data['current_file_path']
        if os.path.exists(file_path):
            os.remove(file_path)

    await update.message.reply_text("Операция отменена.")
    context.user_data.clear()
    return ConversationHandler.END


async def process_image_with_prompt(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Обработка изображения через PiAPI"""
    prompt = update.message.text
    user = update.effective_user
    logger.info(f"Получен промпт от {user.id}: {prompt}")

    image_id = context.user_data.get('current_image_id')
    file_path = context.user_data.get('current_file_path')

    if not image_id or not file_path:
        await update.message.reply_text("❌ Ошибка: изображение не найдено")
        return ConversationHandler.END

    status_msg = await update.message.reply_text(
        f"🎨 Обрабатываю...\n\nПромпт: {prompt}\n\n⏳ Обычно это занимает 10-30 секунд")

    try:
        async with aiosqlite.connect("bot_database.db") as db:
            await db.execute("UPDATE images SET prompt = ? WHERE id = ?", (prompt, image_id))
            await db.commit()

        result_url = await call_piapi_image_edit(file_path, prompt)

        if result_url:
            result_filename = f"user_{user.id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}_result.jpg"
            result_path = os.path.join(PROCESSED_DIR, result_filename)

            current_session = await get_session()
            async with current_session.get(result_url) as resp:
                if resp.status == 200:
                    with open(result_path, "wb") as f:
                        f.write(await resp.read())

                    # Удаляем сообщение о статусе
                    await status_msg.delete()

                    # Отправляем результат
                    with open(result_path, "rb") as photo:
                        await update.message.reply_photo(
                            photo=photo,
                            caption=f"✅ Готово!\n\n✨ Промпт: {prompt}"
                        )

                    # Сохраняем file_id
                    msg = await update.message.reply_photo(photo=open(result_path, "rb"))
                    processed_file_id = msg.photo[-1].file_id

                    await update_image_record(
                        image_id=image_id,
                        processed_file_id=processed_file_id,
                        processed_url=result_path,
                        status='completed'
                    )
                else:
                    raise Exception("Ошибка загрузки результата")
        else:
            raise Exception("Не получен URL результата")

    except Exception as e:
        error_msg = str(e)
        logger.error(f"Ошибка: {error_msg}")
        await status_msg.edit_text(f"❌ Ошибка: {error_msg}\n\nПопробуйте еще раз или отправьте другое фото.")
        await update_image_record(image_id=image_id, status='failed')

    context.user_data.clear()
    return ConversationHandler.END


async def call_piapi_image_edit(image_path: str, prompt: str) -> str:
    """Вызов PiAPI"""
    if not PIApi_API_KEY:
        raise Exception("PIApi_API_KEY не настроен")

    headers = {
        "x-api-key": PIApi_API_KEY,
        "Content-Type": "application/json"
    }

    with open(image_path, "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')

    payload = {
        "model": "black-forest-labs/FLUX.1-dev",
        "task_type": "image-to-image",
        "input": {
            "image": f"data:image/jpeg;base64,{image_base64}",
            "prompt": prompt,
            "num_inference_steps": 28,
            "guidance_scale": 7.5,
            "strength": 0.8
        }
    }

    current_session = await get_session()

    async with current_session.post(PIApi_BASE_URL, headers=headers, json=payload, timeout=30) as response:
        if response.status != 200:
            raise Exception(f"HTTP {response.status}")

        data = await response.json()
        if data.get('code') != 200:
            raise Exception(f"API Error: {data.get('message')}")

        task_id = data['data']['task_id']

    result_url = await poll_task_status(task_id, headers)
    return result_url


async def poll_task_status(task_id: str, headers: dict, max_attempts: int = 60) -> str:
    """Ожидание завершения"""
    get_url = f"https://api.piapi.ai/api/v1/task/{task_id}"
    current_session = await get_session()

    for attempt in range(max_attempts):
        try:
            async with current_session.get(get_url, headers=headers) as response:
                if response.status != 200:
                    await asyncio.sleep(2)
                    continue

                data = await response.json()
                status = data.get('data', {}).get('status')
                logger.info(f"Статус задачи {task_id}: {status} (попытка {attempt + 1})")

                if status == 'completed':
                    output = data.get('data', {}).get('output', {})
                    result_url = output.get('url') or output.get('image')
                    if result_url:
                        return result_url
                elif status == 'failed':
                    raise Exception("Задача не выполнена")

                await asyncio.sleep(2)
        except Exception as e:
            logger.error(f"Ошибка: {e}")
            await asyncio.sleep(2)

    raise Exception("Превышено время ожидания")


async def history_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """История обработок"""
    user = update.effective_user
    images = await get_user_images(user.id, limit=5)

    if not images:
        await update.message.reply_text(
            "📭 У вас пока нет обработанных изображений.\n\nОтправьте фото для начала работы!")
        return

    text = "🖼 **Ваши последние обработки:**\n\n"
    for i, img in enumerate(images, 1):
        status_emoji = "✅" if img['status'] == 'completed' else "❌" if img['status'] == 'failed' else "⏳"
        date = img['created_at'][:16].replace('T', ' ')
        prompt_preview = img['prompt'][:40] + "..." if img['prompt'] and len(img['prompt']) > 40 else img[
                                                                                                          'prompt'] or "без промпта"
        text += f"{status_emoji} **{i}.** {date}\n   📝 {prompt_preview}\n\n"

    text += "💡 Отправьте новое фото для обработки!"
    await update.message.reply_text(text, parse_mode='Markdown')


async def error_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Обработчик ошибок"""
    logger.error(f"Ошибка: {context.error}")
    if update and update.effective_message:
        await update.effective_message.reply_text("😵 Произошла ошибка. Попробуйте позже.")


async def main():
    """Запуск бота"""
    print("\n" + "=" * 50)
    print("🤖 ЗАПУСК TELEGRAM БОТА")
    print("=" * 50)

    if PROXY_URL:
        print(f"🔌 Прокси: {PROXY_URL.split('@')[-1] if '@' in PROXY_URL else PROXY_URL}")
    else:
        print("🔌 Прокси: Не используется")

    print("=" * 50)

    await init_db()

    # Создаем приложение
    request = get_telegram_request()
    application = Application.builder().token(TELEGRAM_BOT_TOKEN).request(request).build()

    # Регистрируем обработчики
    application.add_handler(CommandHandler("start", start))
    application.add_handler(CommandHandler("help", help_command))
    application.add_handler(CommandHandler("history", history_command))
    application.add_handler(CommandHandler("test", test_command))

    conv_handler = ConversationHandler(
        entry_points=[MessageHandler(filters.PHOTO, handle_photo)],
        states={AWAITING_PROMPT: [MessageHandler(filters.TEXT & ~filters.COMMAND, process_image_with_prompt)]},
        fallbacks=[CommandHandler("cancel", cancel)],
    )
    application.add_handler(conv_handler)
    application.add_error_handler(error_handler)

    print("\n✅ Бот запускается...")

    try:
        await application.initialize()
        await application.start()
        await application.updater.start_polling()

        print("\n" + "=" * 50)
        print("🤖 БОТ УСПЕШНО ЗАПУЩЕН!")
        print("=" * 50)
        print("\n📱 Отправь команду /test в Telegram для проверки\n")

        # Бесконечное ожидание
        await asyncio.Event().wait()

    except KeyboardInterrupt:
        print("\n🛑 Бот остановлен")
    except Exception as e:
        logger.error(f"Ошибка: {e}")
        print(f"\n❌ Ошибка: {e}")
    finally:
        await close_session()
        await application.updater.stop()
        await application.stop()


if __name__ == "__main__":
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

    asyncio.run(main())