"""
MVP 数据库准备脚本

1. 检查 test_zqyy_app 库中 test schema 和 xcx-test 表是否存在
2. 如果不存在则创建
3. 插入测试数据 "t91"
"""

import os
import sys
from pathlib import Path

# 加载根 .env
from dotenv import load_dotenv
load_dotenv(Path(__file__).resolve().parents[2] / ".env", override=False)

import psycopg2

DB_HOST = os.getenv("DB_HOST", "localhost")
DB_PORT = os.getenv("DB_PORT", "5432")
DB_USER = os.getenv("DB_USER", "")
DB_PASSWORD = os.getenv("DB_PASSWORD", "")
APP_DB_NAME = os.getenv("APP_DB_NAME", "test_zqyy_app")


def main():
    print(f"连接数据库: {DB_HOST}:{DB_PORT}/{APP_DB_NAME} (用户: {DB_USER})")

    conn = psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        dbname=APP_DB_NAME,
    )
    conn.autocommit = True

    try:
        with conn.cursor() as cur:
            # 1. 检查 test schema
            cur.execute(
                "SELECT 1 FROM information_schema.schemata WHERE schema_name = 'test'"
            )
            if cur.fetchone():
                print("✓ test schema 已存在")
            else:
                cur.execute("CREATE SCHEMA test")
                print("✓ test schema 已创建")

            # 2. 检查 xcx-test 表（注意表名含连字符，需要双引号）
            cur.execute("""
                SELECT 1 FROM information_schema.tables
                WHERE table_schema = 'test' AND table_name = 'xcx-test'
            """)
            if cur.fetchone():
                print('✓ test."xcx-test" 表已存在')
            else:
                cur.execute("""
                    CREATE TABLE test."xcx-test" (
                        ti TEXT
                    )
                """)
                print('✓ test."xcx-test" 表已创建')

            # 3. 查看现有数据
            cur.execute('SELECT ti FROM test."xcx-test" LIMIT 5')
            rows = cur.fetchall()
            if rows:
                print(f"  现有数据: {[r[0] for r in rows]}")

            # 4. 插入 "t91"
            cur.execute(
                'INSERT INTO test."xcx-test" (ti) VALUES (%s)', ("t91",)
            )
            print('✓ 已插入 ti = "t91"')

            # 5. 验证
            cur.execute('SELECT ti FROM test."xcx-test" ORDER BY ti')
            rows = cur.fetchall()
            print(f"  当前全部数据: {[r[0] for r in rows]}")

    finally:
        conn.close()

    print("\n数据库准备完成。")


if __name__ == "__main__":
    main()
