"""
统一存储服务 (Professional Storage Service)
支持 Local (TOML), Redis, MySQL, PostgreSQL

特性:
- 全异步 I/O (Async I/O)
- 连接池管理 (Connection Pooling)
- 分布式/本地锁 (Distributed/Local Locking)
- 内存优化 (序列化性能优化)
"""

import abc
import os
import asyncio
import hashlib
import time
import tomllib
import re
from typing import Any, Dict, Optional
from pathlib import Path
from enum import Enum

try:
    import fcntl
except ImportError:  # pragma: no cover - non-posix platforms
    fcntl = None
from contextlib import asynccontextmanager

import orjson
import aiofiles
from app.core.logger import logger

# 数据目录（支持通过环境变量覆盖）
DEFAULT_DATA_DIR = Path(__file__).parent.parent.parent / "data"
DATA_DIR = Path(os.getenv("DATA_DIR", str(DEFAULT_DATA_DIR))).expanduser()

# 配置文件路径
CONFIG_FILE = DATA_DIR / "config.toml"
TOKEN_FILE = DATA_DIR / "token.json"
LOCK_DIR = DATA_DIR / ".locks"


# JSON 序列化优化助手函数
def json_dumps(obj: Any) -> str:
    return orjson.dumps(obj).decode("utf-8")


def json_loads(obj: str | bytes) -> Any:
    return orjson.loads(obj)


_TOML_BARE_KEY_RE = re.compile(r"^[A-Za-z0-9_-]+$")


def _toml_escape_string(value: str) -> str:
    """转义 TOML 字符串。"""
    return (
        value.replace("\\", "\\\\")
        .replace('"', '\\"')
        .replace("\n", "\\n")
        .replace("\r", "\\r")
        .replace("\t", "\\t")
    )


def _toml_format_key(key: Any) -> str:
    """格式化 TOML 键名，必要时自动加引号。"""
    text = str(key)
    if _TOML_BARE_KEY_RE.fullmatch(text):
        return text
    return f'"{_toml_escape_string(text)}"'


def _toml_format_value(value: Any) -> str:
    """将 Python 值序列化为 TOML 字面量。"""
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, (int, float)) and not isinstance(value, bool):
        return str(value)
    if isinstance(value, str):
        return f'"{_toml_escape_string(value)}"'
    if isinstance(value, list):
        return "[" + ", ".join(_toml_format_value(item) for item in value) + "]"
    if isinstance(value, dict):
        parts = [
            f"{_toml_format_key(key)} = {_toml_format_value(item)}"
            for key, item in value.items()
        ]
        return "{ " + ", ".join(parts) + " }"
    if value is None:
        return '""'
    return f'"{_toml_escape_string(str(value))}"'


class StorageError(Exception):
    """存储服务基础异常"""

    pass


class BaseStorage(abc.ABC):
    """存储基类"""

    @abc.abstractmethod
    async def load_config(self) -> Dict[str, Any]:
        """加载配置"""
        pass

    @abc.abstractmethod
    async def save_config(self, data: Dict[str, Any]):
        """保存配置"""
        pass

    @abc.abstractmethod
    async def load_tokens(self) -> Dict[str, Any]:
        """加载所有 Token"""
        pass

    @abc.abstractmethod
    async def save_tokens(self, data: Dict[str, Any]):
        """保存所有 Token"""
        pass

    @abc.abstractmethod
    async def close(self):
        """关闭资源"""
        pass

    @asynccontextmanager
    async def acquire_lock(self, name: str, timeout: int = 10):
        """
        获取锁 (互斥访问)
        用于读写操作的临界区保护

        Args:
            name: 锁名称
            timeout: 超时时间 (秒)
        """
        # 默认空实现，用于 fallback
        yield

    async def verify_connection(self) -> bool:
        """健康检查"""
        return True


