#!/usr/bin/env python3
"""
一键迁移：支持 Docker 或本机执行，需在项目根目录运行。
  python deploy/scripts/migrate.py          # 默认用 Docker
  python deploy/scripts/migrate.py --local  # 本机执行（需已安装依赖且数据库可连）
"""
import os
import subprocess
import sys
import time

PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
BACKEND_DIR = os.path.join(PROJECT_ROOT, "backend")


def run(cmd, cwd=None, env=None):
    r = subprocess.run(cmd, cwd=cwd or PROJECT_ROOT, env=env or os.environ.copy(), shell=True)
    if r.returncode != 0:
        sys.exit(r.returncode)


def migrate_docker():
    os.chdir(PROJECT_ROOT)
    print("[1/3] 启动数据库...")
    run("docker compose up -d db")
    print("[2/3] 等待数据库就绪...")
    for i in range(30):
        r = subprocess.run(
            "docker compose exec -T db pg_isready -U wecom -d wecom_ai",
            shell=True,
            cwd=PROJECT_ROOT,
            capture_output=True,
        )
        if r.returncode == 0:
            break
        time.sleep(1)
    else:
        print("数据库未在 30s 内就绪，请检查 docker compose logs db")
        sys.exit(1)
    print("[3/3] 执行 Alembic 迁移...")
    run('docker compose run --rm backend sh -c "alembic upgrade head"')
    print("迁移完成。")


def migrate_local():
    # 加载 .env
    env_path = os.path.join(PROJECT_ROOT, ".env")
    if os.path.isfile(env_path):
        with open(env_path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    k, _, v = line.partition("=")
                    os.environ[k.strip()] = v.strip()
    if "DATABASE_URL_SYNC" not in os.environ:
        os.environ.setdefault("DATABASE_URL_SYNC", "postgresql://wecom:wecom_secret@localhost:5432/wecom_ai")
    os.chdir(BACKEND_DIR)
    os.environ["PYTHONPATH"] = BACKEND_DIR
    print("本机执行迁移（backend 目录）...")
    run("alembic upgrade head", cwd=BACKEND_DIR)
    print("迁移完成。")


def main():
    if "--local" in sys.argv:
        migrate_local()
    else:
        migrate_docker()


if __name__ == "__main__":
    main()