"""Tests for data ingestion utilities."""

import pytest
import pandas as pd
import duckdb
from walkingpandas.io.ingest import (
    validate_schema,
    convert_to_superset_format,
    write_walks,
    write_duckdb,
    read_dataset,
    read_duckdb,
    connect,
)
from walkingpandas import Network


def test_validate_schema(sample_walks_simple_paths):
    """Test schema validation."""
    df, _ = sample_walks_simple_paths
    assert validate_schema(df, strict=True)


def test_convert_to_superset_format():
    """Test converting DataFrame to superset format."""
    df = pd.DataFrame({
        'id': [1, 1, 2],
        'step': [0, 1, 0],
        'node': ['A', 'B', 'C']
    })
    
    converted = convert_to_superset_format(df)
    assert 'walk_id' in converted.columns
    assert 'sequence_idx' in converted.columns
    assert 'node_id' in converted.columns


def test_write_walks(temp_dir, sample_walks_simple_paths):
    """Test writing walks to Parquet."""
    df, _ = sample_walks_simple_paths
    output_path = temp_dir / 'output' / 'walks.parquet'
    
    write_walks(df, output_path)
    assert output_path.exists()


def test_read_dataset(sample_nodes, sample_edges, sample_walks_simple_paths, temp_dir):
    """Test read_dataset convenience function."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    _, walks_path = sample_walks_simple_paths
    
    walkframe = read_dataset(nodes_path, edges_path, walks_path)
    assert walkframe is not None
    result = walkframe.compute()
    assert len(result) > 0


def test_connect_with_walks(sample_nodes, sample_edges, sample_walks_simple_paths, temp_dir):
    """Test connect() function with walks path."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    _, walks_path = sample_walks_simple_paths
    
    # First save a network
    network = Network(nodes=nodes_path, edges=edges_path)
    db_path = temp_dir / 'test_connect.duckdb'
    network.save(db_path)
    
    # Now use connect() to load both
    walkframe = connect(db_path, walks_path=str(walks_path))
    assert walkframe is not None
    result = walkframe.compute()
    assert len(result) > 0


def test_connect_without_walks(sample_nodes, sample_edges, temp_dir):
    """Test connect() without walks_path or walks_table raises error."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    
    network = Network(nodes=nodes_path, edges=edges_path)
    db_path = temp_dir / 'test_connect_no_walks.duckdb'
    network.save(db_path)
    
    with pytest.raises(ValueError, match="Provide either walks_path"):
        connect(db_path)


def _duckdb_with_network_and_walks_table(temp_dir, nodes_path, edges_path, walks_df, table_name='wp_walks'):
    """Create a .duckdb containing network (from save) and a walks table."""
    network = Network(nodes=nodes_path, edges=edges_path)
    db_path = temp_dir / 'one_file.duckdb'
    network.save(db_path)
    conn = duckdb.connect(str(db_path))
    conn.register('walks_df', walks_df)
    conn.execute(f'CREATE TABLE {table_name} AS SELECT * FROM walks_df')
    conn.close()
    return db_path


def test_read_duckdb(sample_nodes, sample_edges, sample_walks_simple_paths, temp_dir):
    """Test read_duckdb loads from a single .duckdb with network + walks table (deprecated)."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    walks_df, _ = sample_walks_simple_paths

    db_path = _duckdb_with_network_and_walks_table(temp_dir, nodes_path, edges_path, walks_df)
    with pytest.warns(DeprecationWarning, match="read_duckdb is deprecated"):
        walkframe = read_duckdb(db_path)
    assert walkframe is not None
    result = walkframe.compute()
    assert len(result) == len(walks_df)
    assert set(result.columns) >= {'walk_id', 'sequence_idx', 'node_id'}