class LocalStorage(BaseStorage):
    """
    本地文件存储
    - 使用 aiofiles 进行异步 I/O
    - 使用 asyncio.Lock 进行进程内并发控制
    - 如果需要多进程安全，需要系统级文件锁 (fcntl)
    """

    def __init__(self):
        self._lock = asyncio.Lock()

    @asynccontextmanager
    async def acquire_lock(self, name: str, timeout: int = 10):
        if fcntl is None:
            try:
                async with asyncio.timeout(timeout):
                    async with self._lock:
                        yield
            except asyncio.TimeoutError:
                logger.warning(f"LocalStorage: 获取锁 '{name}' 超时 ({timeout}s)")
                raise StorageError(f"无法获取锁 '{name}'")
            return

        lock_path = LOCK_DIR / f"{name}.lock"
        lock_path.parent.mkdir(parents=True, exist_ok=True)
        fd = None
        locked = False
        start = time.monotonic()

        async with self._lock:
            try:
                fd = open(lock_path, "a+")
                while True:
                    try:
                        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
                        locked = True
                        break
                    except BlockingIOError:
                        if time.monotonic() - start >= timeout:
                            raise StorageError(f"无法获取锁 '{name}'")
                        await asyncio.sleep(0.05)
                yield
            except StorageError:
                logger.warning(f"LocalStorage: 获取锁 '{name}' 超时 ({timeout}s)")
                raise
            finally:
                if fd:
                    if locked:
                        try:
                            fcntl.flock(fd, fcntl.LOCK_UN)
                        except Exception:
                            pass
                    try:
                        fd.close()
                    except Exception:
                        pass

    async def load_config(self) -> Dict[str, Any]:
        if not CONFIG_FILE.exists():
            return {}
        try:
            async with aiofiles.open(CONFIG_FILE, "rb") as f:
                content = await f.read()
                return tomllib.loads(content.decode("utf-8"))
        except Exception as e:
            logger.error(f"LocalStorage: 加载配置失败: {e}")
            return {}

    async def save_config(self, data: Dict[str, Any]):
        try:
            lines = []
            for section, items in data.items():
                if not isinstance(items, dict):
                    continue
                lines.append(f"[{_toml_format_key(section)}]")
                for key, val in items.items():
                    lines.append(
                        f"{_toml_format_key(key)} = {_toml_format_value(val)}"
                    )
                lines.append("")

            content = "\n".join(lines)

            CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
            async with aiofiles.open(CONFIG_FILE, "w", encoding="utf-8") as f:
                await f.write(content)
        except Exception as e:
            logger.error(f"LocalStorage: 保存配置失败: {e}")
            raise StorageError(f"保存配置失败: {e}")

    async def load_tokens(self) -> Dict[str, Any]:
        if not TOKEN_FILE.exists():
            return {}
        try:
            async with aiofiles.open(TOKEN_FILE, "rb") as f:
                content = await f.read()
                return json_loads(content)
        except Exception as e:
            logger.error(f"LocalStorage: 加载 Token 失败: {e}")
            return {}

    async def save_tokens(self, data: Dict[str, Any]):
        try:
            TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True)
            temp_path = TOKEN_FILE.with_suffix(".tmp")

            # 原子写操作: 写入临时文件 -> 重命名
            async with aiofiles.open(temp_path, "wb") as f:
                await f.write(orjson.dumps(data, option=orjson.OPT_INDENT_2))

            # 使用 os.replace 保证原子性
            os.replace(temp_path, TOKEN_FILE)

        except Exception as e:
            logger.error(f"LocalStorage: 保存 Token 失败: {e}")
            raise StorageError(f"保存 Token 失败: {e}")

    async def close(self):
        pass


