#!/usr/bin/env python3
import os
import re
import json
import asyncio
import logging
import tempfile
from datetime import datetime
from typing import Dict, List, Optional, Any

import aiohttp
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import (
    Application,
    CommandHandler,
    MessageHandler,
    ContextTypes,
    filters,
    CallbackQueryHandler
)
from telegram.request import HTTPXRequest
from google import genai
import firebase_admin
from firebase_admin import credentials, firestore

# ============= КОНФИГУРАЦИЯ =============
TG_BOT_TOKEN = "8531345451:AAE8qJlFEwIcuQdUiIFk2viZS9xLLMk2c_o"
ADMIN_ID = 8746779101
BOOST_LINK = "https://t.me/boost/chat_crx"

GEMINI_API_KEY = "AIzaSyC59FksQAoeVeBKzVpxpL8zByEg3DAkKgE"
GEMINI_MODEL_NAME = "gemini-2.5-flash"

RAPIDAPI_KEY = "5626a251afmsh7fc1baa1c873297p11bf87jsn3267c0694304"
RAPIDAPI_HOST = "veo-3-video.p.rapidapi.com"

HF_MUSICGEN_URL = "https://api-inference.huggingface.co/models/facebook/musicgen-medium"
HF_TOKEN = "hf_public"

# Cloudflare Worker прокси
CF_WORKER_URL = "https://tg-proxy.artyomanisimov37.workers.dev"

# ============= FIRESTORE =============
db = None
try:
    if not firebase_admin._apps:
        cred = credentials.Certificate("firebase_admin.json")
        firebase_admin.initialize_app(cred)
    db = firestore.client()
    logging.info("✅ Firebase подключён")
except Exception as e:
    logging.error(f"❌ Firebase ошибка: {e}")

# ============= GEMINI =============
genai_client = genai.Client(api_key=GEMINI_API_KEY)

# ============= CONTENT AGGREGATOR =============
class ContentAggregator:
    @staticmethod
    async def search_web(query: str) -> List[Dict[str, str]]:
        results = []
        try:
            url = f"https://html.duckduckgo.com/html/?q={query.replace(' ', '+')}"
            async with aiohttp.ClientSession() as session:
                async with session.get(url, headers={"User-Agent": "Mozilla/5.0"}) as resp:
                    html = await resp.text()
                    links = re.findall(r'<a[^>]*href="//([^"]+)"[^>]*>([^<]+)</a>', html)
                    for link, title in links[:5]:
                        if not any(x in link for x in ["youtube.com", "facebook.com", "twitter.com"]):
                            results.append({"title": title.strip(), "link": f"https://{link}"})
        except Exception as e:
            logging.error(f"Ошибка поиска: {e}")
        return results

    @staticmethod
    async def fetch_page_content(url: str) -> str:
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(url, timeout=10) as resp:
                    html = await resp.text()
                    text = re.sub(r'<[^>]+>', ' ', html)
                    text = re.sub(r'\s+', ' ', text)
                    return text[:4000]
        except:
            return ""

    @staticmethod
    async def generate_unique_content(topic: str) -> Dict[str, Any]:
        logging.info(f"Агент собирает контент: {topic}")
        search_results = await ContentAggregator.search_web(topic)
        contexts = []
        for res in search_results[:3]:
            content = await ContentAggregator.fetch_page_content(res["link"])
            if content:
                contexts.append(content)

        aggregated = "\n\n---\n\n".join(contexts)
        prompt = f"""
        Ты – контент-мейкер. На основе данных по теме "{topic}":
        {aggregated[:6000]}

        Создай JSON с полями:
        - "article": уникальная статья 800-1500 слов
        - "video_ideas": массив из 10 идей для видео
        - "insights": массив из 5 инсайтов, которых нет в исходных данных
        """
        try:
            # ИСПОЛЬЗУЕМ .aio ДЛЯ АСИНХРОННЫХ ВЫЗОВОВ
            response = await genai_client.aio.models.generate_content(
                model=GEMINI_MODEL_NAME,
                contents=prompt
            )
            match = re.search(r'\{.*\}', response.text, re.DOTALL)
            if match:
                return json.loads(match.group())
            return {"error": "Не удалось распарсить JSON", "raw": response.text}
        except Exception as e:
            return {"error": str(e)}

# ============= MUSIC GENERATOR =============
class MusicGenerator:
    @staticmethod
    async def generate_music(prompt: str) -> Optional[str]:
        headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN and HF_TOKEN != "hf_public" else {}
        payload = {"inputs": prompt, "parameters": {"max_new_tokens": 256}}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(HF_MUSICGEN_URL, headers=headers, json=payload, timeout=60) as resp:
                    if resp.status == 200:
                        audio_data = await resp.read()
                        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                            tmp.write(audio_data)
                            return tmp.name
                    else:
                        logging.warning(f"MusicGen вернул {resp.status}, используем демо")
                        return None
        except Exception as e:
            logging.error(f"MusicGen ошибка: {e}")
            return None

