"""
Environment selection for the migration scripts.

Every transform is environment-agnostic: it reads the staged Parquet (same for
all environments) and writes to whichever database `--env` selects. The target
is defined entirely by `deploy/.env.<env>` (the same file Portainer is fed at
deploy time), whose `DATABASE_URL` is the single source of truth for host /
port / credentials / database.

Reproduce the whole migration in a new environment (e.g. prod) by:
  1. deploy the DB stack for that env (deploy/jorgecuadros-db.stack.yml)
  2. write deploy/.env.<env> with its DATABASE_URL
  3. push the schema:  DATABASE_URL=... npx prisma@5 db push --schema=packages/database/prisma/schema.prisma
  4. run:              ./.venv/bin/python run_all.py --env <env>

Usage in a script:
    from dbenv import connect, env_arg
    env = env_arg()                 # --env dev|prod (default dev)
    conn = connect(env)
"""

from __future__ import annotations

import argparse
from pathlib import Path
from urllib.parse import unquote, urlparse

import pymysql

REPO = Path(__file__).resolve().parents[1]


def load_env(env: str) -> dict:
    f = REPO / "deploy" / f".env.{env}"
    if not f.exists():
        raise SystemExit(
            f"missing {f} — deploy the '{env}' DB stack and write its .env first "
            f"(see dbenv.py header)."
        )
    out = {}
    for line in f.read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            k, v = line.split("=", 1)
            out[k] = v
    if "DATABASE_URL" not in out:
        raise SystemExit(f"{f} has no DATABASE_URL")
    return out


def connect(env: str):
    url = load_env(env)["DATABASE_URL"]
    u = urlparse(url)  # mysql://user:pass@host:port/db
    return pymysql.connect(
        host=u.hostname,
        port=u.port or 3306,
        user=unquote(u.username or ""),
        password=unquote(u.password or ""),
        database=(u.path or "/").lstrip("/"),
        autocommit=False,
        charset="utf8mb4",
    )


def env_arg(extra_args=None) -> str:
    """Parse --env (default 'dev') and return it. Scripts that need more args
    can pass an argparse parser via extra_args(parser)."""
    p = argparse.ArgumentParser()
    p.add_argument("--env", default="dev", help="target environment (dev|prod|...): reads deploy/.env.<env>")
    if extra_args:
        extra_args(p)
    return p.parse_args().env