
# 素材分类表 - PostgreSQL 版本

import uuid
from typing import Dict, List, Any, Optional
from datetime import datetime
from contextlib import contextmanager
from python_core.config import settings
from python_core.utils.logger import setup_logger
from .types import ResourceCategory

# 尝试导入 psycopg2，如果失败则提供友好的错误信息
try:
    import psycopg2
    import psycopg2.extras
    PSYCOPG2_AVAILABLE = True
except ImportError as e:
    PSYCOPG2_AVAILABLE = False
    PSYCOPG2_ERROR = str(e)

logger = setup_logger(__name__)

class ResourceCategoryPostgres:
    """
    素材分类表类 - PostgreSQL 版本
    基于 PostgreSQL 数据库实现的素材分类管理功能
    """

    def __init__(self):
        # 检查 psycopg2 是否可用
        if not PSYCOPG2_AVAILABLE:
            error_msg = f"""
PostgreSQL 驱动 psycopg2 未安装。请安装：

方案1（推荐）：
    pip install psycopg2-binary

方案2：
    # Ubuntu/Debian
    sudo apt-get install postgresql libpq-dev python3-dev
    pip install psycopg2

    # CentOS/RHEL
    sudo yum install postgresql postgresql-devel python3-devel
    pip install psycopg2

    # macOS
    brew install postgresql
    pip install psycopg2

原始错误: {PSYCOPG2_ERROR}
"""
            logger.error(error_msg)
            raise ImportError(error_msg)

        self.db_url = settings.db
        self.table_name = "resource_categories"

        # 初始化素材分类表
        self._init_resource_category_table()

    @contextmanager
    def _get_connection(self):
        """获取数据库连接的上下文管理器"""
        conn = None
        try:
            conn = psycopg2.connect(self.db_url)
            yield conn
        except Exception as e:
            if conn:
                conn.rollback()
            logger.error(f"Database connection error: {e}")
            raise e
        finally:
            if conn:
                conn.close()

    def _init_resource_category_table(self):
        """初始化素材分类表"""
        try:
            with self._get_connection() as conn:
                with conn.cursor() as cursor:
                    # 创建素材分类表（如果不存在）
                    create_table_sql = """
                    CREATE TABLE IF NOT EXISTS resource_categories (
                        id VARCHAR(36) PRIMARY KEY,
                        title VARCHAR(255) NOT NULL,
                        ai_prompt TEXT DEFAULT '',
                        color VARCHAR(7) DEFAULT '#FF5733',
                        is_cloud BOOLEAN DEFAULT FALSE,
                        user_id VARCHAR(36) NOT NULL,
                        is_active BOOLEAN DEFAULT TRUE,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    );
                    """
                    cursor.execute(create_table_sql)

                    # 创建索引
                    indexes = [
                        "CREATE INDEX IF NOT EXISTS idx_resource_categories_title ON resource_categories(title);",
                        "CREATE INDEX IF NOT EXISTS idx_resource_categories_user_id ON resource_categories(user_id);",
                        "CREATE INDEX IF NOT EXISTS idx_resource_categories_is_cloud ON resource_categories(is_cloud);",
                        "CREATE INDEX IF NOT EXISTS idx_resource_categories_is_active ON resource_categories(is_active);",
                        "CREATE INDEX IF NOT EXISTS idx_resource_categories_created_at ON resource_categories(created_at);",
                        "CREATE INDEX IF NOT EXISTS idx_resource_categories_user_title ON resource_categories(user_id, title);"
                    ]

                    for index_sql in indexes:
                        cursor.execute(index_sql)

                    conn.commit()
                    logger.info("Resource category table initialized")

        except Exception as e:
            logger.error(f"Failed to initialize resource category table: {e}")
            raise e

    def create_category(self, category: ResourceCategory) -> str:
        """
        创建素材分类

        Args:
            category: 素材分类信息

        Returns:
            分类ID
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor() as cursor:
                    # 检查分类名称是否已存在（同一用户下）
                    cursor.execute(
                        "SELECT id FROM resource_categories WHERE title = %s AND user_id = %s AND is_active = TRUE",
                        (category.title, category.user_id)
                    )
                    if cursor.fetchone():
                        raise ValueError(f"Category title '{category.title}' already exists for this user")

                    # 生成分类ID（如果没有提供）
                    category_id = category.id or str(uuid.uuid4())

                    # 插入分类记录
                    insert_sql = """
                    INSERT INTO resource_categories (
                        id, title, ai_prompt, color, is_cloud, user_id, is_active, created_at, updated_at
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
                    """

                    now = datetime.now()
                    cursor.execute(insert_sql, (
                        category_id,
                        category.title,
                        category.ai_prompt,
                        category.color,
                        category.is_cloud,
                        category.user_id,
                        category.is_active,
                        now,
                        now
                    ))

                    conn.commit()
                    logger.info(f"Created resource category: {category.title} (ID: {category_id})")
                    return category_id

        except Exception as e:
            logger.error(f"Failed to create resource category '{category.title}': {e}")
            raise e

    def get_category_by_id(self, category_id: str) -> Optional[ResourceCategory]:
        """
        根据ID获取素材分类

        Args:
            category_id: 分类ID

        Returns:
            分类信息，如果不存在返回None
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    cursor.execute(
                        "SELECT * FROM resource_categories WHERE id = %s AND is_active = TRUE",
                        (category_id,)
                    )
                    row = cursor.fetchone()

                    if row:
                        return self._row_to_category(row)
                    return None

        except Exception as e:
            logger.error(f"Failed to get resource category by ID '{category_id}': {e}")
            return None

    def get_category_by_title(self, title: str, user_id: str = None) -> Optional[ResourceCategory]:
        """
        根据标题获取素材分类

        Args:
            title: 分类标题
            user_id: 用户ID（可选，用于过滤用户分类）

        Returns:
            分类信息，如果不存在返回None
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    if user_id:
                        cursor.execute(
                            "SELECT * FROM resource_categories WHERE title = %s AND user_id = %s AND is_active = TRUE LIMIT 1",
                            (title, user_id)
                        )
                    else:
                        cursor.execute(
                            "SELECT * FROM resource_categories WHERE title = %s AND is_active = TRUE LIMIT 1",
                            (title,)
                        )

                    row = cursor.fetchone()
                    if row:
                        return self._row_to_category(row)
                    return None

        except Exception as e:
            logger.error(f"Failed to get resource category by title '{title}': {e}")
            return None

    def update_category(self, category_id: str, updates: Dict[str, Any]) -> bool:
        """
        更新素材分类信息

        Args:
            category_id: 分类ID
            updates: 要更新的字段

        Returns:
            更新成功返回True
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor() as cursor:
                    # 检查分类是否存在（不限制 is_active 状态）
                    cursor.execute(
                        "SELECT id FROM resource_categories WHERE id = %s",
                        (category_id,)
                    )
                    if not cursor.fetchone():
                        logger.warning(f"Resource category not found: {category_id}")
                        return False

                    # 移除不应该直接更新的字段
                    protected_fields = ["id", "created_at"]
                    filtered_updates = {k: v for k, v in updates.items() if k not in protected_fields}

                    if not filtered_updates:
                        logger.warning("No valid fields to update")
                        return False

                    # 添加更新时间
                    filtered_updates["updated_at"] = datetime.now()

                    # 构建更新SQL
                    set_clauses = []
                    values = []
                    for key, value in filtered_updates.items():
                        set_clauses.append(f"{key} = %s")
                        values.append(value)

                    values.append(category_id)  # WHERE条件的参数

                    update_sql = f"UPDATE resource_categories SET {', '.join(set_clauses)} WHERE id = %s"
                    cursor.execute(update_sql, values)

                    conn.commit()

                    if cursor.rowcount > 0:
                        logger.info(f"Updated resource category: {category_id}")
                        return True
                    else:
                        logger.warning(f"No rows updated for resource category: {category_id}")
                        return False

        except Exception as e:
            logger.error(f"Failed to update resource category '{category_id}': {e}")
            return False

    def delete_category(self, category_id: str, hard_delete: bool = True) -> bool:
        """
        删除素材分类

        Args:
            category_id: 分类ID
            hard_delete: 是否硬删除（True）或软删除（False）

        Returns:
            删除成功返回True
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor() as cursor:
                    if hard_delete:
                        # 硬删除：直接从数据库中删除记录
                        cursor.execute("DELETE FROM resource_categories WHERE id = %s", (category_id,))
                    else:
                        # 软删除：设置 is_active = FALSE
                        cursor.execute(
                            "UPDATE resource_categories SET is_active = FALSE, updated_at = %s WHERE id = %s",
                            (datetime.now(), category_id)
                        )

                    conn.commit()

                    if cursor.rowcount > 0:
                        delete_type = "hard deleted" if hard_delete else "soft deleted"
                        logger.info(f"Resource category {delete_type}: {category_id}")
                        return True
                    else:
                        logger.warning(f"No resource category found to delete: {category_id}")
                        return False

        except Exception as e:
            logger.error(f"Failed to delete resource category '{category_id}': {e}")
            return False

    def get_categories_by_user(self, user_id: str, include_cloud: bool = True, limit: int = 100) -> List[ResourceCategory]:
        """
        获取用户的素材分类列表

        Args:
            user_id: 用户ID
            include_cloud: 是否包含云端公共分类
            limit: 最大返回数量

        Returns:
            分类列表
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    if include_cloud:
                        sql = """
                        SELECT * FROM resource_categories
                        WHERE (user_id = %s OR is_cloud = TRUE) AND is_active = TRUE
                        ORDER BY created_at DESC
                        LIMIT %s
                        """
                        cursor.execute(sql, (user_id, limit))
                    else:
                        sql = """
                        SELECT * FROM resource_categories
                        WHERE user_id = %s AND is_active = TRUE
                        ORDER BY created_at DESC
                        LIMIT %s
                        """
                        cursor.execute(sql, (user_id, limit))

                    rows = cursor.fetchall()
                    categories = []

                    for row in rows:
                        category = self._row_to_category(row)
                        categories.append(category)

                    return categories

        except Exception as e:
            logger.error(f"Failed to get resource categories for user '{user_id}': {e}")
            return []

    def search_categories(self, query: str, user_id: str = None, include_cloud: bool = True, limit: int = 50) -> List[ResourceCategory]:
        """
        搜索素材分类

        Args:
            query: 搜索关键词（匹配标题、AI提示词）
            user_id: 用户ID（可选，用于过滤用户分类）
            include_cloud: 是否包含云端公共分类
            limit: 最大返回数量

        Returns:
            匹配的分类列表
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    search_pattern = f"%{query}%"

                    if user_id and include_cloud:
                        sql = """
                        SELECT * FROM resource_categories
                        WHERE (user_id = %s OR is_cloud = TRUE) AND is_active = TRUE
                        AND (title ILIKE %s OR ai_prompt ILIKE %s)
                        ORDER BY created_at DESC
                        LIMIT %s
                        """
                        cursor.execute(sql, (user_id, search_pattern, search_pattern, limit))
                    elif user_id:
                        sql = """
                        SELECT * FROM resource_categories
                        WHERE user_id = %s AND is_active = TRUE
                        AND (title ILIKE %s OR ai_prompt ILIKE %s)
                        ORDER BY created_at DESC
                        LIMIT %s
                        """
                        cursor.execute(sql, (user_id, search_pattern, search_pattern, limit))
                    else:
                        sql = """
                        SELECT * FROM resource_categories
                        WHERE is_active = TRUE
                        AND (title ILIKE %s OR ai_prompt ILIKE %s)
                        ORDER BY created_at DESC
                        LIMIT %s
                        """
                        cursor.execute(sql, (search_pattern, search_pattern, limit))

                    rows = cursor.fetchall()
                    categories = []

                    for row in rows:
                        category = self._row_to_category(row)
                        categories.append(category)

                    return categories

        except Exception as e:
            logger.error(f"Failed to search resource categories with query '{query}': {e}")
            return []

    def get_cloud_categories(self, limit: int = 100) -> List[ResourceCategory]:
        """
        获取云端公共分类

        Args:
            limit: 最大返回数量

        Returns:
            云端分类列表
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    sql = """
                    SELECT * FROM resource_categories
                    WHERE is_cloud = TRUE AND is_active = TRUE
                    ORDER BY created_at DESC
                    LIMIT %s
                    """
                    cursor.execute(sql, (limit,))

                    rows = cursor.fetchall()
                    categories = []

                    for row in rows:
                        category = self._row_to_category(row)
                        categories.append(category)

                    return categories

        except Exception as e:
            logger.error(f"Failed to get cloud resource categories: {e}")
            return []

    def get_category_count(self, user_id: str = None, include_cloud: bool = True) -> int:
        """
        获取分类数量

        Args:
            user_id: 用户ID（可选）
            include_cloud: 是否包含云端公共分类

        Returns:
            分类数量
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor() as cursor:
                    if user_id is None:
                        cursor.execute("SELECT COUNT(*) FROM resource_categories WHERE is_active = TRUE")
                    elif include_cloud:
                        cursor.execute(
                            "SELECT COUNT(*) FROM resource_categories WHERE (user_id = %s OR is_cloud = TRUE) AND is_active = TRUE",
                            (user_id,)
                        )
                    else:
                        cursor.execute(
                            "SELECT COUNT(*) FROM resource_categories WHERE user_id = %s AND is_active = TRUE",
                            (user_id,)
                        )

                    result = cursor.fetchone()
                    return result[0] if result else 0

        except Exception as e:
            logger.error(f"Failed to get resource category count: {e}")
            return 0

    def batch_create_categories(self, categories: List[ResourceCategory]) -> Dict[str, Any]:
        """
        批量创建分类

        Args:
            categories: 分类列表

        Returns:
            创建结果统计
        """
        try:
            success_count = 0
            failed_count = 0
            failed_categories = []

            for category in categories:
                try:
                    # 检查是否已存在（根据标题和用户ID）
                    existing = self.get_category_by_title(category.title, category.user_id)
                    if existing:
                        logger.warning(f"Category already exists, skipping: {category.title}")
                        continue

                    # 创建分类
                    self.create_category(category)
                    success_count += 1

                except Exception as e:
                    logger.error(f"Failed to create category '{category.title}': {e}")
                    failed_count += 1
                    failed_categories.append({
                        "title": category.title,
                        "error": str(e)
                    })

            result = {
                "total": len(categories),
                "success": success_count,
                "failed": failed_count,
                "failed_categories": failed_categories
            }

            logger.info(f"Batch create completed: {success_count} success, {failed_count} failed")
            return result

        except Exception as e:
            logger.error(f"Failed to batch create categories: {e}")
            return {
                "total": len(categories),
                "success": 0,
                "failed": len(categories),
                "error": str(e)
            }

    def get_categories_by_color(self, color: str, user_id: str = None, include_cloud: bool = True, limit: int = 50) -> List[ResourceCategory]:
        """
        根据颜色获取分类

        Args:
            color: 颜色值（hex格式）
            user_id: 用户ID（可选）
            include_cloud: 是否包含云端公共分类
            limit: 最大返回数量

        Returns:
            匹配的分类列表
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
                    if user_id and include_cloud:
                        sql = """
                        SELECT * FROM resource_categories
                        WHERE (user_id = %s OR is_cloud = TRUE) AND is_active = TRUE
                        AND color = %s
                        ORDER BY created_at DESC
                        LIMIT %s
                        """
                        cursor.execute(sql, (user_id, color, limit))
                    elif user_id:
                        sql = """
                        SELECT * FROM resource_categories
                        WHERE user_id = %s AND is_active = TRUE
                        AND color = %s
                        ORDER BY created_at DESC
                        LIMIT %s
                        """
                        cursor.execute(sql, (user_id, color, limit))
                    else:
                        sql = """
                        SELECT * FROM resource_categories
                        WHERE is_active = TRUE AND color = %s
                        ORDER BY created_at DESC
                        LIMIT %s
                        """
                        cursor.execute(sql, (color, limit))

                    rows = cursor.fetchall()
                    categories = []

                    for row in rows:
                        category = self._row_to_category(row)
                        categories.append(category)

                    return categories

        except Exception as e:
            logger.error(f"Failed to get categories by color '{color}': {e}")
            return []

    def activate_category(self, category_id: str) -> bool:
        """
        激活分类（设置 is_active = TRUE）

        Args:
            category_id: 分类ID

        Returns:
            激活成功返回True
        """
        try:
            return self.update_category(category_id, {"is_active": True})
        except Exception as e:
            logger.error(f"Failed to activate category '{category_id}': {e}")
            return False

    def deactivate_category(self, category_id: str) -> bool:
        """
        停用分类（设置 is_active = FALSE）

        Args:
            category_id: 分类ID

        Returns:
            停用成功返回True
        """
        try:
            return self.update_category(category_id, {"is_active": False})
        except Exception as e:
            logger.error(f"Failed to deactivate category '{category_id}': {e}")
            return False

    # 辅助方法
    def _row_to_category(self, row: Dict[str, Any]) -> ResourceCategory:
        """将数据库行转换为ResourceCategory对象"""
        return ResourceCategory(
            id=row['id'],
            title=row['title'],
            ai_prompt=row['ai_prompt'],
            color=row['color'],
            is_cloud=row['is_cloud'],
            user_id=row['user_id'],
            is_active=row['is_active'],
            created_at=row['created_at'].isoformat() if row['created_at'] else "",
            updated_at=row['updated_at'].isoformat() if row['updated_at'] else ""
        )


# 创建全局素材分类表实例
resource_category_table = ResourceCategoryPostgres()