"""
Tests for Dataset Observatory

Comprehensive tests for:
- Entity creation and hashing
- Activity tracking
- Provenance graph operations
- Serialization round-trip
- Schema observation
- Accountability bundle
"""

import json
import time
import pytest
from cascade.data import (
    DatasetObserver,
    ProvenanceGraph,
    DatasetEntity,
    Activity,
    Agent,
    Relationship,
    RelationType,
    ActivityType,
    AgentType,
    SchemaObserver,
    DatasetSchema,
    FieldSchema,
    AccountabilityBundle,
    export_to_croissant,
    create_system_agent,
    create_model_agent,
    create_user_agent,
    hash_content,
)


class TestEntities:
    """Test PROV-O entity classes."""
    
    def test_dataset_entity_creation(self):
        """Test DatasetEntity can be created with all fields."""
        entity = DatasetEntity(
            id="entity:test:001",
            name="test_dataset",
            content_hash="abc123",
            schema_hash="def456",
            version="1.0",
            previous_version="0.9",
            source_type="hf_hub",
            source_uri="hf://squad",
            record_count=1000,
            size_bytes=50000,
            splits={"train": 800, "test": 200},
            attributes={"license": "MIT"},
        )
        
        assert entity.id == "entity:test:001"
        assert entity.name == "test_dataset"
        assert entity.content_hash == "abc123"
        assert entity.record_count == 1000
        assert entity.splits["train"] == 800
    
    def test_dataset_entity_auto_id(self):
        """Test DatasetEntity generates ID if not provided."""
        entity = DatasetEntity(id="", name="auto_id_test")
        assert entity.id.startswith("entity:auto_id_test:")
    
    def test_activity_creation(self):
        """Test Activity can be created and tracks timing."""
        activity = Activity(
            id="activity:test:001",
            activity_type=ActivityType.TRANSFORM,
            name="test_transform",
        )
        
        assert activity.id == "activity:test:001"
        assert activity.activity_type == ActivityType.TRANSFORM
        assert activity.started_at is not None
        assert activity.ended_at is None
        
        # End activity
        activity.end()
        assert activity.ended_at is not None
        assert activity.duration >= 0
    
    def test_activity_inputs_outputs(self):
        """Test Activity tracks inputs and outputs."""
        activity = Activity(
            id="activity:test:002",
            activity_type=ActivityType.INGEST,
            name="test_ingest",
        )
        
        activity.add_input("entity:a")
        activity.add_input("entity:b")
        activity.add_input("entity:a")  # Duplicate
        activity.add_output("entity:c")
        
        assert len(activity.inputs) == 2  # No duplicate
        assert "entity:a" in activity.inputs
        assert "entity:b" in activity.inputs
        assert len(activity.outputs) == 1
    
    def test_agent_creation(self):
        """Test Agent types and factory functions."""
        system = create_system_agent("cascade")
        assert system.agent_type == AgentType.SYSTEM
        assert system.id == "agent:system:cascade"
        
        model = create_model_agent("bert-base-uncased", version="1.0")
        assert model.agent_type == AgentType.MODEL
        assert model.version == "1.0"
        
        user = create_user_agent("researcher", org="cascade-lab")
        assert user.agent_type == AgentType.PERSON
        assert user.parent_agent_id == "agent:organization:cascade-lab"
    
    def test_relationship_creation(self):
        """Test Relationship tracking."""
        rel = Relationship(
            relation_type=RelationType.WAS_DERIVED_FROM,
            source_id="entity:derived",
            target_id="entity:source",
        )
        
        assert rel.relation_type == RelationType.WAS_DERIVED_FROM
        assert rel.source_id == "entity:derived"
        assert rel.target_id == "entity:source"
        assert rel.timestamp > 0