def test_read_duckdb_deprecation_same_as_connect(sample_nodes, sample_edges, sample_walks_simple_paths, temp_dir):
    """Test read_duckdb issues DeprecationWarning and returns same result as connect(..., walks_table=...)."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    walks_df, _ = sample_walks_simple_paths

    db_path = _duckdb_with_network_and_walks_table(temp_dir, nodes_path, edges_path, walks_df)
    with pytest.warns(DeprecationWarning, match="read_duckdb is deprecated"):
        wf_deprecated = read_duckdb(db_path, walks_table='wp_walks')
    wf_connect = connect(db_path, walks_table='wp_walks')
    result_deprecated = wf_deprecated.compute().sort_values(['walk_id', 'sequence_idx']).reset_index(drop=True)
    result_connect = wf_connect.compute().sort_values(['walk_id', 'sequence_idx']).reset_index(drop=True)
    pd.testing.assert_frame_equal(result_deprecated, result_connect)


def test_read_duckdb_custom_table(sample_nodes, sample_edges, sample_walks_simple_paths, temp_dir):
    """Test read_duckdb with custom walks_table name (deprecated API)."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    walks_df, _ = sample_walks_simple_paths

    db_path = _duckdb_with_network_and_walks_table(
        temp_dir, nodes_path, edges_path, walks_df, table_name='trajectories'
    )
    with pytest.warns(DeprecationWarning, match="read_duckdb is deprecated"):
        walkframe = read_duckdb(db_path, walks_table='trajectories')
    assert walkframe is not None
    result = walkframe.compute()
    assert len(result) == len(walks_df)


def test_connect_with_walks_table(sample_nodes, sample_edges, sample_walks_simple_paths, temp_dir):
    """Test connect() with walks_table for single-file usage."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    walks_df, _ = sample_walks_simple_paths
    
    db_path = _duckdb_with_network_and_walks_table(temp_dir, nodes_path, edges_path, walks_df)
    walkframe = connect(db_path, walks_table='wp_walks')
    assert walkframe is not None
    result = walkframe.compute()
    assert len(result) == len(walks_df)


def test_connect_both_walks_path_and_table_raises(sample_nodes, sample_edges, sample_walks_simple_paths, temp_dir):
    """Test connect() with both walks_path and walks_table raises."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    _, walks_path = sample_walks_simple_paths
    
    network = Network(nodes=nodes_path, edges=edges_path)
    db_path = temp_dir / 'test_both.duckdb'
    network.save(db_path)
    
    with pytest.raises(ValueError, match="exactly one of walks_path or walks_table"):
        connect(db_path, walks_path=str(walks_path), walks_table='wp_walks')


def test_write_duckdb(sample_nodes, sample_edges, sample_walks_simple_paths, temp_dir):
    """Test write_duckdb creates a file that connect() or WalkFrame.from_duckdb can load."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    walks_df, _ = sample_walks_simple_paths

    network = Network(nodes=nodes_path, edges=edges_path)
    out_path = temp_dir / 'out.duckdb'
    write_duckdb(network, walks_df, out_path)
    assert out_path.exists()

    walkframe = connect(out_path, walks_table='wp_walks')
    result = walkframe.compute()
    assert len(result) == len(walks_df)
    assert set(result.columns) >= {'walk_id', 'sequence_idx', 'node_id'}


def test_write_duckdb_custom_table(sample_nodes, sample_edges, sample_walks_simple_paths, temp_dir):
    """Test write_duckdb with custom walks_table name."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    walks_df, _ = sample_walks_simple_paths

    network = Network(nodes=nodes_path, edges=edges_path)
    out_path = temp_dir / 'out_custom.duckdb'
    write_duckdb(network, walks_df, out_path, walks_table='trajectories')
    assert out_path.exists()

    walkframe = connect(out_path, walks_table='trajectories')
    result = walkframe.compute()
    assert len(result) == len(walks_df)


def test_write_duckdb_schema_validate(sample_nodes, sample_edges, temp_dir):
    """Test write_duckdb with schema_validate converts alternative column names."""
    _, nodes_path = sample_nodes
    _, edges_path = sample_edges
    df = pd.DataFrame({'id': ['a', 'a'], 'step': [0, 1], 'node': [0, 1]})

    network = Network(nodes=nodes_path, edges=edges_path)
    out_path = temp_dir / 'out_convert.duckdb'
    write_duckdb(network, df, out_path, schema_validate=True)
    assert out_path.exists()

    walkframe = connect(out_path, walks_table='wp_walks')
    result = walkframe.compute()
    assert 'walk_id' in result.columns
    assert 'sequence_idx' in result.columns
    assert 'node_id' in result.columns
    assert len(result) == 2


