"""
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'

#: The last revision the squashed baseline actually covers.
#:
#: Everything AFTER this point is existence-guarded (INFORMATION_SCHEMA checks,
#: or an idempotent ENUM MODIFY), so it can be replayed over the baseline
#: safely. Everything BEFORE it is not — several pre-phase33 migrations would
#: collide with the full-schema baseline, which is why bootstrap stamps past
#: them instead of running them.
#:
#: If the baseline is ever re-squashed to a later point, move this with it.
BASELINE_COVERS_THROUGH = 'phase32_device_token_columns' 

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 bring it to head.

    Three steps, and the middle one is the load-bearing part:

      1. `upgrade(baseline)` — the squashed baseline builds the bulk of the
         schema in one go.
      2. `stamp(BASELINE_COVERS_THROUGH)` — declare the DB to be at the last
         revision the baseline actually covers, WITHOUT running the pre-phase33
         migrations. Those are not idempotent and would collide with the
         baseline; skipping them is the whole reason bootstrap exists.
      3. `upgrade(head)` — replay the guarded tail (phase33 onward). Every one
         of those checks INFORMATION_SCHEMA before touching anything (the two
         ENUM widenings use an idempotent MODIFY), so this adds exactly what
         the baseline lacks and skips the rest.

    This used to stop after step 1 and `stamp('head')` instead — claiming the
    database was current when it was missing every column and table added since
    phase32. Because the baseline was last refreshed around phase33, a tenant
    provisioned that way lacked `users.mfa_enabled` and `users.ui_theme`, and
    SQLAlchemy emits every mapped column in its SELECT — so the new workspace
    could not even log in. It failed at the first query, not at some optional
    feature, and only for freshly provisioned tenants (tenant-zero was adopted
    in place with a real schema), which is what kept it hidden.

    Use upgrade_tenant() for ongoing incremental migrations of a tenant that
    already exists.

    Returns the head revision the database ends up at.
    """
    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)              # 1. baseline schema
            command.stamp(cfg, BASELINE_COVERS_THROUGH)     # 2. skip the unguarded past
            command.upgrade(cfg, 'head')                    # 3. replay the guarded tail
        applied = current_revision(db_uri)

        # A bootstrap that does not end at head has silently produced a broken
        # tenant — exactly the failure this sequence exists to prevent. Say so
        # here rather than letting it surface as a missing column later.
        expected = chain_head(script_location, version_locations)
        if applied != expected:
            raise RuntimeError(
                f'bootstrap ended at {applied!r}, expected head {expected!r} — '
                f'the tenant database is incomplete')

        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 '
                             f'{BASELINE_COVERS_THROUGH}, 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 _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 _cmd_stamp(args):
    """Stamp tenant DB(s) to a specific revision without running migrations.

    Use this to rewind a mistakenly-stamped tenant so a subsequent `upgrade`
    will actually replay the missing migration.  Example:

        python -m control.tenant_migrate stamp --tenant gov phase32_device_token_columns
        python -m control.tenant_migrate upgrade --tenant gov
    """
    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:
            cfg = _make_config(t.db_uri)
            with contextlib.redirect_stdout(io.StringIO()):
                command.stamp(cfg, args.revision)
            rev = current_revision(t.db_uri)
            print(f'  [ok]   {t.slug} (id={t.id}) stamped -> {rev}')
        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 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)

    p_stamp = sub.add_parser('stamp', help='Stamp tenant DB(s) to a specific revision (no migrations run)')
    p_stamp.add_argument('--tenant', required=True, help="'all', or a tenant id or slug")
    p_stamp.add_argument('revision', help='Alembic revision id (e.g. phase32_device_token_columns)')
    p_stamp.set_defaults(func=_cmd_stamp)

    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()