# ============= VIDEO GENERATOR =============
class VideoGenerator:
    @staticmethod
    async def generate_video(prompt: str) -> Optional[str]:
        url = "https://veo-3-video.p.rapidapi.com/v1/jobs"
        payload = {
            "model": "veo-3",
            "prompt": prompt,
            "duration": 8,
            "aspect_ratio": "16:9"
        }
        headers = {
            "x-rapidapi-key": RAPIDAPI_KEY,
            "x-rapidapi-host": RAPIDAPI_HOST,
            "Content-Type": "application/json"
        }
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(url, json=payload, headers=headers) as resp:
                    data = await resp.json()
                    job_id = data.get("job_id")
                    if not job_id:
                        return None
                await asyncio.sleep(15)
                status_url = f"https://veo-3-video.p.rapidapi.com/v1/jobs/{job_id}"
                async with session.get(status_url, headers=headers) as status_resp:
                    status_data = await status_resp.json()
                    return status_data.get("video_url")
        except Exception as e:
            logging.error(f"Видео ошибка: {e}")
            return "https://www.w3schools.com/html/mov_bbb.mp4"

# ============= BOT STATE =============
class BotState:
    def __init__(self):
        self.whitelist = set()
        self.premium_users = set()
        self.load_data()

    def load_data(self):
        if db:
            try:
                doc = db.collection("bot_config").document("whitelist").get()
                if doc.exists:
                    self.whitelist = set(doc.to_dict().get("users", []))
                prem = db.collection("bot_config").document("premium").get()
                if prem.exists:
                    self.premium_users = set(prem.to_dict().get("users", []))
                self.whitelist.add(ADMIN_ID)
            except Exception as e:
                logging.error(f"Загрузка БД: {e}")

    def is_allowed(self, user_id: int) -> bool:
        return user_id in self.whitelist or user_id == ADMIN_ID or user_id in self.premium_users

bot_state = BotState()

# ============= HANDLERS =============
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user = update.effective_user
    if not bot_state.is_allowed(user.id):
        await update.message.reply_text("⛔ Доступ ограничен. Обратитесь к админу.")
        return
    keyboard = [
        [InlineKeyboardButton("💎 Premium", url=BOOST_LINK)],
        [InlineKeyboardButton("🤖 Уникальный контент", callback_data="gen_content")],
        [InlineKeyboardButton("🎵 Музыка", callback_data="gen_music")],
        [InlineKeyboardButton("🎬 Видео", callback_data="gen_video")],
        [InlineKeyboardButton("📚 Помощь", callback_data="help")]
    ]
    await update.message.reply_text(
        f"👋 Привет, {user.first_name}!\nЯ ИИ бот на **Gemini 2.5 Flash**\n\n"
        "/ask [вопрос] – задать вопрос\n"
        "/content [тема] – создать статью + идеи\n"
        "/music [описание] – сгенерировать трек (Premium)\n"
        "/video [описание] – сгенерировать видео (Premium)",
        reply_markup=InlineKeyboardMarkup(keyboard)
    )