class TestProvenanceGraph:
    """Test provenance graph operations."""
    
    def test_graph_creation(self):
        """Test ProvenanceGraph initializes correctly."""
        graph = ProvenanceGraph(name="test_graph")
        
        assert graph.name == "test_graph"
        assert len(graph.list_entities()) == 0
        assert len(graph.list_activities()) == 0
        assert len(graph.list_agents()) == 1  # System agent
    
    def test_add_entity(self):
        """Test adding entities to graph."""
        graph = ProvenanceGraph()
        entity = DatasetEntity(id="entity:test", name="test")
        
        graph.add_entity(entity)
        
        assert graph.get_entity("entity:test") is not None
        assert len(graph.list_entities()) == 1
    
    def test_lineage_tracking(self):
        """Test lineage queries."""
        graph = ProvenanceGraph()
        
        # Create chain: A -> B -> C
        entity_a = DatasetEntity(id="entity:a", name="source")
        entity_b = DatasetEntity(id="entity:b", name="derived1")
        entity_c = DatasetEntity(id="entity:c", name="derived2")
        
        graph.add_entity(entity_a)
        graph.add_entity(entity_b)
        graph.add_entity(entity_c)
        
        graph.link_derivation("entity:b", "entity:a")
        graph.link_derivation("entity:c", "entity:b")
        
        # Trace upstream
        upstream = graph.get_lineage("entity:c", "upstream")
        assert "entity:b" in upstream
        assert "entity:a" in upstream
        
        # Trace downstream
        downstream = graph.get_lineage("entity:a", "downstream")
        assert "entity:b" in downstream
        assert "entity:c" in downstream
    
    def test_root_hash(self):
        """Test Merkle root hash updates."""
        graph = ProvenanceGraph()
        initial_hash = graph.root_hash
        
        entity = DatasetEntity(id="entity:test", name="test")
        graph.add_entity(entity)
        
        # Hash should change
        assert graph.root_hash != initial_hash
    
    def test_integrity_verification(self):
        """Test integrity verification."""
        graph = ProvenanceGraph()
        entity = DatasetEntity(id="entity:test", name="test")
        graph.add_entity(entity)
        
        is_valid, invalid = graph.verify_integrity()
        assert is_valid
        assert len(invalid) == 0
    
    def test_serialization_roundtrip(self):
        """Test to_dict/from_dict preserves all data."""
        graph = ProvenanceGraph(name="roundtrip_test")
        
        # Add entity with all fields
        entity = DatasetEntity(
            id="entity:full",
            name="full_entity",
            content_hash="hash123",
            schema_hash="schema456",
            version="1.0",
            previous_version="0.9",
            source_type="hf_hub",
            source_uri="hf://test",
            record_count=100,
            size_bytes=5000,
            splits={"train": 80, "test": 20},
            attributes={"custom": "value"},
        )
        graph.add_entity(entity)
        
        # Add activity with all fields
        activity = Activity(
            id="activity:full",
            activity_type=ActivityType.TRANSFORM,
            name="full_activity",
            parameters={"param1": "value1"},
            attributes={"custom": "attr"},
        )
        activity.add_input(entity.id)
        activity.end()
        graph.add_activity(activity)
        
        # Add custom agent
        agent = create_user_agent("testuser", org="testorg")
        graph.add_agent(agent)
        
        # Add relationship
        graph.add_relationship(
            RelationType.WAS_ATTRIBUTED_TO,
            entity.id,
            agent.id,
            attributes={"role": "creator"},
        )
        
        # Serialize
        data = graph.to_dict()
        
        # Deserialize
        restored = ProvenanceGraph.from_dict(data)
        
        # Verify
        assert restored.name == "roundtrip_test"
        
        restored_entity = restored.get_entity("entity:full")
        assert restored_entity is not None
        assert restored_entity.content_hash == "hash123"
        assert restored_entity.size_bytes == 5000
        assert restored_entity.attributes.get("custom") == "value"
        
        restored_activity = restored.get_activity("activity:full")
        assert restored_activity is not None
        assert restored_activity.parameters.get("param1") == "value1"
        assert restored_activity.attributes.get("custom") == "attr"
        
        restored_agent = restored.get_agent(agent.id)
        assert restored_agent is not None
        assert restored_agent.parent_agent_id == "agent:organization:testorg"


