"""
Test storage layer functionality.
"""

import pytest
from pathlib import Path
import pandas as pd
from uuid import uuid4

from app.storage.file_storage import FileStorage
from app.storage.metadata_store import MetadataStore
from app.storage.lineage import LineageTracker, LineageNode, LineageEdge
from app.domain.models.dataset import Dataset, DatasetStatus
from app.domain.models.pipeline import Pipeline, PipelineStatus


class TestFileStorage:
    """Tests for FileStorage."""

    def test_read_csv(self, file_storage, sample_csv_file):
        """Test reading CSV file."""
        df = file_storage.read_dataset(sample_csv_file, "csv", engine="pandas")
        
        assert isinstance(df, pd.DataFrame)
        assert len(df) == 6
        assert "customer_id" in df.columns

    def test_write_csv(self, file_storage, sample_dataframe, tmp_path):
        """Test writing CSV file."""
        output_path = tmp_path / "output.csv"
        file_storage.write_dataset(sample_dataframe, output_path, "csv")
        
        assert output_path.exists()
        
        # Read back and verify
        df = pd.read_csv(output_path)
        assert len(df) == len(sample_dataframe)

    def test_read_parquet(self, file_storage, sample_dataframe, tmp_path):
        """Test Parquet read/write."""
        parquet_path = tmp_path / "test.parquet"
        sample_dataframe.to_parquet(parquet_path)
        
        df = file_storage.read_dataset(parquet_path, "parquet", engine="pandas")
        assert len(df) == len(sample_dataframe)

    def test_engine_selection(self, file_storage, tmp_path):
        """Test automatic engine selection based on size."""
        # Small data should use pandas
        small_df = pd.DataFrame({"a": range(100)})
        small_path = tmp_path / "small.csv"
        small_df.to_csv(small_path, index=False)
        
        result = file_storage.read_dataset(small_path, "csv", engine="auto")
        assert isinstance(result, pd.DataFrame)


class TestMetadataStore:
    """Tests for MetadataStore."""

    def test_save_and_get_dataset(self, metadata_store):
        """Test saving and retrieving dataset."""
        dataset = Dataset(
            name="Test Dataset",
            original_filename="test.csv",
            file_path=Path("/tmp/test.csv"),
            file_format="csv",
            file_size_bytes=1024,
            row_count=100,
            column_count=5,
        )
        
        metadata_store.save_dataset(dataset)
        
        retrieved = metadata_store.get_dataset(dataset.id)
        assert retrieved is not None
        assert retrieved.name == "Test Dataset"
        assert retrieved.row_count == 100

    def test_list_datasets(self, metadata_store):
        """Test listing datasets."""
        # Create multiple datasets
        for i in range(3):
            dataset = Dataset(
                name=f"Dataset {i}",
                original_filename=f"file{i}.csv",
                file_path=Path(f"/tmp/file{i}.csv"),
                file_format="csv",
                file_size_bytes=1024 * i,
            )
            metadata_store.save_dataset(dataset)
        
        datasets = metadata_store.list_datasets(limit=10)
        assert len(datasets) >= 3

    def test_delete_dataset(self, metadata_store):
        """Test deleting dataset."""
        dataset = Dataset(
            name="To Delete",
            original_filename="delete.csv",
            file_path=Path("/tmp/delete.csv"),
            file_format="csv",
            file_size_bytes=1024,
        )
        
        metadata_store.save_dataset(dataset)
        dataset_id = dataset.id
        
        metadata_store.delete_dataset(dataset_id)
        
        retrieved = metadata_store.get_dataset(dataset_id)
        assert retrieved is None

    def test_save_and_get_pipeline(self, metadata_store):
        """Test saving and retrieving pipeline."""
        from app.domain.models.ir import IRSchema
        
        pipeline = Pipeline(
            name="Test Pipeline",
            source_dataset_id=uuid4(),
            natural_language_input="Test transformation",
            ir=IRSchema(
                source={"dataset_id": uuid4(), "schema": None},
                dag={"nodes": [], "execution_order": []},
            ),
        )
        
        metadata_store.save_pipeline(pipeline)
        
        retrieved = metadata_store.get_pipeline(pipeline.id)
        assert retrieved is not None
        assert retrieved.name == "Test Pipeline"

    def test_update_pipeline_status(self, metadata_store):
        """Test updating pipeline status."""
        from app.domain.models.ir import IRSchema
        
        pipeline = Pipeline(
            name="Status Test",
            source_dataset_id=uuid4(),
            natural_language_input="Test",
            ir=IRSchema(
                source={"dataset_id": uuid4(), "schema": None},
                dag={"nodes": [], "execution_order": []},
            ),
        )
        
        metadata_store.save_pipeline(pipeline)
        
        pipeline.update_status(PipelineStatus.EXECUTING)
        metadata_store.save_pipeline(pipeline)
        
        retrieved = metadata_store.get_pipeline(pipeline.id)
        assert retrieved.status == PipelineStatus.EXECUTING


class TestLineageTracker:
    """Tests for LineageTracker."""

    def test_add_and_get_lineage(self):
        """Test adding and retrieving lineage."""
        tracker = LineageTracker()
        
        source_id = uuid4()
        target_id = uuid4()
        pipeline_id = uuid4()
        
        tracker.add_lineage(
            source_dataset_id=source_id,
            target_dataset_id=target_id,
            pipeline_id=pipeline_id,
            transformation_type="cleaning",
        )
        
        ancestors = tracker.get_ancestors(target_id)
        assert len(ancestors) == 1
        assert ancestors[0].dataset_id == source_id

    def test_lineage_chain(self):
        """Test lineage chain tracking."""
        tracker = LineageTracker()
        
        ds1 = uuid4()
        ds2 = uuid4()
        ds3 = uuid4()
        
        # ds1 -> ds2 -> ds3
        tracker.add_lineage(ds1, ds2, uuid4(), "step1")
        tracker.add_lineage(ds2, ds3, uuid4(), "step2")
        
        # Get full ancestry
        ancestors = tracker.get_ancestors(ds3)
        assert len(ancestors) == 2  # ds2 and ds1
        
        # Get descendants
        descendants = tracker.get_descendants(ds1)
        assert len(descendants) == 2  # ds2 and ds3

    def test_mermaid_generation(self):
        """Test Mermaid diagram generation."""
        tracker = LineageTracker()
        
        ds1 = uuid4()
        ds2 = uuid4()
        
        tracker.add_lineage(ds1, ds2, uuid4(), "transform")
        
        mermaid = tracker.to_mermaid()
        assert "graph TD" in mermaid
        assert "-->" in mermaid