async def ask_ai(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    if not bot_state.is_allowed(user_id):
        return
    query = " ".join(context.args)
    if not query:
        await update.message.reply_text("Пример: /ask Как работает ИИ?")
        return
    msg = await update.message.reply_text("🤔 Думаю...")
    try:
        # ИСПОЛЬЗУЕМ .aio
        resp = await genai_client.aio.models.generate_content(model=GEMINI_MODEL_NAME, contents=query)
        await msg.edit_text(resp.text[:4000], parse_mode='Markdown')
    except Exception as e:
        await msg.edit_text("❌ Извините, сервера Google сейчас перегружены или произошла ошибка. Попробуйте позже.")
        logging.error(f"Ошибка Gemini: {e}")

async def create_content(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    if not bot_state.is_allowed(user_id):
        return
    topic = " ".join(context.args)
    if not topic:
        await update.message.reply_text("Укажите тему: /content космос")
        return
    msg = await update.message.reply_text("🔍 Агент ищет информацию... (30-60 сек)")
    data = await ContentAggregator.generate_unique_content(topic)
    if "error" in data:
        await msg.edit_text(f"Ошибка: {data['error']}")
        return
    answer = f"📄 *{topic}*\n\n"
    answer += f"**Статья:**\n{data.get('article', '')[:1500]}...\n\n"
    answer += "**💡 Идеи для видео:**\n" + "\n".join(f"{i+1}. {idea}" for i, idea in enumerate(data.get('video_ideas', [])[:5]))
    answer += "\n\n**🔍 Инсайты:**\n" + "\n".join(f"• {insight}" for insight in data.get('insights', [])[:3])
    await msg.edit_text(answer[:4000], parse_mode='Markdown')

async def generate_music(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    if user_id not in bot_state.premium_users and user_id != ADMIN_ID:
        await update.message.reply_text(f"🌟 Premium функция. Бустите канал: {BOOST_LINK}")
        return
    prompt = " ".join(context.args)
    if not prompt:
        await update.message.reply_text("Опишите музыку: /music грустная фортепианная мелодия")
        return
    msg = await update.message.reply_text("🎵 Сочиняю... до 30 сек")
    audio_path = await MusicGenerator.generate_music(prompt)
    if audio_path and os.path.exists(audio_path):
        await update.message.reply_audio(audio=open(audio_path, 'rb'), caption=f"🎵 {prompt[:100]}")
        os.unlink(audio_path)
    else:
        await update.message.reply_audio(audio="https://www.soundhelix.com/examples/mp3/SoundHelix-Song-1.mp3",
                                         caption="Демо-трек (API занят)")
    await msg.delete()

async def generate_video(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    if user_id not in bot_state.premium_users and user_id != ADMIN_ID:
        await update.message.reply_text(f"🎬 Premium функция. Бустите: {BOOST_LINK}")
        return
    prompt = " ".join(context.args)
    if not prompt:
        await update.message.reply_text("Опишите видео: /video кот в космосе")
        return
    msg = await update.message.reply_text("🎬 Генерирую видео... 1-2 мин")
    video_url = await VideoGenerator.generate_video(prompt)
    if video_url:
        await update.message.reply_video(video=video_url, caption=f"🎬 {prompt[:100]}")
    else:
        await update.message.reply_text("Не удалось создать видео, попробуйте позже")
    await msg.delete()

async def handle_text(update: Update, context: ContextTypes.DEFAULT_TYPE):
    text = update.message.text
    user_id = update.effective_user.id
    if not bot_state.is_allowed(user_id):
        return
        
    # Вспомогательная функция с обработкой ошибок API
    async def process_prompt(prompt_text):
        try:
            # ИСПОЛЬЗУЕМ .aio
            resp = await genai_client.aio.models.generate_content(model=GEMINI_MODEL_NAME, contents=prompt_text)
            await update.message.reply_text(resp.text[:4000], parse_mode='Markdown')
        except Exception as e:
            logging.error(f"Gemini API Error: {e}")
            await update.message.reply_text("⚠️ Ошибка: Сервера Google сейчас перегружены. Пожалуйста, попробуйте чуть позже.")

    if text.lower().startswith(("ии ", "ai ")):
        prompt = text[3:].strip()
        if prompt:
            await process_prompt(prompt)
    elif len(text) > 20 and not text.startswith("/"):
        await process_prompt(text)

async def premium_info(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(f"💎 Premium даёт видео, музыку, приоритет.\nПолучите: {BOOST_LINK}")

async def admin_whitelist(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if update.effective_user.id != ADMIN_ID:
        return
    if len(context.args) < 2:
        await update.message.reply_text("/w add|rem|prem [user_id]")
        return
    cmd, uid = context.args[0], int(context.args[1])
    if cmd == "add":
        bot_state.whitelist.add(uid)
        if db:
            db.collection("bot_config").document("whitelist").set({"users": list(bot_state.whitelist)})
        await update.message.reply_text(f"✅ {uid} в белом списке")
    elif cmd == "prem":
        bot_state.premium_users.add(uid)
        if db:
            db.collection("bot_config").document("premium").set({"users": list(bot_state.premium_users)})
        await update.message.reply_text(f"🌟 {uid} теперь Premium")

async def callback_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
    query = update.callback_query
    await query.answer()
    data = query.data
    if data == "help":
        await query.edit_message_text("/ask – вопрос\n/content – статья+идеи\n/music – трек (Premium)\n/video – клип (Premium)")
    elif data == "gen_content":
        await query.edit_message_text("Отправьте: /content [тема]")
    elif data == "gen_music":
        await query.edit_message_text("Premium команда: /music [описание]")
    elif data == "gen_video":
        await query.edit_message_text("Premium команда: /video [описание]")

# ============= MAIN =============
def main():
    logging.basicConfig(level=logging.INFO)
    
    # Оставляем только настройки таймаутов, proxy_url убран
    request = HTTPXRequest(
        connect_timeout=30.0,
        read_timeout=40.0,
        write_timeout=40.0,
        pool_timeout=15.0,
        connection_pool_size=20
    )
    
    # CF_WORKER_URL передан в base_url
    application = (
        Application.builder()
        .token(TG_BOT_TOKEN)
        .base_url(f"{CF_WORKER_URL}/bot")
        .request(request)
        .build()
    )
    
    application.add_handler(CommandHandler("start", start))
    application.add_handler(CommandHandler("ask", ask_ai))
    application.add_handler(CommandHandler("content", create_content))
    application.add_handler(CommandHandler("music", generate_music))
    application.add_handler(CommandHandler("video", generate_video))
    application.add_handler(CommandHandler("premium", premium_info))
    application.add_handler(CommandHandler("w", admin_whitelist))
    application.add_handler(CallbackQueryHandler(callback_handler))
    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_text))
    
    application.run_polling()

if __name__ == "__main__":
    main()