database.py

import asyncio
import json
import os
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Any, Dict, Iterable, List, Optional, Sequence

import aiomysql
from aiomysql.cursors import DictCursor

from config import (
    IMAGES_DIR,
    logger,
    MYSQL_HOST,
    MYSQL_PORT,
    MYSQL_DB,
    MYSQL_USER,
    MYSQL_PASSWORD,
    MYSQL_POOL_MIN_SIZE,
    MYSQL_POOL_MAX_SIZE,
)

_pool: Optional[aiomysql.Pool] = None
_pool_lock = asyncio.Lock()


async def init_mysql_pool() -> aiomysql.Pool:
  """初始化 MySQL 连接池"""
  global _pool
  if _pool is not None:
    return _pool

  async with _pool_lock:
    if _pool is not None:
      return _pool
    try:
      _pool = await aiomysql.create_pool(
          host=MYSQL_HOST,
          port=MYSQL_PORT,
          user=MYSQL_USER,
          password=MYSQL_PASSWORD,
          db=MYSQL_DB,
          minsize=MYSQL_POOL_MIN_SIZE,
          maxsize=MYSQL_POOL_MAX_SIZE,
          autocommit=True,
          charset="utf8mb4",
          cursorclass=DictCursor,
      )
      logger.info(
          "MySQL 连接池初始化成功，host=%s db=%s",
          MYSQL_HOST,
          MYSQL_DB,
      )
    except Exception as exc:
      logger.error(f"初始化 MySQL 连接池失败: {exc}")
      raise
  return _pool


async def close_mysql_pool() -> None:
  """关闭 MySQL 连接池"""
  global _pool
  if _pool is None:
    return

  async with _pool_lock:
    if _pool is None:
      return
    _pool.close()
    await _pool.wait_closed()
    _pool = None
    logger.info("MySQL 连接池已关闭")


@asynccontextmanager
async def get_connection():
  """获取连接池中的连接"""
  if _pool is None:
    await init_mysql_pool()
  assert _pool is not None
  conn = await _pool.acquire()
  try:
    yield conn
  finally:
    _pool.release(conn)


async def execute(query: str,
    params: Sequence[Any] | Dict[str, Any] | None = None) -> None:
  """执行写入类 SQL"""
  async with get_connection() as conn:
    async with conn.cursor() as cursor:
      await cursor.execute(query, params or ())


async def fetch_all(
    query: str, params: Sequence[Any] | Dict[str, Any] | None = None
) -> List[Dict[str, Any]]:
  """执行查询并返回全部结果"""
  async with get_connection() as conn:
    async with conn.cursor() as cursor:
      await cursor.execute(query, params or ())
      rows = await cursor.fetchall()
      return list(rows)


def _serialize_extra(extra: Optional[Dict[str, Any]]) -> Optional[str]:
  if extra is None:
    return None
  try:
    return json.dumps(extra, ensure_ascii=False)
  except Exception:
    logger.warning("无法序列化 extra_metadata，已忽略")
    return None


async def upsert_image_record(
    *,
    file_path: str,
    category: str,
    nickname: Optional[str],
    score: float,
    is_cropped_face: bool,
    size_bytes: int,
    last_modified: datetime,
    bos_uploaded: bool,
    hostname: Optional[str] = None,
    extra_metadata: Optional[Dict[str, Any]] = None,
) -> None:
  """写入或更新图片记录"""
  query = """
        INSERT INTO tpl_app_processed_images (
            file_path,
            category,
            nickname,
            score,
            is_cropped_face,
            size_bytes,
            last_modified,
            bos_uploaded,
            hostname,
            extra_metadata
        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
            category = VALUES(category),
            nickname = VALUES(nickname),
            score = VALUES(score),
            is_cropped_face = VALUES(is_cropped_face),
            size_bytes = VALUES(size_bytes),
            last_modified = VALUES(last_modified),
            bos_uploaded = VALUES(bos_uploaded),
            hostname = VALUES(hostname),
            extra_metadata = VALUES(extra_metadata),
            updated_at = CURRENT_TIMESTAMP
    """
  extra_value = _serialize_extra(extra_metadata)
  await execute(
      query,
      (
        file_path,
        category,
        nickname,
        score,
        1 if is_cropped_face else 0,
        size_bytes,
        last_modified,
        1 if bos_uploaded else 0,
        hostname,
        extra_value,
      ),
  )


async def fetch_paged_image_records(
    *,
    category: Optional[str],
    nickname: Optional[str],
    offset: int,
    limit: int,
) -> List[Dict[str, Any]]:
  """按条件分页查询图片记录"""
  where_clauses: List[str] = []
  params: List[Any] = []
  if category and category != "all":
    where_clauses.append("category = %s")
    params.append(category)
  if nickname:
    where_clauses.append("nickname = %s")
    params.append(nickname)
  where_sql = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
  query = f"""
        SELECT
            file_path,
            category,
            nickname,
            score,
            is_cropped_face,
            size_bytes,
            last_modified,
            bos_uploaded,
            hostname
        FROM tpl_app_processed_images
        {where_sql}
        ORDER BY last_modified DESC, id DESC
        LIMIT %s OFFSET %s
    """
  params.extend([limit, offset])
  return await fetch_all(query, params)