class RedisStorage(BaseStorage):
    """
    Redis 存储
    - 使用 redis-py 异步客户端 (自带连接池)
    - 支持分布式锁 (redis.lock)
    - 扁平化数据结构优化性能
    """

    def __init__(self, url: str):
        try:
            from redis import asyncio as aioredis
        except ImportError:
            raise ImportError("需要安装 redis 包: pip install redis")

        # 显式配置连接池
        # 使用 decode_responses=True 简化字符串处理，但在处理复杂对象时使用 orjson
        self.redis = aioredis.from_url(
            url, decode_responses=True, health_check_interval=30
        )
        self.config_key = "grok2api:config"  # Hash: section.key -> value_json
        self.key_pools = "grok2api:pools"  # Set: pool_names
        self.prefix_pool_set = "grok2api:pool:"  # Set: pool -> token_ids
        self.prefix_token_hash = "grok2api:token:"  # Hash: token_id -> token_data
        self.lock_prefix = "grok2api:lock:"

    @asynccontextmanager
    async def acquire_lock(self, name: str, timeout: int = 10):
        # 使用 Redis 分布式锁
        lock_key = f"{self.lock_prefix}{name}"
        lock = self.redis.lock(lock_key, timeout=timeout, blocking_timeout=5)
        acquired = False
        try:
            acquired = await lock.acquire()
            if not acquired:
                raise StorageError(f"RedisStorage: 无法获取锁 '{name}'")
            yield
        finally:
            if acquired:
                try:
                    await lock.release()
                except Exception:
                    # 锁可能已过期或被意外释放，忽略异常
                    pass

    async def verify_connection(self) -> bool:
        try:
            return await self.redis.ping()
        except Exception:
            return False

    async def load_config(self) -> Dict[str, Any]:
        """从 Redis Hash 加载配置"""
        try:
            raw_data = await self.redis.hgetall(self.config_key)
            if not raw_data:
                return None

            config = {}
            for composite_key, val_str in raw_data.items():
                if "." not in composite_key:
                    continue
                section, key = composite_key.split(".", 1)

                if section not in config:
                    config[section] = {}

                try:
                    val = json_loads(val_str)
                except Exception:
                    val = val_str
                config[section][key] = val
            return config
        except Exception as e:
            logger.error(f"RedisStorage: 加载配置失败: {e}")
            return None

    async def save_config(self, data: Dict[str, Any]):
        """保存配置到 Redis Hash"""
        if not data:
            return
        try:
            mapping = {}
            for section, items in data.items():
                if not isinstance(items, dict):
                    continue
                for key, val in items.items():
                    composite_key = f"{section}.{key}"
                    mapping[composite_key] = json_dumps(val)

            if mapping:
                await self.redis.hset(self.config_key, mapping=mapping)
        except Exception as e:
            logger.error(f"RedisStorage: 保存配置失败: {e}")
            raise

    async def load_tokens(self) -> Dict[str, Any]:
        """加载所有 Token"""
        try:
            pool_names = await self.redis.smembers(self.key_pools)
            if not pool_names:
                return None

            pools = {}
            async with self.redis.pipeline() as pipe:
                for pool_name in pool_names:
                    # 获取该池下所有 Token ID
                    pipe.smembers(f"{self.prefix_pool_set}{pool_name}")
                pool_tokens_res = await pipe.execute()

            # 收集所有 Token ID 以便批量查询
            all_token_ids = []
            pool_map = {}  # pool_name -> list[token_id]

            for i, pool_name in enumerate(pool_names):
                tids = list(pool_tokens_res[i])
                pool_map[pool_name] = tids
                all_token_ids.extend(tids)

            if not all_token_ids:
                return {name: [] for name in pool_names}

            # 批量获取 Token 详情 (Hash)
            async with self.redis.pipeline() as pipe:
                for tid in all_token_ids:
                    pipe.hgetall(f"{self.prefix_token_hash}{tid}")
                token_data_list = await pipe.execute()

            # 重组数据结构
            token_lookup = {}
            for i, tid in enumerate(all_token_ids):
                t_data = token_data_list[i]
                if not t_data:
                    continue

                # 恢复 tags (JSON -> List)
                if "tags" in t_data:
                    try:
                        t_data["tags"] = json_loads(t_data["tags"])
                    except Exception:
                        t_data["tags"] = []

                # 类型转换 (Redis 返回全 string)
                for int_field in [
                    "quota",
                    "created_at",
                    "use_count",
                    "fail_count",
                    "last_used_at",
                    "last_fail_at",
                    "last_sync_at",
                ]:
                    if t_data.get(int_field) and t_data[int_field] != "None":
                        try:
                            t_data[int_field] = int(t_data[int_field])
                        except Exception:
                            pass

                token_lookup[tid] = t_data

            # 按 Pool 分组返回
            for pool_name in pool_names:
                pools[pool_name] = []
                for tid in pool_map[pool_name]:
                    if tid in token_lookup:
                        pools[pool_name].append(token_lookup[tid])

            return pools

        except Exception as e:
            logger.error(f"RedisStorage: 加载 Token 失败: {e}")
            return None

    async def save_tokens(self, data: Dict[str, Any]):
        """保存所有 Token"""
        if data is None:
            return
        try:
            new_pools = set(data.keys()) if isinstance(data, dict) else set()
            pool_tokens_map = {}
            new_token_ids = set()

            for pool_name, tokens in (data or {}).items():
                tids_in_pool = []
                for t in tokens:
                    token_str = t.get("token")
                    if not token_str:
                        continue
                    tids_in_pool.append(token_str)
                    new_token_ids.add(token_str)
                pool_tokens_map[pool_name] = tids_in_pool

            existing_pools = await self.redis.smembers(self.key_pools)
            existing_pools = set(existing_pools) if existing_pools else set()

            existing_token_ids = set()
            if existing_pools:
                async with self.redis.pipeline() as pipe:
                    for pool_name in existing_pools:
                        pipe.smembers(f"{self.prefix_pool_set}{pool_name}")
                    pool_tokens_res = await pipe.execute()
                for tokens in pool_tokens_res:
                    existing_token_ids.update(list(tokens or []))

            tokens_to_delete = existing_token_ids - new_token_ids
            all_pools = existing_pools.union(new_pools)

            async with self.redis.pipeline() as pipe:
                # Reset pool index
                pipe.delete(self.key_pools)
                if new_pools:
                    pipe.sadd(self.key_pools, *new_pools)

                # Reset pool sets
                for pool_name in all_pools:
                    pipe.delete(f"{self.prefix_pool_set}{pool_name}")
                for pool_name, tids_in_pool in pool_tokens_map.items():
                    if tids_in_pool:
                        pipe.sadd(f"{self.prefix_pool_set}{pool_name}", *tids_in_pool)

                # Remove deleted token hashes
                for token_str in tokens_to_delete:
                    pipe.delete(f"{self.prefix_token_hash}{token_str}")

                # Upsert token hashes
                for pool_name, tokens in (data or {}).items():
                    for t in tokens:
                        token_str = t.get("token")
                        if not token_str:
                            continue
                        t_flat = t.copy()
                        if "tags" in t_flat:
                            t_flat["tags"] = json_dumps(t_flat["tags"])
                        status = t_flat.get("status")
                        if isinstance(status, str) and status.startswith(
                            "TokenStatus."
                        ):
                            t_flat["status"] = status.split(".", 1)[1].lower()
                        elif isinstance(status, Enum):
                            t_flat["status"] = status.value
                        t_flat = {k: str(v) for k, v in t_flat.items() if v is not None}
                        pipe.hset(
                            f"{self.prefix_token_hash}{token_str}", mapping=t_flat
                        )

                await pipe.execute()

        except Exception as e:
            logger.error(f"RedisStorage: 保存 Token 失败: {e}")
            raise

    async def close(self):
        try:
            await self.redis.close()
        except (RuntimeError, asyncio.CancelledError, Exception):
            # 忽略关闭时的 Event loop is closed 错误
            pass


