import os
import sys
import gc
import json
import logging
import asyncio
import traceback
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional
from contextlib import asynccontextmanager
import tempfile

# تنظیم مسیرها - استفاده از مسیر temp برای حل مشکل دسترسی
BASE_DIR = Path(tempfile.gettempdir()) / "ai_assistant"
CACHE_DIR = BASE_DIR / "cache"
LOG_DIR = BASE_DIR / "logs"
STATIC_DIR = BASE_DIR / "static"
TEMPLATES_DIR = BASE_DIR / "templates"

# ایجاد مسیرها
for path in [BASE_DIR, CACHE_DIR, LOG_DIR, STATIC_DIR, TEMPLATES_DIR]:
    path.mkdir(exist_ok=True, parents=True)

# تنظیم متغیرهای محیطی
os.environ["TRANSFORMERS_CACHE"] = str(CACHE_DIR)
os.environ["HF_HOME"] = str(CACHE_DIR)
os.environ["TORCH_HOME"] = str(CACHE_DIR / "torch")

# کتابخانه‌های هوش مصنوعی
import torch
import transformers
from transformers import pipeline

# کتابخانه‌های وب و تلگرام
import uvicorn
from fastapi import FastAPI, WebSocket, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from telegram import Update, Bot, InlineKeyboardButton, InlineKeyboardMarkup, BotCommand
from telegram.ext import Application, CommandHandler, MessageHandler, CallbackQueryHandler, filters

# کتابخانه‌های دیگر
import hazm
import httpx