async def count_image_records(
    *, category: Optional[str], nickname: Optional[str]
) -> int:
  """按条件统计图片记录数量"""
  where_clauses: List[str] = []
  params: List[Any] = []
  if category and category != "all":
    where_clauses.append("category = %s")
    params.append(category)
  if nickname:
    where_clauses.append("nickname = %s")
    params.append(nickname)
  where_sql = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
  query = f"SELECT COUNT(*) AS total FROM tpl_app_processed_images {where_sql}"
  rows = await fetch_all(query, params)
  if not rows:
    return 0
  return int(rows[0].get("total", 0) or 0)


async def fetch_today_category_counts() -> List[Dict[str, Any]]:
  """统计当天按类别分组的数量"""
  query = """
      SELECT
          COALESCE(category, 'unknown') AS category,
          COUNT(*) AS count
      FROM tpl_app_processed_images
      WHERE last_modified >= CURDATE()
        AND last_modified < DATE_ADD(CURDATE(), INTERVAL 1 DAY)
      GROUP BY COALESCE(category, 'unknown')
    """
  rows = await fetch_all(query)
  return [
    {
      "category": str(row.get("category") or "unknown"),
      "count": int(row.get("count") or 0),
    }
    for row in rows
  ]


async def fetch_records_by_paths(file_paths: Iterable[str]) -> Dict[
  str, Dict[str, Any]]:
  """根据文件名批量查询图片记录"""
  paths = list({path for path in file_paths if path})
  if not paths:
    return {}

  placeholders = ", ".join(["%s"] * len(paths))
  query = f"""
        SELECT
            file_path,
            category,
            nickname,
            score,
            is_cropped_face,
            size_bytes,
            last_modified,
            bos_uploaded,
            hostname
        FROM tpl_app_processed_images
        WHERE file_path IN ({placeholders})
    """
  rows = await fetch_all(query, paths)
  return {row["file_path"]: row for row in rows}


_IMAGES_DIR_ABS = os.path.abspath(os.path.expanduser(IMAGES_DIR))


def _normalize_file_path(file_path: str) -> Optional[str]:
  """将绝对路径转换为相对 IMAGES_DIR 的文件名"""
  try:
    abs_path = os.path.abspath(os.path.expanduser(file_path))
    if os.path.isdir(abs_path):
      return None
    if os.path.commonpath([_IMAGES_DIR_ABS, abs_path]) != _IMAGES_DIR_ABS:
      return os.path.basename(abs_path)
    rel_path = os.path.relpath(abs_path, _IMAGES_DIR_ABS)
    return rel_path.replace("\\", "/")
  except Exception:
    return None


def infer_category_from_filename(filename: str, default: str = "other") -> str:
  """根据文件名推断类别"""
  lower_name = filename.lower()
  if "_face_" in lower_name:
    return "face"
  if lower_name.endswith("_original.webp") or "_original" in lower_name:
    return "original"
  if "_restore" in lower_name:
    return "restore"
  if "_upcolor" in lower_name:
    return "upcolor"
  if "_compress" in lower_name:
    return "compress"
  if "_upscale" in lower_name:
    return "upscale"
  if "_anime_style_" in lower_name:
    return "anime_style"
  if "_grayscale" in lower_name:
    return "grayscale"
  if "_id_photo" in lower_name or "_save_id_photo" in lower_name:
    return "id_photo"
  if "_grid_" in lower_name:
    return "grid"
  if "_rvm_id_photo" in lower_name:
    return "rvm"
  if "_celebrity_" in lower_name or "_celebrity" in lower_name:
    return "celebrity"
  return default


from config import HOSTNAME

async def record_image_creation(
    *,
    file_path: str,
    nickname: Optional[str],
    score: float = 0.0,
    category: Optional[str] = None,
    bos_uploaded: bool = False,
    extra_metadata: Optional[Dict[str, Any]] = None,
) -> None:
  """
  记录图片元数据到数据库，如果数据库不可用则静默忽略。
  :param file_path: 绝对或相对文件路径
  :param nickname: 用户昵称
  :param score: 关联得分
  :param category: 文件类别，未提供时自动根据文件名推断
  :param bos_uploaded: 是否已上传至 BOS
  :param extra_metadata: 额外信息
  """
  normalized = _normalize_file_path(file_path)
  if normalized is None:
    logger.info("record_image_creation: 无法计算文件名，路径=%s", file_path)
    return

  abs_path = os.path.join(_IMAGES_DIR_ABS, normalized)
  if not os.path.isfile(abs_path):
    logger.info("record_image_creation: 文件不存在，跳过记录 file=%s", abs_path)
    return

  try:
    stat = os.stat(abs_path)
    category_name = category or infer_category_from_filename(normalized)
    is_cropped_face = "_face_" in normalized and normalized.count("_") >= 2
    last_modified = datetime.fromtimestamp(stat.st_mtime)

    nickname_value = nickname.strip() if isinstance(nickname,
                                                    str) and nickname.strip() else None

    await upsert_image_record(
        file_path=normalized,
        category=category_name,
        nickname=nickname_value,
        score=score,
        is_cropped_face=is_cropped_face,
        size_bytes=stat.st_size,
        last_modified=last_modified,
        bos_uploaded=bos_uploaded,
        hostname=HOSTNAME,
        extra_metadata=extra_metadata,
    )
  except Exception as exc:
    logger.warning(f"写入图片记录失败: {exc}")