"""
舷号-型号映射数据库模块
用于管理舰船型号、舷号、别名、舰名的映射关系
数据存储在 PostgreSQL 的 ship_model_mapping 表中
"""
import json
from typing import Dict, List, Any, Optional

from config import POSTGRES_CONNECTION_STRING


async def init_ship_model_mapping_table():
    """
    初始化 ship_model_mapping 表并插入初始数据
    """
    try:
        from psycopg_pool import AsyncConnectionPool
    except ImportError:
        print("[ship_model_db] psycopg_pool 未安装，无法初始化表")
        return

    async with AsyncConnectionPool(
        POSTGRES_CONNECTION_STRING,
        kwargs={"autocommit": True}
    ) as pool:
        async with pool.connection() as conn:
            async with conn.cursor() as cur:
                await cur.execute("""
                    CREATE TABLE IF NOT EXISTS ship_model_mapping (
                        id SERIAL PRIMARY KEY,
                        model_name TEXT NOT NULL UNIQUE,
                        numbers JSONB NOT NULL DEFAULT '[]',
                        aliases JSONB NOT NULL DEFAULT '[]',
                        names JSONB NOT NULL DEFAULT '[]'
                    )
                """)
                await cur.execute("""
                    CREATE INDEX IF NOT EXISTS ship_model_mapping_model_name_idx ON ship_model_mapping(model_name)
                """)
                print("[ship_model_db] ship_model_mapping 表初始化完成")

                # 检查是否已有数据，没有则插入初始数据
                await cur.execute("SELECT COUNT(*) FROM ship_model_mapping")
                count = (await cur.fetchone())[0]
                if count == 0:
                    initial_data = [
                        ("055型驱逐舰", ["102", "103", "104", "105", "106", "107", "108"],
                         ["055", "055型", "055驱逐舰", "万吨大驱"],
                         ["拉萨舰", "鞍山舰", "无锡舰", "大连舰", "延安舰", "遵义舰", "咸阳舰"]),
                        ("052D型驱逐舰", ["172", "173", "174", "175", "163"],
                         ["052D", "052D型", "052D驱逐舰"],
                         ["昆明舰", "长沙舰", "合肥舰", "银川舰", "南昌舰"]),
                        ("054A型护卫舰", ["529", "530", "547", "568", "570"],
                         ["054A", "054A型", "054A护卫舰"],
                         ["舟山舰", "徐州舰", "临沂舰", "衡阳舰", "黄山舰"]),
                    ]
                    for model_name, numbers, aliases, names in initial_data:
                        await cur.execute(
                            """
                            INSERT INTO ship_model_mapping (model_name, numbers, aliases, names)
                            VALUES (%s, %s, %s, %s)
                            """,
                            (model_name,
                             json.dumps(numbers, ensure_ascii=False),
                             json.dumps(aliases, ensure_ascii=False),
                             json.dumps(names, ensure_ascii=False))
                        )
                    print(f"[ship_model_db] 已插入 {len(initial_data)} 条初始数据")


async def load_ship_model_mapping() -> Dict[str, Dict[str, List[str]]]:
    """
    从数据库加载全部舷号-型号映射数据

    Returns:
        与原 config.SHIP_MODEL_MAPPING 格式一致的字典
        {
            "055型驱逐舰": {
                "numbers": ["102", ...],
                "aliases": ["055", ...],
                "names": ["拉萨舰", ...]
            },
            ...
        }
    """
    try:
        from psycopg_pool import AsyncConnectionPool
    except ImportError:
        print("[ship_model_db] psycopg_pool 未安装，返回空映射")
        return {}

    try:
        async with AsyncConnectionPool(
            POSTGRES_CONNECTION_STRING,
            kwargs={"autocommit": True}
        ) as pool:
            async with pool.connection() as conn:
                async with conn.cursor() as cur:
                    await cur.execute(
                        "SELECT model_name, numbers, aliases, names FROM ship_model_mapping"
                    )
                    rows = await cur.fetchall()
                    result = {}
                    for row in rows:
                        model_name, numbers, aliases, names = row
                        # JSONB 返回的可能是 list 或 str，统一处理
                        if isinstance(numbers, str):
                            numbers = json.loads(numbers)
                        if isinstance(aliases, str):
                            aliases = json.loads(aliases)
                        if isinstance(names, str):
                            names = json.loads(names)
                        result[model_name] = {
                            "numbers": numbers,
                            "aliases": aliases,
                            "names": names,
                        }
                    return result
    except Exception as e:
        print(f"[ship_model_db] 加载映射数据失败: {str(e)}")
        return {}


async def add_ship_model(
    model_name: str,
    numbers: List[str],
    aliases: List[str],
    names: List[str]
) -> bool:
    """
    添加舰船型号映射

    Args:
        model_name: 型号名称
        numbers: 舷号列表
        aliases: 别名列表
        names: 舰名列表

    Returns:
        是否添加成功
    """
    try:
        from psycopg_pool import AsyncConnectionPool
    except ImportError:
        print("[ship_model_db] psycopg_pool 未安装，无法添加")
        return False

    try:
        async with AsyncConnectionPool(
            POSTGRES_CONNECTION_STRING,
            kwargs={"autocommit": True}
        ) as pool:
            async with pool.connection() as conn:
                async with conn.cursor() as cur:
                    await cur.execute(
                        """
                        INSERT INTO ship_model_mapping (model_name, numbers, aliases, names)
                        VALUES (%s, %s, %s, %s)
                        ON CONFLICT (model_name) DO UPDATE SET
                            numbers = EXCLUDED.numbers,
                            aliases = EXCLUDED.aliases,
                            names = EXCLUDED.names
                        """,
                        (model_name,
                         json.dumps(numbers, ensure_ascii=False),
                         json.dumps(aliases, ensure_ascii=False),
                         json.dumps(names, ensure_ascii=False))
                    )
        print(f"[ship_model_db] 舰船型号 '{model_name}' 已保存")
        return True
    except Exception as e:
        print(f"[ship_model_db] 添加舰船型号失败: {str(e)}")
        return False


async def delete_ship_model(model_name: str) -> bool:
    """
    删除舰船型号映射

    Args:
        model_name: 型号名称

    Returns:
        是否删除成功
    """
    try:
        from psycopg_pool import AsyncConnectionPool
    except ImportError:
        print("[ship_model_db] psycopg_pool 未安装，无法删除")
        return False

    try:
        async with AsyncConnectionPool(
            POSTGRES_CONNECTION_STRING,
            kwargs={"autocommit": True}
        ) as pool:
            async with pool.connection() as conn:
                async with conn.cursor() as cur:
                    await cur.execute(
                        "DELETE FROM ship_model_mapping WHERE model_name = %s",
                        (model_name,)
                    )
        print(f"[ship_model_db] 舰船型号 '{model_name}' 已删除")
        return True
    except Exception as e:
        print(f"[ship_model_db] 删除舰船型号失败: {str(e)}")
        return False