#!/usr/bin/env python3
"""
补全高价值双星/多星系统数据
包含8-10个科学价值最高的多恒星系统
"""
import asyncio
import asyncpg
import json
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 数据库连接配置
DB_CONFIG = {
    "host": "localhost",
    "port": 5432,
    "user": "postgres",
    "password": "postgres",
    "database": "cosmo_db"
}

# 高价值多恒星系统数据（基于天文学资料）
MULTI_STAR_SYSTEMS = {
    # 1. Alpha Centauri (比邻星系统) - 已完成，保留用于验证
    479: {
        "stars": [
            {
                "id": "star-479-primary",
                "name": "Alpha Centauri A",
                "name_zh": "南门二A",
                "description": "该恒星系主序星，光谱类型: G2V, 表面温度: 5790K",
                "extra_data": {
                    "spectral_type": "G2V",
                    "mass_solar": 1.1,
                    "radius_solar": 1.22,
                    "temperature_k": 5790
                }
            },
            {
                "id": "star-479-secondary",
                "name": "Alpha Centauri B",
                "name_zh": "南门二B",
                "description": "与南门二A相互绕转的明亮双星，是该系统的主体。",
                "extra_data": {
                    "spectral_type": "K1V",
                    "mass_solar": 0.93,
                    "radius_solar": 0.86,
                    "temperature_k": 5260
                }
            },
            {
                "id": "star-479-tertiary",
                "name": "Proxima Centauri",
                "name_zh": "比邻星",
                "description": "一颗质量小、光度弱的红矮星，距离南门二A/B约0.2光年，围绕它们公转。",
                "extra_data": {
                    "spectral_type": "M5.5V",
                    "mass_solar": 0.12,
                    "radius_solar": 0.14,
                    "temperature_k": 2900
                }
            }
        ]
    },

    # 2. 55 Cancri (巨蟹座55) - 双星系统
    11: {
        "stars": [
            {
                "id": "star-11-primary",
                "name": "55 Cancri A",
                "name_zh": "巨蟹座55A",
                "description": "类太阳黄矮星，拥有5颗已确认行星，包括著名的超级地球55 Cnc e。",
                "extra_data": {
                    "spectral_type": "G8V",
                    "mass_solar": 0.95,
                    "radius_solar": 0.94,
                    "temperature_k": 5196
                }
            },
            {
                "id": "star-11-secondary",
                "name": "55 Cancri B",
                "name_zh": "巨蟹座55B",
                "description": "红矮星伴星，距离A星约1065 AU，轨道周期约1000年。",
                "extra_data": {
                    "spectral_type": "M4V",
                    "mass_solar": 0.13,
                    "radius_solar": 0.30,
                    "temperature_k": 3200,
                    "separation_au": 1065,
                    "orbital_period_years": 1000
                }
            }
        ]
    },

    # 3. 16 Cygni (天鹅座16) - 双星系统
    5: {
        "stars": [
            {
                "id": "star-5-primary",
                "name": "16 Cygni A",
                "name_zh": "天鹅座16A",
                "description": "类太阳黄矮星，该双星系统的主星。",
                "extra_data": {
                    "spectral_type": "G1.5V",
                    "mass_solar": 1.11,
                    "radius_solar": 1.24,
                    "temperature_k": 5825
                }
            },
            {
                "id": "star-5-secondary",
                "name": "16 Cygni B",
                "name_zh": "天鹅座16B",
                "description": "类太阳黄矮星，拥有一颗高偏心率轨道的行星16 Cyg B b，展示了双星引力对行星轨道的影响。",
                "extra_data": {
                    "spectral_type": "G2.5V",
                    "mass_solar": 1.07,
                    "radius_solar": 1.14,
                    "temperature_k": 5750,
                    "separation_au": 850,
                    "orbital_period_years": 18200
                }
            }
        ]
    },

    # 4. Epsilon Indi (天园增四) - 三体系统 (1恒星 + 2棕矮星)
    40: {
        "stars": [
            {
                "id": "star-40-primary",
                "name": "Epsilon Indi A",
                "name_zh": "天园增四A",
                "description": "橙矮星，第五近的恒星系统，伴有两颗棕矮星（Ba和Bb）。",
                "extra_data": {
                    "spectral_type": "K5V",
                    "mass_solar": 0.76,
                    "radius_solar": 0.73,
                    "temperature_k": 4630
                }
            },
            {
                "id": "star-40-secondary",
                "name": "Epsilon Indi Ba",
                "name_zh": "天园增四Ba",
                "description": "T1V型棕矮星，距离A星约1460 AU，与Bb组成棕矮星双星系统。",
                "extra_data": {
                    "spectral_type": "T1V",
                    "mass_jupiter": 47,
                    "radius_jupiter": 0.91,
                    "temperature_k": 1300,
                    "separation_from_A_au": 1460,
                    "is_brown_dwarf": True
                }
            },
            {
                "id": "star-40-tertiary",
                "name": "Epsilon Indi Bb",
                "name_zh": "天园增四Bb",
                "description": "T6V型棕矮星，与Ba互绕，周期约15年，是最近的棕矮星双星系统。",
                "extra_data": {
                    "spectral_type": "T6V",
                    "mass_jupiter": 28,
                    "radius_jupiter": 0.80,
                    "temperature_k": 880,
                    "orbital_period_years": 15,
                    "is_brown_dwarf": True
                }
            }
        ]
    },

    # 5. Gamma Cephei (仙王座γ) - 双星系统
    49: {
        "stars": [
            {
                "id": "star-49-primary",
                "name": "Gamma Cephei A",
                "name_zh": "仙王座γA",
                "description": "亚巨星，最早被怀疑有行星的恒星之一（1988年），拥有一颗类木行星。",
                "extra_data": {
                    "spectral_type": "K1IV",
                    "mass_solar": 1.59,
                    "radius_solar": 4.9,
                    "temperature_k": 4800
                }
            },
            {
                "id": "star-49-secondary",
                "name": "Gamma Cephei B",
                "name_zh": "仙王座γB",
                "description": "红矮星伴星，距离A星约20 AU，轨道周期约66年，形成紧密双星系统。",
                "extra_data": {
                    "spectral_type": "M4V",
                    "mass_solar": 0.4,
                    "radius_solar": 0.40,
                    "temperature_k": 3200,
                    "separation_au": 20,
                    "orbital_period_years": 66
                }
            }
        ]
    },

    # 6. Upsilon Andromedae (仙女座υ) - 双星系统
    572: {
        "stars": [
            {
                "id": "star-572-primary",
                "name": "Upsilon Andromedae A",
                "name_zh": "仙女座υA",
                "description": "黄白主序星，第一个被发现有多颗行星的主序星（1999年），拥有4颗已确认行星。",
                "extra_data": {
                    "spectral_type": "F8V",
                    "mass_solar": 1.27,
                    "radius_solar": 1.63,
                    "temperature_k": 6212
                }
            },
            {
                "id": "star-572-secondary",
                "name": "Upsilon Andromedae B",
                "name_zh": "仙女座υB",
                "description": "红矮星伴星，距离A星约750 AU。",
                "extra_data": {
                    "spectral_type": "M4.5V",
                    "mass_solar": 0.25,
                    "radius_solar": 0.28,
                    "temperature_k": 3100,
                    "separation_au": 750
                }
            }
        ]
    },

    # 7. HD 41004 - 双星系统（两个独立的system_id需要合并）
    347: {
        "stars": [
            {
                "id": "star-347-primary",
                "name": "HD 41004 A",
                "name_zh": "HD 41004 A",
                "description": "橙矮星，拥有一颗类木行星HD 41004 A b。",
                "extra_data": {
                    "spectral_type": "K1V",
                    "mass_solar": 0.70,
                    "radius_solar": 0.67,
                    "temperature_k": 5000
                }
            },
            {
                "id": "star-347-secondary",
                "name": "HD 41004 B",
                "name_zh": "HD 41004 B",
                "description": "红矮星伴星，距离A星约23 AU，可能拥有棕矮星伴星。",
                "extra_data": {
                    "spectral_type": "M2V",
                    "mass_solar": 0.40,
                    "radius_solar": 0.39,
                    "temperature_k": 3400,
                    "separation_au": 23
                }
            }
        ]
    },

    # 8. GJ 86 (格利泽86) - 双星系统（橙矮星 + 白矮星）
    128: {
        "stars": [
            {
                "id": "star-128-primary",
                "name": "GJ 86 A",
                "name_zh": "格利泽86A",
                "description": "橙矮星，拥有一颗类木行星GJ 86 b，伴星是罕见的白矮星。",
                "extra_data": {
                    "spectral_type": "K1V",
                    "mass_solar": 0.79,
                    "radius_solar": 0.77,
                    "temperature_k": 5100
                }
            },
            {
                "id": "star-128-secondary",
                "name": "GJ 86 B",
                "name_zh": "格利泽86B",
                "description": "白矮星伴星，距离A星约21 AU，是研究恒星演化对行星影响的重要案例。",
                "extra_data": {
                    "spectral_type": "DA (白矮星)",
                    "mass_solar": 0.55,
                    "radius_solar": 0.01,
                    "temperature_k": 8000,
                    "separation_au": 21,
                    "is_white_dwarf": True
                }
            }
        ]
    },

    # 9. HD 196885 - 双星系统
    267: {
        "stars": [
            {
                "id": "star-267-primary",
                "name": "HD 196885 A",
                "name_zh": "HD 196885 A",
                "description": "黄白主序星，拥有一颗行星HD 196885 A b。",
                "extra_data": {
                    "spectral_type": "F8V",
                    "mass_solar": 1.33,
                    "radius_solar": 1.68,
                    "temperature_k": 6172
                }
            },
            {
                "id": "star-267-secondary",
                "name": "HD 196885 B",
                "name_zh": "HD 196885 B",
                "description": "红矮星伴星，距离A星约25 AU。",
                "extra_data": {
                    "spectral_type": "M",
                    "mass_solar": 0.45,
                    "radius_solar": 0.43,
                    "temperature_k": 3500,
                    "separation_au": 25
                }
            }
        ]
    }
}


