from datetime import datetime, timezone, timedelta
from decimal import Decimal
from pathlib import Path
import os
import subprocess
import sys

from psycopg import sql
import pytest

from btcdata_archive.archive import run
from btcdata_archive.postgres import Postgres
from btcdata_archive.schema import TABLES, encode_row, decode_row
from btcdata_archive.sqlite_archive import connect, validate
from btcdata_archive.weeks import boundaries, WEEK

pytestmark = pytest.mark.integration
CUTOFF = boundaries(datetime.now(timezone.utc))[1]
START = CUTOFF - 3 * WEEK


def seed(infra):
    for offset, pk in [(0, 1), (14, 3), (21, 4)]:
        instant = START + timedelta(days=offset, microseconds=123456)
        infra.admin.execute('INSERT INTO binance_price VALUES (%s,%s,%s)', (pk, Decimal('1234567890.12345678'), instant))
        infra.admin.execute("INSERT INTO strike_price VALUES (%s,'{\"a\":12345678901234567890.12345678901234567890,\"null\":null}'::jsonb,%s,NULL)", (pk, instant))
        infra.admin.execute('INSERT INTO kraken_minute VALUES (%s,1.00000001,2,3,4,5,1234567890123456.12345678,7,%s)', (int(instant.timestamp()), instant))


def invoke(infra, mode='archive-and-prune', max_weeks=None):
    return run(infra.pg, infra.store, infra.config, CUTOFF, mode, max_weeks)


def test_export_gap_precision_cutoff_restore(infra, tmp_path):
    seed(infra)
    original = {}
    for table in TABLES:
        projection = ','.join(n + '::text' if typ == 'jsonb' else n for n, typ, _ in table.columns)
        original[table.name] = [encode_row(table, row) for row in infra.admin.execute(f'SELECT {projection} FROM public.{table.name} ORDER BY {table.pk}')]
    # Exactly at cutoff must stay, regardless of current week.
    infra.admin.execute('INSERT INTO binance_price VALUES (5,5,%s)', (CUTOFF,))
    invoke(infra, 'verify-only')
    assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 4
    assert infra.admin.execute("SELECT count(*) FROM btcdata_archive.weeks WHERE status='verified'").fetchone()[0] == 3
    invoke(infra)
    assert infra.admin.execute('SELECT id FROM binance_price ORDER BY id').fetchall() == [(4,), (5,)]
    assert infra.admin.execute('SELECT count(*) FROM kraken_minute').fetchone()[0] == 1
    for offset in (0, 7, 14):
        start = START + timedelta(days=offset)
        key = infra.config.prefix + start.date().isoformat() + '.sqlite'
        head = infra.store.head(key)
        path = tmp_path / ('restored-' + str(offset) + '.sqlite')
        infra.store.download(key, head['VersionId'], path)
        meta = validate(path, start, infra.config.pg['dbname'], 2)
        assert all(v['count'] == (0 if offset == 7 else 1) for v in meta['tables'].values())
        # Restore every column through the documented type mapping into isolated PostgreSQL tables.
        db = connect(path, readonly=True)
        for table in TABLES:
            restored = 'restored_' + str(offset) + '_' + table.name
            infra.admin.execute(sql.SQL('CREATE TABLE {} (LIKE public.{})').format(sql.Identifier(restored), sql.Identifier(table.name)))
            rows = db.execute(f'SELECT * FROM {table.name} ORDER BY {table.pk}').fetchall()
            for row in rows:
                placeholders = ','.join('%s::jsonb' if typ == 'jsonb' else '%s' for _, typ, _ in table.columns)
                infra.admin.execute(sql.SQL(f'INSERT INTO {{}} VALUES ({placeholders})').format(sql.Identifier(restored)), decode_row(table, row))
            projection = ','.join(n + '::text' if typ == 'jsonb' else n for n, typ, _ in table.columns)
            roundtrip = [encode_row(table, r) for r in infra.admin.execute(sql.SQL(f'SELECT {projection} FROM {{}} ORDER BY {table.pk}').format(sql.Identifier(restored)))]
            assert roundtrip == rows
            assert all(r in original[table.name] for r in rows)
        db.close()
    assert infra.admin.execute("SELECT count(*) FROM btcdata_archive.weeks WHERE status='complete'").fetchone()[0] == 3
    # Application latest-value reads still work with the retained rows.
    assert infra.admin.execute('SELECT id FROM binance_price ORDER BY recorded_at DESC LIMIT 1').fetchone()[0] == 4


def test_schema_mismatch(infra):
    seed(infra)
    infra.admin.execute('ALTER TABLE strike_price ADD COLUMN unexpected text')
    with pytest.raises(RuntimeError, match='migration required'):
        invoke(infra)
    assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 3


def test_concurrent_lock_and_loss(infra, monkeypatch):
    seed(infra)
    second = Postgres(infra.config)
    infra.pg.lock()
    try:
        with pytest.raises(RuntimeError, match='advisory lock'):
            second.lock()
    finally:
        second.close()
    # Closing the coordinating session after upload must prevent ANY deletion.
    upload = infra.store.upload
    def lose_lock(*args):
        result = upload(*args)
        infra.pg.close()
        return result
    monkeypatch.setattr(infra.store, 'upload', lose_lock)
    # Don't acquire a second lock on the same session in run.
    monkeypatch.setattr(infra.pg, 'lock', lambda: None)
    with pytest.raises(Exception):
        invoke(infra)
    assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 3