class SQLStorage(BaseStorage):
    """
    SQL 数据库存储 (MySQL/PgSQL)
    - 使用 SQLAlchemy 异步引擎
    - 自动 Schema 初始化
    - 内置连接池 (QueuePool)
    """

    def __init__(self, url: str):
        try:
            from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
        except ImportError:
            raise ImportError(
                "需要安装 sqlalchemy 和 async 驱动: pip install sqlalchemy[asyncio]"
            )

        self.dialect = url.split(":", 1)[0].split("+", 1)[0].lower()

        # 配置 robust 的连接池
        self.engine = create_async_engine(
            url,
            echo=False,
            pool_size=20,
            max_overflow=10,
            pool_recycle=3600,
            pool_pre_ping=True,
        )
        self.async_session = async_sessionmaker(self.engine, expire_on_commit=False)
        self._initialized = False

    async def _ensure_schema(self):
        """确保数据库表存在"""
        if self._initialized:
            return
        try:
            async with self.engine.begin() as conn:
                from sqlalchemy import text

                # Tokens 表 (通用 SQL)
                await conn.execute(
                    text("""
                    CREATE TABLE IF NOT EXISTS tokens (
                        token VARCHAR(512) PRIMARY KEY,
                        pool_name VARCHAR(64) NOT NULL,
                        data TEXT,
                        updated_at BIGINT
                    )
                """)
                )

                # 配置表
                await conn.execute(
                    text("""
                    CREATE TABLE IF NOT EXISTS app_config (
                        section VARCHAR(64) NOT NULL,
                        key_name VARCHAR(64) NOT NULL,
                        value TEXT,
                        PRIMARY KEY (section, key_name)
                    )
                """)
                )

                # 索引
                try:
                    await conn.execute(
                        text("CREATE INDEX idx_tokens_pool ON tokens (pool_name)")
                    )
                except Exception:
                    pass

                # 尝试兼容旧表结构
                try:
                    if self.dialect in ("mysql", "mariadb"):
                        await conn.execute(
                            text("ALTER TABLE tokens MODIFY token VARCHAR(512)")
                        )
                        await conn.execute(text("ALTER TABLE tokens MODIFY data TEXT"))
                    elif self.dialect in ("postgres", "postgresql", "pgsql"):
                        await conn.execute(
                            text(
                                "ALTER TABLE tokens ALTER COLUMN token TYPE VARCHAR(512)"
                            )
                        )
                        await conn.execute(
                            text("ALTER TABLE tokens ALTER COLUMN data TYPE TEXT")
                        )
                except Exception:
                    pass

            self._initialized = True
        except Exception as e:
            logger.error(f"SQLStorage: Schema 初始化失败: {e}")
            raise

    @asynccontextmanager
    async def acquire_lock(self, name: str, timeout: int = 10):
        # SQL 分布式锁: MySQL GET_LOCK / PG advisory_lock
        from sqlalchemy import text

        lock_name = f"g2a:{hashlib.sha1(name.encode('utf-8')).hexdigest()[:24]}"
        if self.dialect in ("mysql", "mariadb"):
            async with self.async_session() as session:
                res = await session.execute(
                    text("SELECT GET_LOCK(:name, :timeout)"),
                    {"name": lock_name, "timeout": timeout},
                )
                got = res.scalar()
                if got != 1:
                    raise StorageError(f"SQLStorage: 无法获取锁 '{name}'")
                try:
                    yield
                finally:
                    try:
                        await session.execute(
                            text("SELECT RELEASE_LOCK(:name)"), {"name": lock_name}
                        )
                        await session.commit()
                    except Exception:
                        pass
        elif self.dialect in ("postgres", "postgresql", "pgsql"):
            lock_key = int.from_bytes(
                hashlib.sha256(name.encode("utf-8")).digest()[:8], "big", signed=False
            )
            async with self.async_session() as session:
                start = time.monotonic()
                while True:
                    res = await session.execute(
                        text("SELECT pg_try_advisory_lock(:key)"), {"key": lock_key}
                    )
                    if res.scalar():
                        break
                    if time.monotonic() - start >= timeout:
                        raise StorageError(f"SQLStorage: 无法获取锁 '{name}'")
                    await asyncio.sleep(0.1)
                try:
                    yield
                finally:
                    try:
                        await session.execute(
                            text("SELECT pg_advisory_unlock(:key)"), {"key": lock_key}
                        )
                        await session.commit()
                    except Exception:
                        pass
        else:
            yield

    async def load_config(self) -> Dict[str, Any]:
        await self._ensure_schema()
        from sqlalchemy import text

        try:
            async with self.async_session() as session:
                res = await session.execute(
                    text("SELECT section, key_name, value FROM app_config")
                )
                rows = res.fetchall()
                if not rows:
                    return None

                config = {}
                for section, key, val_str in rows:
                    if section not in config:
                        config[section] = {}
                    try:
                        val = json_loads(val_str)
                    except Exception:
                        val = val_str
                    config[section][key] = val
                return config
        except Exception as e:
            logger.error(f"SQLStorage: 加载配置失败: {e}")
            return None

    async def save_config(self, data: Dict[str, Any]):
        await self._ensure_schema()
        from sqlalchemy import text

        try:
            async with self.async_session() as session:
                for section, items in data.items():
                    if not isinstance(items, dict):
                        continue
                    for key, val in items.items():
                        val_str = json_dumps(val)

                        # Upsert 逻辑 (简单实现: Delete + Insert)
                        await session.execute(
                            text(
                                "DELETE FROM app_config WHERE section=:s AND key_name=:k"
                            ),
                            {"s": section, "k": key},
                        )
                        await session.execute(
                            text(
                                "INSERT INTO app_config (section, key_name, value) VALUES (:s, :k, :v)"
                            ),
                            {"s": section, "k": key, "v": val_str},
                        )
                await session.commit()
        except Exception as e:
            logger.error(f"SQLStorage: 保存配置失败: {e}")
            raise

    async def load_tokens(self) -> Dict[str, Any]:
        await self._ensure_schema()
        from sqlalchemy import text

        try:
            async with self.async_session() as session:
                res = await session.execute(text("SELECT pool_name, data FROM tokens"))
                rows = res.fetchall()
                if not rows:
                    return None

                pools = {}
                for pool_name, data_json in rows:
                    if pool_name not in pools:
                        pools[pool_name] = []

                    try:
                        if isinstance(data_json, str):
                            t_data = json_loads(data_json)
                        else:
                            t_data = data_json
                        pools[pool_name].append(t_data)
                    except Exception:
                        pass
                return pools
        except Exception as e:
            logger.error(f"SQLStorage: 加载 Token 失败: {e}")
            return None

    async def save_tokens(self, data: Dict[str, Any]):
        await self._ensure_schema()
        from sqlalchemy import text

        try:
            async with self.async_session() as session:
                await session.execute(text("DELETE FROM tokens"))

                params = []
                for pool_name, tokens in data.items():
                    for t in tokens:
                        params.append(
                            {
                                "token": t.get("token"),
                                "pool_name": pool_name,
                                "data": json_dumps(t),
                                "updated_at": 0,
                            }
                        )

                if params:
                    # 批量插入
                    await session.execute(
                        text(
                            "INSERT INTO tokens (token, pool_name, data, updated_at) VALUES (:token, :pool_name, :data, :updated_at)"
                        ),
                        params,
                    )
                await session.commit()
        except Exception as e:
            logger.error(f"SQLStorage: 保存 Token 失败: {e}")
            raise

    async def close(self):
        await self.engine.dispose()