def test_validate_schema_non_strict_missing_cols():
    """Test validate_schema with strict=False warns and returns False."""
    df = pd.DataFrame([{'walk_id': 1, 'node_id': 'A'}])  # missing sequence_idx
    with pytest.warns(UserWarning, match="Missing recommended columns"):
        assert validate_schema(df, strict=False) is False


def test_validate_schema_strict_missing_cols():
    """Test validate_schema with strict=True raises on missing columns."""
    df = pd.DataFrame([{'walk_id': 1}])
    with pytest.raises(ValueError, match="Missing required"):
        validate_schema(df, strict=True)


def test_convert_to_superset_format_path_id():
    """Test convert_to_superset_format with path_id column."""
    df = pd.DataFrame({'path_id': [1, 1], 'step': [0, 1], 'node': ['A', 'B']})
    converted = convert_to_superset_format(df)
    assert 'walk_id' in converted.columns
    assert list(converted['walk_id']) == [1, 1]


def test_convert_to_superset_format_extra_cols_become_dynamic():
    """Test convert_to_superset_format adds dynamic_attributes from extra columns."""
    df = pd.DataFrame({
        'id': [1, 1], 'step': [0, 1], 'node': ['A', 'B'],
        'purpose': ['commute', 'commute'], 'speed': [1.0, 1.2]
    })
    converted = convert_to_superset_format(df)
    assert 'dynamic_attributes' in converted.columns


def test_convert_to_superset_format_missing_walk_id_raises():
    """Test convert_to_superset_format raises when walk_id cannot be mapped."""
    df = pd.DataFrame({'unknown_id': [1], 'step': [0], 'node': ['A']})
    with pytest.raises(ValueError, match="walk_id"):
        convert_to_superset_format(df)


def test_write_walks_with_partitioning(temp_dir, sample_walks_simple_paths):
    """Test write_walks with partition_by."""
    df, _ = sample_walks_simple_paths
    df['year'] = 2024
    df['month'] = 1
    output_path = temp_dir / 'partitioned_walks'
    write_walks(df, output_path, partition_by=['year', 'month'])
    assert output_path.exists()
    # Should have created partition dirs
    assert any(output_path.iterdir())


def test_write_walks_without_schema_validate(temp_dir):
    """Test write_walks with schema_validate=False."""
    df = pd.DataFrame({
        'walk_id': [1], 'sequence_idx': [0], 'node_id': ['A'],
        'extra': [1]
    })
    output_path = temp_dir / 'no_validate.parquet'
    write_walks(df, output_path, schema_validate=False)
    assert output_path.exists()


def test_convert_unmappable_sequence_raises():
    """Test convert raises ValueError when sequence_idx cannot be mapped."""
    df = pd.DataFrame({"id": [1], "unmapped_step": [0], "node": ["A"]})
    with pytest.raises(ValueError, match="sequence_idx"):
        convert_to_superset_format(df)


def test_convert_unmappable_node_raises():
    """Test convert raises ValueError when node_id cannot be mapped."""
    df = pd.DataFrame({"id": [1], "step": [0], "unmapped_node": ["A"]})
    with pytest.raises(ValueError, match="node_id"):
        convert_to_superset_format(df)


def test_write_walks_partition_by_string(temp_dir, sample_walks_simple_paths):
    """Test write_walks with partition_by as a single string (not list)."""
    df, _ = sample_walks_simple_paths
    df["year"] = 2024
    output_path = temp_dir / "partitioned_string"
    write_walks(df, output_path, partition_by="year")
    assert output_path.exists()
    assert any(output_path.iterdir())


def test_convert_to_superset_format_sequence_idx_float():
    """Test convert_to_superset_format converts non-integer sequence_idx."""
    df = pd.DataFrame({'id': [1, 1], 'step': [0.0, 1.0], 'node': ['A', 'B']})
    converted = convert_to_superset_format(df)
    assert 'sequence_idx' in converted.columns
    assert converted['sequence_idx'].dtype in (pd.Int64Dtype(), 'int64', 'int32')