def test_changed_and_late_rows(infra, monkeypatch):
    seed(infra)
    upload = infra.store.upload
    def change(*args):
        version = upload(*args)
        infra.admin.execute('UPDATE binance_price SET price=9 WHERE id=1')
        infra.admin.execute('INSERT INTO binance_price VALUES (2,2,%s)', (START,))
        return version
    monkeypatch.setattr(infra.store, 'upload', change)
    with pytest.raises(RuntimeError, match='reconciliation required'):
        invoke(infra)
    assert infra.admin.execute('SELECT id FROM binance_price ORDER BY id').fetchall() == [(1,), (2,), (3,), (4,)]
    # Other unchanged archived rows were safely pruned; their data must survive the merge.
    assert infra.admin.execute('SELECT count(*) FROM strike_price').fetchone()[0] == 2
    monkeypatch.setattr(infra.store, 'upload', upload)
    invoke(infra)
    key = infra.config.prefix + START.date().isoformat() + '.sqlite'
    path = infra.config.work_dir / 'late.sqlite'
    infra.store.download(key, infra.store.head(key)['VersionId'], path)
    db = connect(path, readonly=True)
    assert db.execute('SELECT id,price FROM binance_price ORDER BY id').fetchall() == [(1,'9.00000000'), (2,'2.00000000')]
    assert db.execute('SELECT count(*) FROM strike_price').fetchone()[0] == 1
    db.close()


@pytest.mark.parametrize('boundary', ['uploaded', 'verified', 'pruning', 'deletion_batch', 'complete'])
def test_crash_restart_boundaries(infra, monkeypatch, boundary):
    seed(infra)
    ledger, upload = infra.pg.ledger, infra.store.upload
    fired = False
    def crash_ledger(start, status, *args, **kwargs):
        nonlocal fired
        result = ledger(start, status, *args, **kwargs)
        target = 'pruning' if boundary == 'deletion_batch' else boundary
        if not fired and status == target and (boundary != 'deletion_batch' or kwargs.get('deletions')):
            fired = True
            raise KeyboardInterrupt()
        return result
    def crash_upload(*args):
        result = upload(*args)
        raise KeyboardInterrupt()
    monkeypatch.setattr(infra.pg, 'ledger', crash_ledger)
    if boundary == 'uploaded':
        monkeypatch.setattr(infra.store, 'upload', crash_upload)
    with pytest.raises(KeyboardInterrupt):
        invoke(infra)
    infra.pg.close()
    infra.pg = Postgres(infra.config)
    monkeypatch.setattr(infra.store, 'upload', upload)
    invoke(infra)
    key = infra.config.prefix + START.date().isoformat() + '.sqlite'
    path = infra.config.work_dir / 'resumed.sqlite'
    infra.store.download(key, infra.store.head(key)['VersionId'], path)
    assert all(s['count'] == 1 for s in validate(path, START, infra.config.pg['dbname'], 2)['tables'].values())
    assert infra.admin.execute('SELECT id FROM binance_price').fetchall() == [(4,)]


@pytest.mark.parametrize('operation', ['upload', 'download'])
def test_network_failure_preserves_data(infra, monkeypatch, operation):
    seed(infra)
    def fail(*args):
        raise OSError('network unavailable')
    monkeypatch.setattr(infra.store, operation, fail)
    with pytest.raises(OSError):
        invoke(infra)
    assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 3


def test_missing_delete_privilege(infra):
    seed(infra)
    infra.admin.execute(sql.SQL('REVOKE DELETE ON binance_price FROM {}').format(sql.Identifier(infra.config.pg['user'])))
    with pytest.raises(Exception):
        invoke(infra)
    assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 3


def test_wrapper_all_weeks_and_bound(infra):
    seed(infra)
    c = infra.config
    env = {**os.environ, 'PGHOST': c.pg['host'], 'PGPORT': str(c.pg['port']), 'PGDATABASE': c.pg['dbname'],
           'PGUSER': c.pg['user'], 'PGPASSWORD': c.pg['password'], 'S3_ENDPOINT_URL': c.endpoint,
           'AWS_ACCESS_KEY_ID': c.access_key, 'AWS_SECRET_ACCESS_KEY': c.secret_key, 'ARCHIVE_BUCKET': c.bucket,
           'ARCHIVE_WORK_DIR': str(c.work_dir), 'PATH': str(Path(sys.executable).parent) + ':' + os.environ['PATH']}
    script = str(Path(__file__).resolve().parents[2] / 'scripts/archive-old-weeks.sh')
    first = subprocess.run([script, '--mode', 'verify-only', '--max-weeks', '1'], env=env, capture_output=True, text=True)
    assert first.returncode == 0, first.stdout
    second = subprocess.run([script, '--mode', 'archive-and-prune'], env=env, capture_output=True, text=True)
    assert second.returncode == 0, second.stdout
    assert infra.admin.execute("SELECT count(*) FROM btcdata_archive.weeks WHERE status='complete'").fetchone()[0] == 3