memory.py

"""
الذاكرة الدائمة - SQLite
يحفظ:
- سياق كل مستخدم (آخر N رسالة)
- تفضيلات المستخدم (النموذج المفضل، اللغة)
- إحصائيات الاستخدام
"""
import asyncio
import json
import logging
import os
import time
from typing import List, Dict, Optional

import aiosqlite

from config import config

logger = logging.getLogger(__name__)


class Memory:
    """إدارة ذاكرة البوت الدائمة عبر SQLite"""

    def __init__(self, db_path: str):
        self.db_path = db_path
        # تأكد من وجود المجلد الأب
        os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True)
        self._lock = asyncio.Lock()

    async def init(self):
        """تهيئة قاعدة البيانات"""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("""
                CREATE TABLE IF NOT EXISTS messages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id INTEGER NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT NOT NULL,
                    model TEXT,
                    timestamp INTEGER NOT NULL
                )
            """)
            await db.execute("""
                CREATE INDEX IF NOT EXISTS idx_messages_user
                ON messages(user_id, timestamp DESC)
            """)
            await db.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    user_id INTEGER PRIMARY KEY,
                    username TEXT,
                    first_name TEXT,
                    language TEXT DEFAULT 'ar',
                    preferred_model TEXT,
                    is_owner INTEGER DEFAULT 0,
                    is_admin INTEGER DEFAULT 0,
                    first_seen INTEGER NOT NULL,
                    last_seen INTEGER NOT NULL,
                    message_count INTEGER DEFAULT 0
                )
            """)
            await db.execute("""
                CREATE TABLE IF NOT EXISTS github_cache (
                    user_id INTEGER NOT NULL,
                    repo_full_name TEXT NOT NULL,
                    data TEXT NOT NULL,
                    timestamp INTEGER NOT NULL,
                    PRIMARY KEY (user_id, repo_full_name)
                )
            """)
            await db.commit()
        logger.info(f"Database initialized at {self.db_path}")

    async def register_user(
        self,
        user_id: int,
        username: str = "",
        first_name: str = "",
    ) -> Dict:
        """تسجيل أو تحديث مستخدم"""
        now = int(time.time())
        is_owner = 1 if config.is_owner(user_id) else 0
        is_admin = 1 if config.is_admin(user_id) else 0

        async with self._lock, aiosqlite.connect(self.db_path) as db:
            # تحقق من وجود المستخدم
            cursor = await db.execute(
                "SELECT user_id, message_count FROM users WHERE user_id = ?",
                (user_id,),
            )
            row = await cursor.fetchone()

            if row is None:
                await db.execute(
                    """INSERT INTO users
                    (user_id, username, first_name, language, is_owner, is_admin,
                     first_seen, last_seen, message_count)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0)""",
                    (user_id, username, first_name, "ar", is_owner, is_admin, now, now),
                )
                await db.commit()
                return {
                    "user_id": user_id, "username": username,
                    "first_name": first_name, "is_new": True,
                    "message_count": 0,
                }
            else:
                await db.execute(
                    """UPDATE users SET
                       username = ?, first_name = ?, is_owner = ?, is_admin = ?,
                       last_seen = ?
                       WHERE user_id = ?""",
                    (username, first_name, is_owner, is_admin, now, user_id),
                )
                await db.commit()
                return {
                    "user_id": user_id, "username": username,
                    "first_name": first_name, "is_new": False,
                    "message_count": row[1],
                }

    async def increment_message_count(self, user_id: int):
        async with self._lock, aiosqlite.connect(self.db_path) as db:
            await db.execute(
                "UPDATE users SET message_count = message_count + 1 WHERE user_id = ?",
                (user_id,),
            )
            await db.commit()

    async def add_message(
        self,
        user_id: int,
        role: str,
        content: str,
        model: str = "",
    ):
        """إضافة رسالة للذاكرة"""
        # استخدام time.time() بدلاً من int(time.time()) لدقة أعلى
        now = time.time()
        async with self._lock, aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """INSERT INTO messages (user_id, role, content, model, timestamp)
                   VALUES (?, ?, ?, ?, ?)""",
                (user_id, role, content, model, now),
            )
            # حذف الرسائل القديمة الزائدة عن الحد
            await db.execute(
                """DELETE FROM messages WHERE user_id = ? AND id NOT IN (
                       SELECT id FROM messages WHERE user_id = ?
                       ORDER BY timestamp DESC, id DESC LIMIT ?
                   )""",
                (user_id, user_id, config.MAX_HISTORY_PER_USER),
            )
            await db.commit()

    async def get_history(self, user_id: int) -> List[Dict[str, str]]:
        """جلب سياق محادثة المستخدم (مرتب زمنياً، مع id كحاكم)"""
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute(
                """SELECT role, content FROM messages
                   WHERE user_id = ? AND role IN ('user', 'assistant')
                   ORDER BY timestamp ASC, id ASC""",
                (user_id,),
            )
            rows = await cursor.fetchall()
            return [{"role": r[0], "content": r[1]} for r in rows]

    async def clear_history(self, user_id: int) -> int:
        """حذف ذاكرة مستخدم - يعيد عدد الرسائل المحذوفة"""
        async with self._lock, aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute(
                "SELECT COUNT(*) FROM messages WHERE user_id = ?",
                (user_id,),
            )
            count = (await cursor.fetchone())[0]
            await db.execute("DELETE FROM messages WHERE user_id = ?", (user_id,))
            await db.commit()
            return count

    async def set_preferred_model(self, user_id: int, model: str):
        async with self._lock, aiosqlite.connect(self.db_path) as db:
            await db.execute(
                "UPDATE users SET preferred_model = ? WHERE user_id = ?",
                (model, user_id),
            )
            await db.commit()

    async def get_preferred_model(self, user_id: int) -> Optional[str]:
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute(
                "SELECT preferred_model FROM users WHERE user_id = ?",
                (user_id,),
            )
            row = await cursor.fetchone()
            return row[0] if row and row[0] else None

    async def get_user_language(self, user_id: int) -> str:
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute(
                "SELECT language FROM users WHERE user_id = ?",
                (user_id,),
            )
            row = await cursor.fetchone()
            return row[0] if row else "ar"

    async def set_user_language(self, user_id: int, lang: str):
        async with self._lock, aiosqlite.connect(self.db_path) as db:
            await db.execute(
                "UPDATE users SET language = ? WHERE user_id = ?",
                (lang, user_id),
            )
            await db.commit()

    async def get_stats(self) -> Dict:
        """إحصائيات عامة"""
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute("SELECT COUNT(*) FROM users")
            users_count = (await cursor.fetchone())[0]
            cursor = await db.execute("SELECT COUNT(*) FROM messages")
            messages_count = (await cursor.fetchone())[0]
            cursor = await db.execute(
                "SELECT COUNT(*) FROM users WHERE is_owner = 1"
            )
            owners = (await cursor.fetchone())[0]
            return {
                "users": users_count,
                "messages": messages_count,
                "owners_configured": owners,
            }


# Singleton
memory = Memory(config.DB_PATH)