async def add_missing_stars(conn):
    """添加缺失的恒星"""
    logger.info("=" * 80)
    logger.info("开始补全多恒星系统数据")
    logger.info("=" * 80)

    added_count = 0
    skipped_count = 0

    for system_id, system_data in MULTI_STAR_SYSTEMS.items():
        # 检查系统是否存在
        system = await conn.fetchrow(
            "SELECT id, name, name_zh FROM star_systems WHERE id = $1",
            system_id
        )

        if not system:
            logger.warning(f"\n⚠️  系统ID={system_id}不存在，跳过")
            continue

        logger.info(f"\n{'='*80}")
        logger.info(f"处理恒星系统: {system['name_zh'] or system['name']} (ID={system_id})")
        logger.info(f"{'='*80}")

        for star in system_data["stars"]:
            # 检查是否已存在
            existing = await conn.fetchrow(
                "SELECT id FROM celestial_bodies WHERE id = $1",
                star["id"]
            )

            if existing:
                logger.info(f"  ✓ 恒星已存在: {star['name_zh']} ({star['id']})")
                skipped_count += 1
            else:
                # 插入新恒星
                await conn.execute("""
                    INSERT INTO celestial_bodies
                    (id, name, name_zh, type, system_id, description, is_active, extra_data, created_at, updated_at)
                    VALUES
                    ($1, $2, $3, 'star', $4, $5, TRUE, $6::jsonb, NOW(), NOW())
                """,
                    star["id"],
                    star["name"],
                    star["name_zh"],
                    system_id,
                    star["description"],
                    json.dumps(star["extra_data"])
                )
                logger.info(f"  ✅ 添加恒星: {star['name_zh']} ({star['id']})")
                added_count += 1

    logger.info(f"\n{'='*80}")
    logger.info(f"恒星数据补全完成！")
    logger.info(f"  新增: {added_count}颗")
    logger.info(f"  跳过: {skipped_count}颗（已存在）")
    logger.info(f"{'='*80}")


