"""
control/tenant_migrate.py
-------------------------
Per-tenant migration runner (MT-2).

Runs the EXISTING tenant Alembic chain (migrations/versions) against each
tenant's own database via the standalone migrations_tenant/env.py. Records the
applied head on tenants.alembic_head and writes a ProvisioningJob row per run.

The existing migrations are re-run-safe (INFORMATION_SCHEMA guards, Rule 14),
so upgrading an already-current tenant (e.g. tenant-zero / LT, whose DB is
already at head) is a no-op.

Usage:
    python -m control.tenant_migrate upgrade --tenant all
    python -m control.tenant_migrate upgrade --tenant lts       # id or slug
    python -m control.tenant_migrate current --tenant all
    python -m control.tenant_migrate heads

`upgrade_tenant()` is also imported by the MT-3 provisioning service to bring a
freshly-created tenant DB up to head.
"""

import argparse
import contextlib
import io
import os
import sys
from collections import namedtuple

from alembic import command
from alembic.config import Config
from alembic.script import ScriptDirectory
from alembic.runtime.migration import MigrationContext
from sqlalchemy import create_engine

from control.base import control_session
from control.models import Tenant, ProvisioningJob
from control.time_utils import now_eastern

# Repo layout: control/ and migrations/ are siblings at the repo root.
_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_DEFAULT_SCRIPT_LOCATION = os.path.join(_REPO_ROOT, 'migrations_tenant')
_DEFAULT_VERSION_LOCATIONS = os.path.join(_REPO_ROOT, 'migrations', 'versions')

# Everything except a deleted tenant should be kept schema-current.
MIGRATABLE_STATUSES = ('provisioning', 'active', 'suspended')

# Squashed baseline that builds the full schema (migrations/versions root).
BASELINE_REVISION = '0003_add_user_active'

TenantRef = namedtuple('TenantRef', ['id', 'slug', 'db_uri'])


def _make_config(db_uri, script_location=None, version_locations=None):
    cfg = Config()
    cfg.set_main_option('script_location', script_location or _DEFAULT_SCRIPT_LOCATION)
    cfg.set_main_option('version_locations', version_locations or _DEFAULT_VERSION_LOCATIONS)
    # Escape % so ConfigParser interpolation doesn't choke on URL-encoded passwords.
    cfg.set_main_option('sqlalchemy.url', db_uri.replace('%', '%%'))
    return cfg


def chain_head(script_location=None, version_locations=None):
    """Return the head revision id of the tenant migration chain."""
    cfg = _make_config('sqlite://', script_location, version_locations)
    script = ScriptDirectory.from_config(cfg)
    heads = script.get_heads()
    return heads[0] if len(heads) == 1 else ','.join(sorted(heads))


def current_revision(db_uri):
    """Return the alembic revision currently stamped on the given database."""
    engine = create_engine(db_uri, future=True)
    try:
        with engine.connect() as conn:
            return MigrationContext.configure(conn).get_current_revision()
    finally:
        engine.dispose()


def upgrade_tenant(tenant, script_location=None, version_locations=None, record_job=True):
    """Upgrade one tenant's database to head.

    `tenant` must expose `.id` and `.db_uri` (a control ORM Tenant, a TenantRef,
    or any object with those attributes). Returns the applied head revision.
    Records tenants.alembic_head and a ProvisioningJob('migrate') unless
    record_job is False.
    """
    db_uri = tenant.db_uri
    cfg = _make_config(db_uri, script_location, version_locations)

    job_id = None
    if record_job:
        with control_session() as s:
            job = ProvisioningJob(tenant_id=tenant.id, action='migrate',
                                  status='running', created_at=now_eastern())
            s.add(job)
            s.flush()
            job_id = job.id

    try:
        # Alembic logs progress to stdout; keep the runner output tidy.
        with contextlib.redirect_stdout(io.StringIO()):
            command.upgrade(cfg, 'head')
        applied = current_revision(db_uri)

        with control_session() as s:
            t = s.get(Tenant, tenant.id)
            if t is not None:
                t.alembic_head = applied
            if job_id is not None:
                j = s.get(ProvisioningJob, job_id)
                if j is not None:
                    j.status = 'ok'
                    j.finished_at = now_eastern()
                    j.log = f'upgraded to {applied}'
        return applied

    except Exception as e:
        if job_id is not None:
            with control_session() as s:
                j = s.get(ProvisioningJob, job_id)
                if j is not None:
                    j.status = 'failed'
                    j.finished_at = now_eastern()
                    j.log = f'{type(e).__name__}: {e}'
        raise


