from pathlib import Path
from unittest.mock import MagicMock

import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from pingcap_datapipelines.config import PipelineConfig
from pingcap_datapipelines.extract import materialize_empty_tidb_tables, tidb_empty_arrow_type


def config():
    return PipelineConfig(Path('source.yaml'), {
        'source': {'database': 'db', 'name': 'test', 'type': 'tidbcloud'},
        'target': {'project': 'project'}, 'export': {'parquet_compression': 'ZSTD'},
    })


@pytest.fixture
def empty_export(monkeypatch):
    prefix = 'gs://bucket/run/attempt=0002/'
    monkeypatch.setattr('pingcap_datapipelines.extract.list_objects', lambda _: [
        {'name': 'run/attempt=0002/db.Empty-schema.sql'},
    ])
    monkeypatch.setattr('pingcap_datapipelines.extract.download_text',
                        lambda _: 'SHOW MASTER STATUS:\n\tPos: 123456\nFinished dump at: now\n')
    monkeypatch.setattr('pingcap_datapipelines.extract.load_table_metadata', lambda *a, **kw: {
        'Empty': {'columns': [
            {'name': 'Id', 'data_type': 'bigint', 'column_type': 'bigint', 'nullable': False},
            {'name': 'Amount', 'data_type': 'decimal', 'column_type': 'decimal(18,2)', 'nullable': True},
        ]},
    })
    connection = MagicMock()
    cursor = connection.__enter__.return_value.cursor.return_value.__enter__.return_value
    cursor.fetchone.return_value = None
    monkeypatch.setattr('pingcap_datapipelines.extract.tidb_connect', lambda _: connection)
    uploaded = []
    client = MagicMock()
    def upload(path, **kwargs):
        table = pq.read_table(path)
        assert table.num_rows == 0
        assert table.schema.field('Id').type == pa.int64()
        assert table.schema.field('Amount').type == pa.decimal128(18, 2)
        assert kwargs['if_generation_match'] == 0
        uploaded.append(table)
    client.bucket.return_value.blob.return_value.upload_from_filename.side_effect = upload
    monkeypatch.setattr('pingcap_datapipelines.extract.storage.Client', lambda **kw: client)
    return prefix, cursor, uploaded


def test_empty_table_has_snapshot_bound_zero_row_parquet(empty_export):
    prefix, cursor, uploaded = empty_export
    result = materialize_empty_tidb_tables(config(), prefix, ['db.Empty'], 'exp-ok')
    assert len(uploaded) == 1
    assert result[0]['snapshot_tso'] == '123456'
    assert result[0]['row_count'] == 0
    assert result[0]['export_id'] == 'exp-ok'
    assert result[0]['uri'] == prefix + 'db.Empty.empty.parquet'
    assert cursor.execute.call_args_list[0].args == ('SET @@tidb_snapshot = %s', ('123456',))


def test_nonempty_missing_parquet_is_not_replaced_with_empty_file(empty_export):
    prefix, cursor, uploaded = empty_export
    cursor.fetchone.return_value = {'1': 1}
    with pytest.raises(RuntimeError, match='has rows at export snapshot'):
        materialize_empty_tidb_tables(config(), prefix, ['db.Empty'], 'exp-ok')
    assert not uploaded


def test_missing_schema_is_not_evidence_of_empty_table(empty_export, monkeypatch):
    prefix, cursor, uploaded = empty_export
    monkeypatch.setattr('pingcap_datapipelines.extract.list_objects', lambda _: [])
    with pytest.raises(RuntimeError, match='missing schema artifact'):
        materialize_empty_tidb_tables(config(), prefix, ['db.Empty'], 'exp-ok')
    assert not uploaded
    cursor.execute.assert_not_called()


@pytest.mark.parametrize('metadata', ['Pos: 123456', 'Finished dump at: now'])
def test_incomplete_snapshot_evidence_fails_closed(empty_export, monkeypatch, metadata):
    prefix, _, uploaded = empty_export
    monkeypatch.setattr('pingcap_datapipelines.extract.download_text', lambda _: metadata)
    with pytest.raises(RuntimeError, match='completed export snapshot'):
        materialize_empty_tidb_tables(config(), prefix, ['db.Empty'], 'exp-ok')
    assert not uploaded


def test_snapshot_expired_fails_closed(empty_export):
    prefix, cursor, uploaded = empty_export
    cursor.execute.side_effect = RuntimeError('snapshot expired')
    with pytest.raises(RuntimeError, match='snapshot expired'):
        materialize_empty_tidb_tables(config(), prefix, ['db.Empty'], 'exp-ok')
    assert not uploaded


@pytest.mark.parametrize(('kind', 'declared', 'expected'), [
    ('date', 'date', pa.string()), ('time', 'time(6)', pa.string()),
    ('bigint', 'bigint unsigned', pa.decimal128(20, 0)),
    ('decimal', 'decimal(65,2)', pa.string()), ('float', 'float', pa.float64()),
    ('timestamp', 'timestamp(6)', pa.timestamp('us')), ('blob', 'blob', pa.binary()),
])
def test_empty_parquet_types_match_tidb_export(kind, declared, expected):
    assert tidb_empty_arrow_type({'data_type': kind, 'column_type': declared}) == expected