import os
import sys
import time
import sqlite3
import logging
from pathlib import Path
from typing import Optional, Dict, Any, List, Tuple
from dataclasses import dataclass, field, asdict

from config import settings

logger = logging.getLogger("AGYMemoryManager")

DATA_DIR = Path("/root/telegram-agy-bot/data")
DB_PATH = DATA_DIR / "memory.db"

@dataclass
class MemoryItem:
    id: int
    type: str  # 'global' or 'private'
    user_id: Optional[int]  # None for global, chat_id for private
    key: str
    category: str  # 'rule', 'preference', 'tech_stack', 'fact', 'workflow', 'general'
    content: str
    importance: int = 1
    created_by: Optional[int] = None
    created_at: float = field(default_factory=time.time)
    updated_at: float = field(default_factory=time.time)

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    @property
    def is_global(self) -> bool:
        return self.type == "global"

    @property
    def is_private(self) -> bool:
        return self.type == "private"


class MemoryManager:
    def __init__(self, db_path: Path = DB_PATH):
        self.db_path = db_path
        self._init_db()

    def _get_connection(self) -> sqlite3.Connection:
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        conn = sqlite3.connect(str(self.db_path), timeout=15.0)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL;")
        conn.execute("PRAGMA synchronous=NORMAL;")
        return conn

    def _init_db(self):
        """Initializes tables and indexes if they do not exist."""
        try:
            with self._get_connection() as conn:
                conn.execute("""
                    CREATE TABLE IF NOT EXISTS memories (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        type TEXT NOT NULL CHECK(type IN ('global', 'private')),
                        user_id INTEGER,
                        key TEXT NOT NULL,
                        category TEXT NOT NULL DEFAULT 'general',
                        content TEXT NOT NULL,
                        importance INTEGER NOT NULL DEFAULT 1,
                        created_by INTEGER,
                        created_at REAL NOT NULL,
                        updated_at REAL NOT NULL
                    );
                """)
                # Partial unique indexes to properly handle NULL user_id for global memories
                conn.execute("""
                    CREATE UNIQUE INDEX IF NOT EXISTS idx_memories_global_key 
                    ON memories(type, key) WHERE type = 'global';
                """)
                conn.execute("""
                    CREATE UNIQUE INDEX IF NOT EXISTS idx_memories_private_key 
                    ON memories(type, user_id, key) WHERE type = 'private';
                """)
                conn.execute("""
                    CREATE INDEX IF NOT EXISTS idx_memories_user_type 
                    ON memories(type, user_id);
                """)
                conn.execute("""
                    CREATE INDEX IF NOT EXISTS idx_memories_category 
                    ON memories(category);
                """)
            logger.info("Memory database initialized successfully at %s", self.db_path)
        except Exception as e:
            logger.error("Failed to initialize memory database: %s", e)
            raise

    def save_or_update(
        self,
        type_: str,
        key: str,
        content: str,
        user_id: Optional[int] = None,
        category: str = "general",
        importance: int = 1,
        created_by: Optional[int] = None,
    ) -> Tuple[MemoryItem, bool]:
        """
        Saves or updates a memory.
        Returns: (MemoryItem, is_created: bool)
        """
        type_ = type_.lower().strip()
        if type_ not in ("global", "private"):
            raise ValueError("Memory type must be 'global' or 'private'")

        if type_ == "global":
            user_id = None
        elif user_id is None:
            raise ValueError("Private memory requires a valid user_id")

        key = key.strip().lower()
        if not key:
            raise ValueError("Memory key cannot be empty")

        content = content.strip()
        if not content:
            raise ValueError("Memory content cannot be empty")

        category = category.strip().lower() or "general"
        importance = max(1, min(5, int(importance)))
        now = time.time()

        with self._get_connection() as conn:
            # Check if exists
            if type_ == "global":
                cur = conn.execute(
                    "SELECT id, created_at, created_by FROM memories WHERE type = 'global' AND key = ?",
                    (key,),
                )
            else:
                cur = conn.execute(
                    "SELECT id, created_at, created_by FROM memories WHERE type = 'private' AND user_id = ? AND key = ?",
                    (user_id, key),
                )
            row = cur.fetchone()

            if row:
                mem_id = row["id"]
                orig_created_at = row["created_at"]
                orig_created_by = row["created_by"] or created_by
                conn.execute(
                    """
                    UPDATE memories 
                    SET category = ?, content = ?, importance = ?, updated_at = ?
                    WHERE id = ?
                    """,
                    (category, content, importance, now, mem_id),
                )
                is_created = False
                created_at = orig_created_at
                creator = orig_created_by
            else:
                cur = conn.execute(
                    """
                    INSERT INTO memories (type, user_id, key, category, content, importance, created_by, created_at, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (type_, user_id, key, category, content, importance, created_by, now, now),
                )
                mem_id = cur.lastrowid
                is_created = True
                created_at = now
                creator = created_by

        item = MemoryItem(
            id=mem_id,
            type=type_,
            user_id=user_id,
            key=key,
            category=category,
            content=content,
            importance=importance,
            created_by=creator,
            created_at=created_at,
            updated_at=now,
        )
        logger.info(
            "%s memory [id=%d, key=%s, type=%s, user=%s]",
            "Created" if is_created else "Updated",
            mem_id,
            key,
            type_,
            user_id,
        )
        return item, is_created

    def get_by_id(self, memory_id: int) -> Optional[MemoryItem]:
        with self._get_connection() as conn:
            cur = conn.execute("SELECT * FROM memories WHERE id = ?", (memory_id,))
            row = cur.fetchone()
            if row:
                return MemoryItem(**dict(row))
        return None

    def get_by_key(self, type_: str, key: str, user_id: Optional[int] = None) -> Optional[MemoryItem]:
        type_ = type_.lower().strip()
        key = key.lower().strip()
        with self._get_connection() as conn:
            if type_ == "global":
                cur = conn.execute("SELECT * FROM memories WHERE type = 'global' AND key = ?", (key,))
            else:
                cur = conn.execute(
                    "SELECT * FROM memories WHERE type = 'private' AND user_id = ? AND key = ?",
                    (user_id, key),
                )
            row = cur.fetchone()
            if row:
                return MemoryItem(**dict(row))
        return None

    def get_global_memories(self, category: Optional[str] = None, limit: int = 100) -> List[MemoryItem]:
        with self._get_connection() as conn:
            if category:
                cur = conn.execute(
                    "SELECT * FROM memories WHERE type = 'global' AND category = ? ORDER BY importance DESC, updated_at DESC LIMIT ?",
                    (category.lower().strip(), limit),
                )
            else:
                cur = conn.execute(
                    "SELECT * FROM memories WHERE type = 'global' ORDER BY importance DESC, updated_at DESC LIMIT ?",
                    (limit,),
                )
            return [MemoryItem(**dict(r)) for r in cur.fetchall()]

    def get_user_memories(
        self, user_id: int, category: Optional[str] = None, limit: int = 100
    ) -> List[MemoryItem]:
        with self._get_connection() as conn:
            if category:
                cur = conn.execute(
                    "SELECT * FROM memories WHERE type = 'private' AND user_id = ? AND category = ? ORDER BY importance DESC, updated_at DESC LIMIT ?",
                    (user_id, category.lower().strip(), limit),
                )
            else:
                cur = conn.execute(
                    "SELECT * FROM memories WHERE type = 'private' AND user_id = ? ORDER BY importance DESC, updated_at DESC LIMIT ?",
                    (user_id, limit),
                )
            return [MemoryItem(**dict(r)) for r in cur.fetchall()]

    def delete_by_id(self, memory_id: int, user_id: Optional[int] = None, is_admin: bool = False) -> bool:
        with self._get_connection() as conn:
            cur = conn.execute("SELECT * FROM memories WHERE id = ?", (memory_id,))
            row = cur.fetchone()
            if not row:
                return False
            mem = MemoryItem(**dict(row))
            if mem.type == "global" and not is_admin:
                raise PermissionError("Only administrators can delete global memories")
            if mem.type == "private" and mem.user_id != user_id and not is_admin:
                raise PermissionError("Cannot delete another user's private memory")

            conn.execute("DELETE FROM memories WHERE id = ?", (memory_id,))
            logger.info("Deleted memory [id=%d, key=%s, type=%s]", memory_id, mem.key, mem.type)
            return True

    def delete_by_key(
        self, type_: str, key: str, user_id: Optional[int] = None, is_admin: bool = False
    ) -> bool:
        type_ = type_.lower().strip()
        key = key.lower().strip()
        with self._get_connection() as conn:
            if type_ == "global":
                if not is_admin:
                    raise PermissionError("Only administrators can delete global memories")
                cur = conn.execute("DELETE FROM memories WHERE type = 'global' AND key = ?", (key,))
            else:
                if user_id is None:
                    return False
                cur = conn.execute(
                    "DELETE FROM memories WHERE type = 'private' AND user_id = ? AND key = ?",
                    (user_id, key),
                )
            deleted = cur.rowcount > 0
            if deleted:
                logger.info("Deleted memory by key [key=%s, type=%s, user=%s]", key, type_, user_id)
            return deleted

    def clear_memories(self, type_: str, user_id: Optional[int] = None, is_admin: bool = False) -> int:
        type_ = type_.lower().strip()
        with self._get_connection() as conn:
            if type_ == "global":
                if not is_admin:
                    raise PermissionError("Only administrators can clear global memories")
                cur = conn.execute("DELETE FROM memories WHERE type = 'global'")
            else:
                if user_id is None:
                    return 0
                cur = conn.execute("DELETE FROM memories WHERE type = 'private' AND user_id = ?", (user_id,))
            count = cur.rowcount
            logger.info("Cleared %d %s memories for user %s", count, type_, user_id)
            return count

    def search_memories(
        self, query: str, user_id: Optional[int] = None, type_: Optional[str] = None, limit: int = 50
    ) -> List[MemoryItem]:
        q = f"%{query.strip().lower()}%"
        with self._get_connection() as conn:
            if type_ == "global":
                cur = conn.execute(
                    """
                    SELECT * FROM memories 
                    WHERE type = 'global' AND (key LIKE ? OR content LIKE ? OR category LIKE ?)
                    ORDER BY importance DESC, updated_at DESC LIMIT ?
                    """,
                    (q, q, q, limit),
                )
            elif type_ == "private" and user_id is not None:
                cur = conn.execute(
                    """
                    SELECT * FROM memories 
                    WHERE type = 'private' AND user_id = ? AND (key LIKE ? OR content LIKE ? OR category LIKE ?)
                    ORDER BY importance DESC, updated_at DESC LIMIT ?
                    """,
                    (user_id, q, q, q, limit),
                )
            elif user_id is not None:
                cur = conn.execute(
                    """
                    SELECT * FROM memories 
                    WHERE (type = 'global' OR (type = 'private' AND user_id = ?)) 
                      AND (key LIKE ? OR content LIKE ? OR category LIKE ?)
                    ORDER BY type DESC, importance DESC, updated_at DESC LIMIT ?
                    """,
                    (user_id, q, q, q, limit),
                )
            else:
                cur = conn.execute(
                    """
                    SELECT * FROM memories 
                    WHERE type = 'global' AND (key LIKE ? OR content LIKE ? OR category LIKE ?)
                    ORDER BY importance DESC, updated_at DESC LIMIT ?
                    """,
                    (q, q, q, limit),
                )
            return [MemoryItem(**dict(r)) for r in cur.fetchall()]

    def format_memories_for_prompt(self, user_id: Optional[int] = None) -> Optional[str]:
        """
        Formats all applicable memories (Global + User's Private) into a concise, token-efficient
        system instruction block for the AI model.
        Returns None if no memories are found.
        """
        globals_list = self.get_global_memories(limit=50)
        privates_list = self.get_user_memories(user_id=user_id, limit=50) if user_id else []

        if not globals_list and not privates_list:
            return None

        lines = [
            "[SYSTEM INSTRUCTION: AI PERSISTENT MEMORY & CONTINUOUS LEARNING]",
            "You have access to persistent long-term memory containing verified system rules and user context.",
            "Use this knowledge to maintain continuity and customize all your responses accordingly.\n"
        ]

        if globals_list:
            lines.append("🌐 GLOBAL SYSTEM MEMORY (قوانین و اطلاعات عمومی سیستم):")
            for m in globals_list:
                cat = f"[{m.category}]" if m.category and m.category != "general" else ""
                lines.append(f"• [key: {m.key}] {cat} {m.content}")
            lines.append("")

        if privates_list:
            lines.append(f"👤 USER PRIVATE MEMORY (خاطرات و ترجیحات اختصاصی این کاربر):")
            for m in privates_list:
                cat = f"[{m.category}]" if m.category and m.category != "general" else ""
                lines.append(f"• [key: {m.key}] {cat} {m.content}")
            lines.append("")

        lines.extend([
            "🧠 AUTONOMOUS MEMORY MANAGEMENT RULES:",
            "1. When the user shares important facts, personal preferences, coding conventions, or rules that should persist, output the action tag:",
            '   [[SAVE_MEMORY: type="private|global", key="<short_slug>", category="preference|rule|tech_stack|fact", content="<concise memory summary>"]]',
            "2. Note: Non-admin users can ONLY save 'private' memories. Only system administrators can set 'global' memories.",
            "3. If a past memory is superseded or changed, update it with the same key or use [[UPDATE_MEMORY: key=\"<slug>\", content=\"<new>\"]].",
            "4. If the user explicitly asks to forget something, use [[DELETE_MEMORY: key=\"<slug>\"]].",
            "5. Do NOT store temporary or trivial small-talk. Only store enduring and valuable knowledge."
        ])

        return "\n".join(lines)

    def get_stats(self) -> Dict[str, Any]:
        with self._get_connection() as conn:
            cur = conn.execute("SELECT COUNT(*) as cnt FROM memories WHERE type = 'global'")
            global_count = cur.fetchone()["cnt"]

            cur = conn.execute("SELECT COUNT(*) as cnt FROM memories WHERE type = 'private'")
            private_count = cur.fetchone()["cnt"]

            cur = conn.execute("SELECT COUNT(DISTINCT user_id) as cnt FROM memories WHERE type = 'private'")
            user_count = cur.fetchone()["cnt"]

            return {
                "total_memories": global_count + private_count,
                "global_memories": global_count,
                "private_memories": private_count,
                "users_with_memory": user_count,
            }


memory_manager = MemoryManager()