# تنظیم لاگینگ
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(LOG_DIR / "app.log", encoding='utf-8'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

class Config:
    """تنظیمات اصلی برنامه"""

    # تنظیمات ربات تلگرام
    TELEGRAM_TOKEN = "7437859619:AAGeGG3ZkLM0OVaw-Exx1uMRE55JtBCZZCY"
    CHAT_ID = "-1002228627548"

    # تنظیمات مسیرها
    BASE_DIR = BASE_DIR
    CACHE_DIR = CACHE_DIR
    LOG_DIR = LOG_DIR
    STATIC_DIR = STATIC_DIR
    TEMPLATES_DIR = TEMPLATES_DIR

    # تنظیمات مدل
    MODEL_NAME = "bigscience/bloom-560m"  # استفاده از مدل کوچکتر
    MAX_LENGTH = 200
    MIN_LENGTH = 20
    TEMPERATURE = 0.7
    TOP_K = 50
    TOP_P = 0.95
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

    # تنظیمات عمومی
    DEBUG = True
    HOST = "0.0.0.0"
    PORT = 8000
    MAX_RETRIES = 3
    TIMEOUT = 30
    CLEANUP_INTERVAL = 3600  # یک ساعت
    MAX_QUEUE_SIZE = 1000
    MAX_CACHE_SIZE = 1000
    MAX_MESSAGE_LENGTH = 500

class BloomAI:
    """کلاس مدیریت مدل هوش مصنوعی"""

    def __init__(self):
        self.pipeline = None
        self.device = Config.DEVICE

    async def initialize(self) -> bool:
        try:
            logger.info(f"Loading AI model: {Config.MODEL_NAME}")
            for attempt in range(Config.MAX_RETRIES):
                try:
                    self.pipeline = pipeline(
                        "text-generation",
                        model=Config.MODEL_NAME,
                        device=self.device,
                        trust_remote_code=True
                    )
                    logger.info("AI model loaded successfully")
                    return True
                except Exception as e:
                    if attempt < Config.MAX_RETRIES - 1:
                        logger.warning(f"Attempt {attempt + 1} failed, retrying in 5 seconds...")
                        await asyncio.sleep(5)
                    else:
                        raise e
        except Exception as e:
            logger.error(f"Failed to load model: {str(e)}")
            logger.error(traceback.format_exc())
            return False

    async def generate_response(self, text: str) -> str:
        """تولید پاسخ با استفاده از مدل"""
        try:
            if not self.pipeline:
                logger.error("Model not initialized")
                return "مدل در دسترس نیست."

            outputs = self.pipeline(
                text,
                max_length=Config.MAX_LENGTH,
                min_length=Config.MIN_LENGTH,
                do_sample=True,
                top_k=Config.TOP_K,
                top_p=Config.TOP_P,
                temperature=Config.TEMPERATURE,
                num_return_sequences=1,
                pad_token_id=self.pipeline.tokenizer.eos_token_id
            )

            generated_text = outputs[0]['generated_text']
            # حذف متن ورودی از پاسخ
            response = generated_text.replace(text, "").strip()

            if not response:
                return "متأسفم، نتوانستم پاسخ مناسبی تولید کنم."

            return response

        except Exception as e:
            logger.error(f"Error generating response: {e}")
            logger.error(traceback.format_exc())
            return "خطا در تولید پاسخ."

    async def cleanup(self):
        """پاکسازی حافظه و کش مدل"""
        try:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            gc.collect()
            logger.info("Model cleanup completed")
        except Exception as e:
            logger.error(f"Error in cleanup: {e}")

class TelegramUI:
    """کلاس مدیریت رابط کاربری تلگرام"""

    @staticmethod
    def get_main_menu() -> InlineKeyboardMarkup:
        """منوی اصلی با دکمه‌های تخت"""
        keyboard = [
            [
                InlineKeyboardButton("💬 چت جدید", callback_data="new_chat"),
                InlineKeyboardButton("🔍 جستجو", callback_data="search")
            ],
            [
                InlineKeyboardButton("📚 آموزش", callback_data="learn"),
                InlineKeyboardButton("🎯 تمرین", callback_data="practice")
            ],
            [
                InlineKeyboardButton("⚙️ تنظیمات", callback_data="settings"),
                InlineKeyboardButton("❓ راهنما", callback_data="help")
            ]
        ]
        return InlineKeyboardMarkup(keyboard)

    @staticmethod
    def get_chat_menu() -> InlineKeyboardMarkup:
        """منوی چت"""
        keyboard = [
            [
                InlineKeyboardButton("💭 چت عمومی", callback_data="general_chat"),
                InlineKeyboardButton("❓ سوال و جواب", callback_data="qa_chat")
            ],
            [
                InlineKeyboardButton("📝 ترجمه متن", callback_data="translate"),
                InlineKeyboardButton("🔄 تصحیح متن", callback_data="correct_text")
            ],
            [
                InlineKeyboardButton("🏠 بازگشت به منوی اصلی", callback_data="main_menu")
            ]
        ]
        return InlineKeyboardMarkup(keyboard)

class TelegramBot:
    """کلاس اصلی ربات تلگرام"""

    def __init__(self, ai_model: BloomAI):
        self.model = ai_model
        self.ui = TelegramUI()
        self.application = (
            Application.builder()
            .token(Config.TELEGRAM_TOKEN)
            .build()
        )
        self.normalizer = hazm.Normalizer()
        self.client = httpx.AsyncClient(timeout=Config.TIMEOUT)
        self._setup_handlers()
        self.message_queue = asyncio.Queue()
        self.response_cache = {}

    def _setup_handlers(self):
        """تنظیم هندلرهای ربات"""
        # دستورات اصلی
        self.application.add_handler(CommandHandler("start", self.start_command))
        self.application.add_handler(CommandHandler("help", self.help_command))
        self.application.add_handler(CommandHandler("menu", self.menu_command))
        self.application.add_handler(CommandHandler("cancel", self.cancel_command))

        # هندلر دکمه‌ها
        self.application.add_handler(CallbackQueryHandler(self.handle_callback))

        # هندلر پیام‌های متنی
        self.application.add_handler(
            MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message)
        )

    async def start_command(self, update: Update, context):
        """پردازش دستور /start"""
        welcome_text = """
🌟 به دستیار هوشمند خوش آمدید!

من می‌تونم در موارد زیر کمکتون کنم:
• پاسخ به سوالات شما
• ترجمه و تصحیح متن
• آموزش و تمرین

برای شروع، از منوی زیر گزینه مورد نظر خود را انتخاب کنید:
"""
        await update.message.reply_text(
            welcome_text,
            reply_markup=self.ui.get_main_menu()
        )

    async def help_command(self, update: Update, context):
        """پردازش دستور /help"""
        help_text = """
📚 راهنمای استفاده از ربات:

🔹 دستورات اصلی:
/start - شروع مجدد ربات
/menu - نمایش منوی اصلی
/help - نمایش این راهنما
/cancel - لغو عملیات فعلی

🔹 نحوه استفاده:
1. از منوی اصلی بخش مورد نظر را انتخاب کنید
2. در بخش چت، سوال یا متن خود را وارد کنید
3. منتظر پاسخ بمانید
"""
        await update.message.reply_text(
            help_text,
            reply_markup=self.ui.get_main_menu()
        )

    async def menu_command(self, update: Update, context):
        """نمایش منوی اصلی"""
        await update.message.reply_text(
            "🏠 منوی اصلی:",
            reply_markup=self.ui.get_main_menu()
        )

    async def cancel_command(self, update: Update, context):
        """لغو عملیات فعلی"""
        # پاک کردن صف پیام‌ها
        while not self.message_queue.empty():
            self.message_queue.get_nowait()

        await update.message.reply_text(
            "❌ عملیات فعلی لغو شد.",
            reply_markup=self.ui.get_main_menu()
        )

    async def handle_callback(self, update: Update, context):
        """پردازش دکمه‌های فشرده شده"""
        query = update.callback_query
        await query.answer()

        if query.data == "main_menu":
            await query.message.edit_text(
                "🏠 منوی اصلی:",
                reply_markup=self.ui.get_main_menu()
            )

        elif query.data == "new_chat":
            await query.message.edit_text(
                "💭 لطفاً نوع چت را انتخاب کنید:",
                reply_markup=self.ui.get_chat_menu()
            )

        elif query.data in ["general_chat", "qa_chat"]:
            chat_type = "عمومی" if query.data == "general_chat" else "سوال و جواب"
            await query.message.edit_text(
                f"💬 وارد بخش چت {chat_type} شدید.\nلطفاً پیام خود را بنویسید..."
            )

    async def handle_message(self, update: Update, context):
        """پردازش پیام‌های دریافتی"""
        try:
            # بررسی طول پیام
            if len(update.message.text) > Config.MAX_MESSAGE_LENGTH:
                await update.message.reply_text("⚠️ پیام شما بیش از حد طولانی است.")
                return

            # بررسی حجم صف پیام‌ها
            if self.message_queue.qsize() >= Config.MAX_QUEUE_SIZE:
                await update.message.reply_text("⚠️ سیستم در حال حاضر شلوغ است. لطفاً کمی صبر کنید.")
                return

            # نرمال‌سازی متن
            text = self.normalizer.normalize(update.message.text)

            # بررسی کش
            if text in self.response_cache:
                response = self.response_cache[text]
                logger.info("Using cached response")
            else:
                # ارسال پیام در حال پردازش
                processing_message = await update.message.reply_text(
                    "🤔 در حال پردازش پیام شما..."
                )

                # دریافت پاسخ از مدل
                response = await self.model.generate_response(text)

                # بررسی حجم کش
                if len(self.response_cache) >= Config.MAX_CACHE_SIZE:
                    # حذف قدیمی‌ترین پاسخ‌ها
                    oldest_keys = list(self.response_cache.keys())[:100]
                    for key in oldest_keys:
                        del self.response_cache[key]

                # ذخیره در کش
                self.response_cache[text] = response

                # حذف پیام در حال پردازش
                await processing_message.delete()

            # ارسال پاسخ با دکمه برگشت به منو
            keyboard = [[InlineKeyboardButton("🏠 بازگشت به منو", callback_data="main_menu")]]
            reply_markup = InlineKeyboardMarkup(keyboard)

            await update.message.reply_text(
                f"🤖 {response}",
                reply_markup=reply_markup
            )

        except Exception as e:
            logger.error(f"خطا در پردازش پیام: {e}")
            logger.error(traceback.format_exc())
            await update.message.reply_text(
                "❌ متأسفانه در پردازش پیام مشکلی پیش آمد. لطفاً دوباره تلاش کنید."
            )

    async def periodic_cleanup(self):
        """پاکسازی دوره‌ای حافظه و کش"""
        try:
            while True:
                await asyncio.sleep(Config.CLEANUP_INTERVAL)
                logger.info("Starting periodic cleanup...")

                # پاکسازی کش پاسخ‌ها
                self.response_cache.clear()

                # پاکسازی حافظه
                await self.model.cleanup()

                logger.info("Periodic cleanup completed")
        except Exception as e:
            logger.error(f"Error in periodic cleanup: {e}")
            logger.error(traceback.format_exc())

    async def run(self):
        """اجرای ربات"""
        try:
            # شروع پاکسازی دوره‌ای
            asyncio.create_task(self.periodic_cleanup())

            # راه‌اندازی ربات
            await self.application.initialize()
            await self.application.start()
            await self.application.run_polling(
                allowed_updates=Update.ALL_TYPES,
                drop_pending_updates=True
            )

        except Exception as e:
            logger.error(f"خطا در اجرای ربات: {e}")
            logger.error(traceback.format_exc())
            raise