class TestDatasetObserver:
    """Test DatasetObserver operations."""
    
    def test_observer_creation(self):
        """Test observer initializes correctly."""
        observer = DatasetObserver(name="test")
        
        assert observer.graph.name == "test"
        assert len(observer.graph.list_agents()) == 1  # Uses graph's system agent
    
    def test_observe_dict_dataset(self):
        """Test observing a dict dataset."""
        observer = DatasetObserver()
        
        data = {"col1": [1, 2, 3], "col2": ["a", "b", "c"]}
        entity = observer.observe_dataset(data, name="dict_data")
        
        assert entity.name == "dict_data"
        assert entity.content_hash is not None
        # For columnar dict, record_count is sum of all column lengths
        # (this is how the observer treats dict data - as splits)
        assert entity.record_count == 6  # 3 + 3
        assert entity.splits == {"col1": 3, "col2": 3}
    
    def test_observe_ingest_context(self):
        """Test ingest context manager."""
        observer = DatasetObserver()
        
        with observer.observe_ingest("load_data") as ctx:
            data = {"x": [1, 2], "y": [3, 4]}
            entity = ctx.output(data, name="loaded")
        
        # Activity should be recorded
        activities = observer.graph.list_activities()
        assert len(activities) == 1
        assert activities[0].activity_type == ActivityType.INGEST
        assert activities[0].duration >= 0
        
        # Entity should be linked
        assert entity.id in activities[0].outputs
    
    def test_observe_transform_with_input(self):
        """Test transform context with input tracking."""
        observer = DatasetObserver()
        
        # First, ingest
        with observer.observe_ingest("load") as ctx:
            input_data = {"x": [1, 2, 3]}
            input_entity = ctx.output(input_data, name="input")
        
        # Then transform
        with observer.observe_transform("filter") as ctx:
            ctx.input(input_entity.id)  # Reference by ID
            output_data = {"x": [1, 2]}
            output_entity = ctx.output(output_data, name="filtered")
        
        # Check lineage
        lineage = observer.graph.get_lineage(output_entity.id, "upstream")
        assert input_entity.id in lineage
    
    def test_observe_entity_resolution(self):
        """Test entity resolution observation."""
        observer = DatasetObserver()
        
        with observer.observe_entity_resolution(
            "match_records",
            model_id="all-MiniLM-L6-v2",
            threshold=0.8,
        ) as ctx:
            data1 = {"id": [1, 2]}
            data2 = {"id": [3, 4]}
            ctx.input(observer.observe_dataset(data1, name="source1"))
            ctx.input(observer.observe_dataset(data2, name="source2"))
            unified = {"matched_id": [1, 3]}
            ctx.output(unified, name="unified")
        
        activities = observer.graph.list_activities()
        er_activity = [a for a in activities if a.activity_type == ActivityType.ENTITY_RESOLUTION]
        assert len(er_activity) == 1
        assert er_activity[0].parameters.get("threshold") == 0.8


class TestSchemaObserver:
    """Test schema observation."""
    
    def test_observe_dict(self):
        """Test observing dict schema."""
        observer = SchemaObserver()
        
        data = {
            "name": ["Alice", "Bob"],
            "age": [25, 30],
            "score": [0.8, 0.9],
        }
        
        schema = observer.observe_dict(data)
        
        assert schema.num_fields == 3
        assert schema.fields["name"].dtype == "string"
        assert schema.fields["age"].dtype == "int64"
        assert schema.fields["score"].dtype == "float64"
    
    def test_schema_hash_stability(self):
        """Test schema hash is stable."""
        observer = SchemaObserver()
        
        data = {"a": [1, 2], "b": ["x", "y"]}
        
        schema1 = observer.observe_dict(data)
        schema2 = observer.observe_dict(data)
        
        assert schema1.hash() == schema2.hash()
    
    def test_schema_diff(self):
        """Test schema comparison."""
        schema1 = DatasetSchema()
        schema1.add_field(FieldSchema(name="col1", dtype="string"))
        schema1.add_field(FieldSchema(name="col2", dtype="int64"))
        
        schema2 = DatasetSchema()
        schema2.add_field(FieldSchema(name="col1", dtype="string"))
        schema2.add_field(FieldSchema(name="col3", dtype="float64"))
        
        diff = schema1.diff(schema2)
        
        assert "col2" in diff["removed"]
        assert "col3" in diff["added"]
        assert diff["compatible"] is False  # col2 was removed