class StorageFactory:
    """存储后端工厂"""

    _instance: Optional[BaseStorage] = None

    @staticmethod
    def _normalize_sql_url(storage_type: str, url: str) -> str:
        if not url or "://" not in url:
            return url
        if storage_type == "mysql":
            if url.startswith("mysql://"):
                return f"mysql+aiomysql://{url[len('mysql://') :]}"
            if url.startswith("mariadb://"):
                return f"mariadb+aiomysql://{url[len('mariadb://') :]}"
        if storage_type == "pgsql":
            if url.startswith("postgres://"):
                return f"postgresql+asyncpg://{url[len('postgres://') :]}"
            if url.startswith("postgresql://"):
                return f"postgresql+asyncpg://{url[len('postgresql://') :]}"
            if url.startswith("pgsql://"):
                return f"postgresql+asyncpg://{url[len('pgsql://') :]}"
        return url

    @classmethod
    def get_storage(cls) -> BaseStorage:
        """获取全局存储实例 (单例)"""
        if cls._instance:
            return cls._instance

        storage_type = os.getenv("SERVER_STORAGE_TYPE", "local").lower()
        storage_url = os.getenv("SERVER_STORAGE_URL", "")

        logger.info(f"StorageFactory: 初始化存储后端: {storage_type}")

        if storage_type == "redis":
            if not storage_url:
                raise ValueError("Redis 存储需要设置 SERVER_STORAGE_URL")
            cls._instance = RedisStorage(storage_url)

        elif storage_type in ("mysql", "pgsql"):
            if not storage_url:
                raise ValueError("SQL 存储需要设置 SERVER_STORAGE_URL")
            storage_url = cls._normalize_sql_url(storage_type, storage_url)
            cls._instance = SQLStorage(storage_url)

        else:
            cls._instance = LocalStorage()

        return cls._instance


def get_storage() -> BaseStorage:
    return StorageFactory.get_storage()