@asynccontextmanager
async def lifespan(app: FastAPI):
    """مدیریت چرخه حیات برنامه"""
    logger.info("Starting application...")

    # مقداردهی متغیرهای حالت برنامه
    app.state.is_ready = asyncio.Event()

    # ایجاد و مقداردهی مدل
    app.state.model = BloomAI()
    if not await app.state.model.initialize():
        logger.error("Failed to initialize AI model")
        sys.exit(1)

    # راه‌اندازی ربات تلگرام
    app.state.bot = TelegramBot(app.state.model)
    asyncio.create_task(app.state.bot.run())

    # تنظیم وضعیت آماده
    app.state.is_ready.set()

    yield

    # خاموش کردن تسک‌ها و ربات
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
    for task in tasks:
        task.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)
    await app.state.bot.application.shutdown()
    logger.info("Application shutdown completed")

# برنامه FastAPI
app = FastAPI(title="AI Assistant", lifespan=lifespan)

# پیوست کردن فایل‌های استاتیک
app.mount("/static", StaticFiles(directory=Config.STATIC_DIR), name="static")
templates = Jinja2Templates(directory=Config.TEMPLATES_DIR)

# HTML پیش‌فرض
DEFAULT_HTML = """
<!DOCTYPE html>
<html dir="rtl" lang="fa">
<head>
    <meta charset="UTF-8">
    <title>دستیار هوشمند</title>
    <style>
        body {
            font-family: Tahoma, Arial;
            background-color: #f5f5f5;
            margin: 0;
            padding: 20px;
        }
        .container {
            max-width: 800px;
            margin: 0 auto;
            background: white;
            padding: 20px;
            border-radius: 10px;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
        h1 { color: #1a73e8; }
        .status { margin: 20px 0; }
        .chat-container {
            margin-top: 20px;
            border: 1px solid #e0e0e0;
            border-radius: 8px;
            padding: 10px;
            min-height: 300px;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>دستیار هوشمند</h1>
        <div class="status">
            <p>برای استفاده از دستیار هوشمند، به ربات تلگرام مراجعه کنید یا از رابط وب استفاده نمایید.</p>
        </div>
        <div class="chat-container" id="chat-container">
        </div>
        <input type="text" id="message-input" placeholder="پیام خود را بنویسید..." style="width: 100%; padding: 10px; margin-top: 10px;">
        <button onclick="sendMessage()" style="margin-top: 10px; padding: 10px 20px;">ارسال</button>
    </div>
    <script>
        const ws = new WebSocket(`ws://${window.location.host}/ws`);
        const chatContainer = document.getElementById('chat-container');
        const messageInput = document.getElementById('message-input');

        ws.onmessage = function(event) {
            const response = JSON.parse(event.data);
            appendMessage('bot', response.response);
        };

        function appendMessage(type, text) {
            const messageDiv = document.createElement('div');
            messageDiv.style.padding = '10px';
            messageDiv.style.margin = '5px';
            messageDiv.style.borderRadius = '5px';

            if (type === 'user') {
                messageDiv.style.backgroundColor = '#e3f2fd';
                messageDiv.style.marginLeft = '20%';
            } else {
                messageDiv.style.backgroundColor = '#f5f5f5';
                messageDiv.style.marginRight = '20%';
            }

            messageDiv.textContent = text;
            chatContainer.appendChild(messageDiv);
            chatContainer.scrollTop = chatContainer.scrollHeight;
        }

        function sendMessage() {
            const message = messageInput.value.trim();
            if (message) {
                appendMessage('user', message);
                ws.send(message);
                messageInput.value = '';
            }
        }

        messageInput.addEventListener('keypress', function(e) {
            if (e.key === 'Enter') {
                sendMessage();
            }
        });
    </script>
</body>
</html>
"""