def bootstrap_tenant(tenant, script_location=None, version_locations=None,
                     baseline_rev=BASELINE_REVISION, record_job=True):
    """Build a FRESH tenant database, then mark it current.

    Runs the squashed baseline (full schema) only, then `stamp head` — the
    historical phase migrations are NOT replayed (several are not idempotent and
    would conflict with the full-schema baseline). Use this for brand-new tenant
    databases; use upgrade_tenant() for ongoing incremental migrations.

    Returns the stamped head revision.
    """
    db_uri = tenant.db_uri
    cfg = _make_config(db_uri, script_location, version_locations)

    job_id = None
    if record_job:
        with control_session() as s:
            job = ProvisioningJob(tenant_id=tenant.id, action='migrate',
                                  status='running', created_at=now_eastern())
            s.add(job)
            s.flush()
            job_id = job.id

    try:
        with contextlib.redirect_stdout(io.StringIO()):
            command.upgrade(cfg, baseline_rev)   # build full schema (baseline only)
            command.stamp(cfg, 'head')           # mark at head without replaying phases
        applied = current_revision(db_uri)

        with control_session() as s:
            t = s.get(Tenant, tenant.id)
            if t is not None:
                t.alembic_head = applied
            if job_id is not None:
                j = s.get(ProvisioningJob, job_id)
                if j is not None:
                    j.status = 'ok'
                    j.finished_at = now_eastern()
                    j.log = f'bootstrapped (baseline {baseline_rev}) + stamped {applied}'
        return applied

    except Exception as e:
        if job_id is not None:
            with control_session() as s:
                j = s.get(ProvisioningJob, job_id)
                if j is not None:
                    j.status = 'failed'
                    j.finished_at = now_eastern()
                    j.log = f'{type(e).__name__}: {e}'
        raise

def _select_tenants(selector):
    with control_session() as s:
        q = s.query(Tenant).filter(Tenant.status.in_(MIGRATABLE_STATUSES))
        if selector != 'all':
            if selector.isdigit():
                q = q.filter(Tenant.id == int(selector))
            else:
                q = q.filter(Tenant.slug == selector)
        return [TenantRef(t.id, t.slug, t.db_uri) for t in q.order_by(Tenant.id).all()]


def _cmd_upgrade(args):
    tenants = _select_tenants(args.tenant)
    if not tenants:
        print(f'No matching tenants for --tenant {args.tenant}.')
        return 0
    head = chain_head()
    print(f'Target head: {head}')
    failures = 0
    for t in tenants:
        try:
            applied = upgrade_tenant(t)
            print(f'  [ok]   {t.slug} (id={t.id}) -> {applied}')
        except Exception as e:
            failures += 1
            print(f'  [FAIL] {t.slug} (id={t.id}): {type(e).__name__}: {e}')
    print(f'Done. {len(tenants) - failures}/{len(tenants)} upgraded.')
    return 1 if failures else 0


def _cmd_current(args):
    tenants = _select_tenants(args.tenant)
    if not tenants:
        print(f'No matching tenants for --tenant {args.tenant}.')
        return 0
    head = chain_head()
    print(f'Chain head: {head}')
    for t in tenants:
        try:
            rev = current_revision(t.db_uri)
        except Exception as e:
            rev = f'<error: {type(e).__name__}>'
        flag = '' if rev == head else '  (BEHIND)'
        print(f'  {t.slug} (id={t.id}): {rev}{flag}')
    return 0


def _cmd_bootstrap(args):
    tenants = _select_tenants(args.tenant)
    if not tenants:
        print(f'No matching tenants for --tenant {args.tenant}.')
        return 0
    failures = 0
    for t in tenants:
        try:
            applied = bootstrap_tenant(t)
            print(f'  [ok]   {t.slug} (id={t.id}) bootstrapped -> {applied}')
        except Exception as e:
            failures += 1
            print(f'  [FAIL] {t.slug} (id={t.id}): {type(e).__name__}: {e}')
    return 1 if failures else 0


def _cmd_heads(_args):
    print(chain_head())
    return 0


def main(argv=None):
    parser = argparse.ArgumentParser(prog='control.tenant_migrate',
                                     description='Per-tenant migration runner')
    sub = parser.add_subparsers(dest='command', required=True)

    p_up = sub.add_parser('upgrade', help='Upgrade tenant DB(s) to head')
    p_up.add_argument('--tenant', required=True, help="'all', or a tenant id or slug")
    p_up.set_defaults(func=_cmd_upgrade)

    p_cur = sub.add_parser('current', help='Show each tenant DB current revision')
    p_cur.add_argument('--tenant', default='all', help="'all', or a tenant id or slug")
    p_cur.set_defaults(func=_cmd_current)

    p_boot = sub.add_parser('bootstrap',
                            help='Build a FRESH tenant DB (baseline + stamp head)')
    p_boot.add_argument('--tenant', required=True, help="'all', or a tenant id or slug")
    p_boot.set_defaults(func=_cmd_bootstrap)

    sub.add_parser('heads', help='Show the chain head revision').set_defaults(func=_cmd_heads)

    args = parser.parse_args(argv)
    sys.exit(args.func(args))


if __name__ == '__main__':
    main()