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', 'user', 'project'
    user_id: Optional[int] = None  # None for global, chat_id for user/project
    project_name: Optional[str] = None  # None for global/user, project_name for project
    key: str = ""
    category: str = "general"  # '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_user(self) -> bool:
        return self.type == "user"

    @property
    def is_project(self) -> bool:
        return self.type == "project"


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, migrating from old schema if needed."""
        try:
            with self._get_connection() as conn:
                # Check if old memories table exists and needs migration
                cur = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='memories';")
                table_exists = cur.fetchone() is not None

                if table_exists:
                    # Check columns in existing table
                    cur = conn.execute("PRAGMA table_info(memories);")
                    cols = [row["name"] for row in cur.fetchall()]
                    if "project_name" not in cols:
                        logger.info("Migrating memories table to support 3-tier hierarchy (adding project_name)...")
                        # Migrate table
                        conn.execute("""
                            CREATE TABLE memories_new (
                                id INTEGER PRIMARY KEY AUTOINCREMENT,
                                type TEXT NOT NULL CHECK(type IN ('global', 'user', 'project', 'private')),
                                user_id INTEGER,
                                project_name TEXT,
                                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
                            );
                        """)
                        conn.execute("""
                            INSERT INTO memories_new (id, type, user_id, project_name, key, category, content, importance, created_by, created_at, updated_at)
                            SELECT id, CASE WHEN type='private' THEN 'user' ELSE type END, user_id, NULL, key, category, content, importance, created_by, created_at, updated_at
                            FROM memories;
                        """)
                        conn.execute("DROP TABLE memories;")
                        conn.execute("ALTER TABLE memories_new RENAME TO memories;")
                else:
                    conn.execute("""
                        CREATE TABLE memories (
                            id INTEGER PRIMARY KEY AUTOINCREMENT,
                            type TEXT NOT NULL CHECK(type IN ('global', 'user', 'project', 'private')),
                            user_id INTEGER,
                            project_name TEXT,
                            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
                        );
                    """)

                # Create partial unique indexes for each tier
                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_user_key 
                    ON memories(type, user_id, key) WHERE type IN ('user', 'private');
                """)
                conn.execute("""
                    CREATE UNIQUE INDEX IF NOT EXISTS idx_memories_project_key 
                    ON memories(type, user_id, project_name, key) WHERE type = 'project';
                """)
                conn.execute("""
                    CREATE INDEX IF NOT EXISTS idx_memories_lookup 
                    ON memories(type, user_id, project_name);
                """)
                conn.execute("""
                    CREATE INDEX IF NOT EXISTS idx_memories_category 
                    ON memories(category);
                """)
            logger.info("3-Tier 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,
        project_name: Optional[str] = None,
        category: str = "general",
        importance: int = 1,
        created_by: Optional[int] = None,
    ) -> Tuple[MemoryItem, bool]:
        """
        Saves or updates a memory in one of the 3 tiers (global, user, project).
        Returns: (MemoryItem, is_created: bool)
        """
        type_ = type_.lower().strip()
        if type_ == "private":
            type_ = "user"
        if type_ not in ("global", "user", "project"):
            raise ValueError("Memory type must be 'global', 'user', or 'project'")

        if type_ == "global":
            user_id = None
            project_name = None
        elif type_ == "user":
            if user_id is None:
                raise ValueError("User memory requires a valid user_id")
            project_name = None
        elif type_ == "project":
            if user_id is None or not project_name:
                raise ValueError("Project memory requires both user_id and project_name")
            project_name = project_name.strip()

        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 existing item exists
            if type_ == "global":
                cur = conn.execute(
                    "SELECT id, created_at, created_by FROM memories WHERE type = 'global' AND key = ?",
                    (key,),
                )
            elif type_ == "user":
                cur = conn.execute(
                    "SELECT id, created_at, created_by FROM memories WHERE type IN ('user', 'private') AND user_id = ? AND key = ?",
                    (user_id, key),
                )
            else:  # project
                cur = conn.execute(
                    "SELECT id, created_at, created_by FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? AND key = ?",
                    (user_id, project_name, 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 type = ?, category = ?, content = ?, importance = ?, updated_at = ?
                    WHERE id = ?
                    """,
                    (type_, 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, project_name, key, category, content, importance, created_by, created_at, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (type_, user_id, project_name, 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,
            project_name=project_name,
            key=key,
            category=category,
            content=content,
            importance=importance,
            created_by=creator,
            created_at=created_at,
            updated_at=now,
        )
        logger.info(
            "%s memory [id=%d, type=%s, key=%s, user=%s, proj=%s]",
            "Created" if is_created else "Updated",
            mem_id,
            type_,
            key,
            user_id,
            project_name,
        )
        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, project_name: Optional[str] = None
    ) -> Optional[MemoryItem]:
        type_ = type_.lower().strip()
        if type_ == "private":
            type_ = "user"
        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,))
            elif type_ == "user":
                cur = conn.execute(
                    "SELECT * FROM memories WHERE type IN ('user', 'private') AND user_id = ? AND key = ?",
                    (user_id, key),
                )
            else:  # project
                cur = conn.execute(
                    "SELECT * FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? AND key = ?",
                    (user_id, project_name, 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 IN ('user', '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 IN ('user', '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 get_project_memories(
        self, user_id: int, project_name: str, 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 = 'project' AND user_id = ? AND project_name = ? AND category = ? ORDER BY importance DESC, updated_at DESC LIMIT ?",
                    (user_id, project_name, category.lower().strip(), limit),
                )
            else:
                cur = conn.execute(
                    "SELECT * FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? ORDER BY importance DESC, updated_at DESC LIMIT ?",
                    (user_id, project_name, 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 in ("user", "project", "private") and mem.user_id != user_id and not is_admin:
                raise PermissionError("Cannot delete another user's 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,
        project_name: Optional[str] = None,
        is_admin: bool = False,
    ) -> bool:
        type_ = type_.lower().strip()
        if type_ == "private":
            type_ = "user"
        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,))
            elif type_ == "user":
                if user_id is None:
                    return False
                cur = conn.execute(
                    "DELETE FROM memories WHERE type IN ('user', 'private') AND user_id = ? AND key = ?",
                    (user_id, key),
                )
            else:  # project
                if user_id is None or not project_name:
                    return False
                cur = conn.execute(
                    "DELETE FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? AND key = ?",
                    (user_id, project_name, key),
                )
            deleted = cur.rowcount > 0
            if deleted:
                logger.info("Deleted memory by key [key=%s, type=%s, user=%s, proj=%s]", key, type_, user_id, project_name)
            return deleted

    def clear_memories(
        self,
        type_: str,
        user_id: Optional[int] = None,
        project_name: Optional[str] = None,
        is_admin: bool = False,
    ) -> int:
        type_ = type_.lower().strip()
        if type_ == "private":
            type_ = "user"
        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'")
            elif type_ == "user":
                if user_id is None:
                    return 0
                cur = conn.execute("DELETE FROM memories WHERE type IN ('user', 'private') AND user_id = ?", (user_id,))
            else:  # project
                if user_id is None or not project_name:
                    return 0
                cur = conn.execute(
                    "DELETE FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ?",
                    (user_id, project_name),
                )
            count = cur.rowcount
            logger.info("Cleared %d %s memories for user %s (proj: %s)", count, type_, user_id, project_name)
            return count

    def search_memories(
        self,
        query: str,
        user_id: Optional[int] = None,
        project_name: Optional[str] = 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_ == "user" and user_id is not None:
                cur = conn.execute(
                    """
                    SELECT * FROM memories 
                    WHERE type IN ('user', '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 type_ == "project" and user_id is not None and project_name:
                cur = conn.execute(
                    """
                    SELECT * FROM memories 
                    WHERE type = 'project' AND user_id = ? AND project_name = ? AND (key LIKE ? OR content LIKE ? OR category LIKE ?)
                    ORDER BY importance DESC, updated_at DESC LIMIT ?
                    """,
                    (user_id, project_name, q, q, q, limit),
                )
            elif user_id is not None:
                cur = conn.execute(
                    """
                    SELECT * FROM memories 
                    WHERE (type = 'global' OR (type IN ('user', 'private') AND user_id = ?) OR (type = 'project' AND user_id = ? AND project_name = ?)) 
                      AND (key LIKE ? OR content LIKE ? OR category LIKE ?)
                    ORDER BY importance DESC, updated_at DESC LIMIT ?
                    """,
                    (user_id, user_id, project_name or "", 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, project_name: Optional[str] = None
    ) -> Optional[str]:
        """
        Formats all 3 tiers of memories into a structured, concise instruction block for AI:
        1. 🌐 Global System Rules
        2. 👤 User Personal Preferences & Profile
        3. 📁 Active Project Specific Knowledge
        """
        globals_list = self.get_global_memories(limit=50)
        users_list = self.get_user_memories(user_id=user_id, limit=50) if user_id else []
        projects_list = (
            self.get_project_memories(user_id=user_id, project_name=project_name, limit=50)
            if user_id and project_name
            else []
        )

        if not globals_list and not users_list and not projects_list:
            return None

        lines = [
            "[SYSTEM INSTRUCTION: AI PERSISTENT HIERARCHICAL MEMORY & CONTINUOUS LEARNING]",
            "You have access to persistent 3-tier long-term memory containing verified system rules, project context, and user profile.",
            "Use this knowledge to maintain continuity and customize all your responses accordingly.\n",
            "⚖️ STRICT PRECEDENCE HIERARCHY (سلسله‌مراتب قطعی اولویت و حل تعارض):",
            "1. 🥇 Tier 1 (Highest Authority): Global System Rules ALWAYS strictly supersede all Project and User instructions.",
            "2. 🥈 Tier 2 (Intermediate Authority): Active Project Knowledge & Rules ALWAYS supersede general User preferences inside this project.",
            "3. 🥉 Tier 3 (Baseline Authority): User Personal Preferences apply across projects as defaults when not overridden by Project or Global rules.\n"
        ]

        if globals_list:
            lines.append("🌐 1. GLOBAL SYSTEM RULES & KNOWLEDGE (قوانین عمومی سیستم - بالاترین اولویت):")
            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 projects_list:
            lines.append(f"📁 2. ACTIVE PROJECT KNOWLEDGE [پروژه: {project_name}] (دانش و تصمیمات اختصاصی این پروژه - اولویت دوم):")
            for m in projects_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 users_list:
            lines.append("👤 3. USER PERSONAL PROFILE & PREFERENCES (ترجیحات و مشخصات شخصی کاربر - اولویت پایه):")
            for m in users_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 user shares durable facts, choose the appropriate tier:",
            '   • User-level (developer habits, preferred tools, personal info): type="user"',
            '   • Project-level (architecture, database schema, libraries, API designs, fixed bugs in this project): type="project"',
            '   Action Tag: [[SAVE_MEMORY: type="user|project", key="<short_slug>", category="preference|rule|tech_stack|fact", content="<concise summary>"]]',
            "2. Note: Global memory is read-only for AI and managed exclusively by system administrators via the Telegram Bot UI.",
            "3. If past memory is superseded, update it: [[UPDATE_MEMORY: type=\"user|project\", key=\"<slug>\", content=\"<new_content>\"]].",
            "4. If user asks to forget something: [[DELETE_MEMORY: type=\"user|project\", 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, user_id: Optional[int] = None, project_name: Optional[str] = None) -> 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 IN ('user', 'private')")
            user_count = cur.fetchone()["cnt"]

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

            cur = conn.execute("SELECT COUNT(DISTINCT user_id) as cnt FROM memories WHERE type IN ('user', 'project', 'private')")
            users_with_mem = cur.fetchone()["cnt"]

            this_user_count = 0
            this_project_count = 0
            if user_id:
                cur = conn.execute("SELECT COUNT(*) as cnt FROM memories WHERE type IN ('user', 'private') AND user_id = ?", (user_id,))
                this_user_count = cur.fetchone()["cnt"]
                if project_name:
                    cur = conn.execute("SELECT COUNT(*) as cnt FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ?", (user_id, project_name))
                    this_project_count = cur.fetchone()["cnt"]

            return {
                "total_memories": global_count + user_count + project_count,
                "global_memories": global_count,
                "user_memories": user_count,
                "project_memories": project_count,
                "users_with_memory": users_with_mem,
                "current_user_memories": this_user_count,
                "current_project_memories": this_project_count,
            }


memory_manager = MemoryManager()