async def verify_results(conn):
    """验证结果"""
    logger.info("\n" + "=" * 80)
    logger.info("验证多星系统数据")
    logger.info("=" * 80)

    system_ids = list(MULTI_STAR_SYSTEMS.keys())

    rows = await conn.fetch("""
        SELECT
            s.id,
            s.name,
            s.name_zh,
            COUNT(CASE WHEN cb.type = 'star' THEN 1 END) as star_count,
            COUNT(CASE WHEN cb.type = 'planet' THEN 1 END) as planet_count,
            COUNT(CASE WHEN cb.is_active = TRUE THEN 1 END) as active_count,
            string_agg(
                CASE WHEN cb.type = 'star' THEN cb.name_zh || ' (' || cb.id || ')' END,
                ', '
                ORDER BY cb.id
            ) as star_names
        FROM star_systems s
        LEFT JOIN celestial_bodies cb ON s.id = cb.system_id
        WHERE s.id = ANY($1)
        GROUP BY s.id, s.name, s.name_zh
        ORDER BY s.id
    """, system_ids)

    print(f"\n{'系统ID':<8} {'系统名称':<30} {'恒星数':<8} {'行星数':<8} {'启用数':<8}")
    print("=" * 100)
    for row in rows:
        system_name = row['name_zh'] or row['name']
        print(f"{row['id']:<8} {system_name:<30} {row['star_count']:<8} {row['planet_count']:<8} {row['active_count']:<8}")

    # 详细显示每个系统的恒星
    print(f"\n{'='*100}")
    print("各系统恒星详情:")
    print(f"{'='*100}")

    for row in rows:
        system_name = row['name_zh'] or row['name']
        stars = await conn.fetch("""
            SELECT id, name, name_zh, extra_data
            FROM celestial_bodies
            WHERE system_id = $1 AND type = 'star'
            ORDER BY id
        """, row['id'])

        print(f"\n{system_name} (ID={row['id']}):")
        for star in stars:
            # Handle both dict and JSON string
            extra = star['extra_data']
            if isinstance(extra, str):
                extra = json.loads(extra) if extra else {}
            elif extra is None:
                extra = {}

            spectral = extra.get('spectral_type', 'N/A')
            mass = extra.get('mass_solar', extra.get('mass_jupiter'))
            mass_unit = 'M☉' if 'mass_solar' in extra else ('MJ' if 'mass_jupiter' in extra else '')
            print(f"  • {star['name_zh']:<25} | 光谱: {spectral:<10} | 质量: {mass}{mass_unit if mass else 'N/A'}")


async def main():
    """主函数"""
    print("\n" + "=" * 80)
    print("多恒星系统数据补全脚本 v2.0")
    print("将补全8-10个高价值双星/多星系统")
    print("=" * 80)

    conn = await asyncpg.connect(**DB_CONFIG)

    try:
        # 1. 添加缺失的恒星
        await add_missing_stars(conn)

        # 2. 验证结果
        await verify_results(conn)

        print("\n" + "=" * 80)
        print("✅ 所有操作完成！")
        print("=" * 80)

    except Exception as e:
        logger.error(f"❌ 发生错误: {e}")
        import traceback
        traceback.print_exc()
    finally:
        await conn.close()


if __name__ == "__main__":
    asyncio.run(main())