class TestAccountabilityBundle:
    """Test accountability bundle generation."""
    
    def test_bundle_creation(self):
        """Test bundle can be created."""
        graph = ProvenanceGraph(name="test")
        bundle = AccountabilityBundle(graph)
        
        summary = bundle.summary()
        assert summary["graph_name"] == "test"
        assert len(summary["files_included"]) == 7
    
    def test_activity_log(self):
        """Test activity log export."""
        observer = DatasetObserver()
        
        with observer.observe_ingest("load") as ctx:
            ctx.output({"x": [1]}, name="data")
        
        bundle = AccountabilityBundle(observer.graph)
        log = bundle.to_activity_log()
        
        assert len(log) == 1
        assert log[0]["type"] == "ingest"
        assert log[0]["started_at"] is not None
    
    def test_agent_attributions(self):
        """Test agent attribution export."""
        observer = DatasetObserver()
        
        with observer.observe_ingest("load") as ctx:
            ctx.output({"x": [1]}, name="data")
        
        bundle = AccountabilityBundle(observer.graph)
        agents = bundle.to_agent_attributions()
        
        assert agents["total_agents"] >= 1
        assert agents["total_attributions"] >= 1
    
    def test_integrity_manifest(self):
        """Test integrity manifest export."""
        observer = DatasetObserver()
        
        with observer.observe_ingest("load") as ctx:
            ctx.output({"x": [1]}, name="data")
        
        bundle = AccountabilityBundle(observer.graph)
        manifest = bundle.to_integrity_manifest()
        
        assert manifest["is_valid"] is True
        assert manifest["root_hash"] is not None


class TestCroissantExport:
    """Test Croissant export."""
    
    def test_basic_export(self):
        """Test basic Croissant export."""
        observer = DatasetObserver()
        
        with observer.observe_ingest("load") as ctx:
            ctx.output({"text": ["hello", "world"]}, name="test_data")
        
        croissant = export_to_croissant(
            observer.graph,
            name="test_dataset",
            description="A test dataset",
        )
        
        assert croissant["name"] == "test_dataset"
        assert croissant["description"] == "A test dataset"
        assert "@context" in croissant
        assert "distribution" in croissant
    
    def test_provenance_extension(self):
        """Test Croissant includes provenance extension."""
        observer = DatasetObserver()
        
        with observer.observe_ingest("load") as ctx:
            ctx.output({"x": [1]}, name="data")
        
        croissant = export_to_croissant(observer.graph, name="test")
        
        assert "cr:provenance" in croissant
        assert croissant["cr:provenance"]["cascade:rootHash"] is not None


class TestHashContent:
    """Test content hashing."""
    
    def test_dict_hash(self):
        """Test hashing dict content."""
        data = {"a": [1, 2, 3], "b": ["x", "y", "z"]}
        
        hash1 = hash_content(data)
        hash2 = hash_content(data)
        
        assert hash1 == hash2
    
    def test_different_data_different_hash(self):
        """Test different data produces different hash."""
        data1 = {"a": [1, 2, 3]}
        data2 = {"a": [1, 2, 4]}
        
        assert hash_content(data1) != hash_content(data2)


if __name__ == "__main__":
    # Run with: python -m pytest tests/test_data_observatory.py -v
    pytest.main([__file__, "-v"])