@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    """صفحه اصلی"""
    try:
        # می‌توانید اطلاعات بیشتری را به تمپلیت ارسال کنید
        return templates.TemplateResponse(
            "index.html",
            {
                "request": request,
                "bot_status": "فعال" if request.app.state.is_ready.is_set() else "در حال آماده‌سازی"
            }
        )
    except Exception as e:
        logger.error(f"خطا در نمایش صفحه اصلی: {e}")
        return HTMLResponse(content=DEFAULT_HTML)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """اتصال WebSocket"""
    await websocket.accept()

    try:
        while True:
            # دریافت پیام
            text = await websocket.receive_text()

            # تولید پاسخ
            response = await app.state.model.generate_response(text)

            # ارسال پاسخ
            await websocket.send_json({
                "response": response
            })
    except WebSocketDisconnect:
        logger.info("اتصال WebSocket قطع شد")
    except Exception as e:
        logger.error(f"خطا در WebSocket: {e}")
        await websocket.close()

@app.exception_handler(404)
async def not_found_handler(request: Request, exc: Exception):
    """پردازش خطای 404"""
    return HTMLResponse(content=DEFAULT_HTML, status_code=404)

@app.exception_handler(500)
async def server_error_handler(request: Request, exc: Exception):
    """پردازش خطای 500"""
    return HTMLResponse(content=DEFAULT_HTML, status_code=500)

if __name__ == "__main__":
    # اجرای برنامه
    uvicorn.run(
        app,
        host=Config.HOST,
        port=Config.PORT,
        log_level="info"
    )