# test_job_tracking.py
"""Tests for job tracking module.

This module tests:
- JobManifest: Job state persistence and recovery detection
- JobLogger: Structured logging for ingestion jobs
"""

import json
import tempfile
from pathlib import Path

import pytest

from earthcatalog.job_tracking import (
    ConsolidationPhaseState,
    DownloadPhaseState,
    JobLogger,
    JobManifest,
    JobStatus,
)
from earthcatalog.storage_backends import LocalStorage


class TestJobStatus:
    """Test JobStatus enum."""

    def test_status_values(self):
        """JobStatus should have expected values."""
        assert JobStatus.PENDING == "pending"
        assert JobStatus.DOWNLOADING == "downloading"
        assert JobStatus.CONSOLIDATING == "consolidating"
        assert JobStatus.COMPLETED == "completed"
        assert JobStatus.FAILED == "failed"


class TestDownloadPhaseState:
    """Test DownloadPhaseState dataclass."""

    def test_default_values(self):
        """DownloadPhaseState should have sensible defaults."""
        state = DownloadPhaseState()
        assert state.completed is False
        assert state.batches_total == 0
        assert state.batches_completed == 0
        assert state.urls_processed == 0
        assert state.urls_failed == 0
        assert state.shards_written == []

    def test_to_dict_from_dict_roundtrip(self):
        """DownloadPhaseState should serialize/deserialize correctly."""
        state = DownloadPhaseState(
            completed=True,
            batches_total=10,
            batches_completed=8,
            urls_processed=1000,
            urls_failed=5,
            shards_written=["shard1.parquet", "shard2.parquet"],
        )

        data = state.to_dict()
        restored = DownloadPhaseState.from_dict(data)

        assert restored.completed == state.completed
        assert restored.batches_total == state.batches_total
        assert restored.batches_completed == state.batches_completed
        assert restored.urls_processed == state.urls_processed
        assert restored.urls_failed == state.urls_failed
        assert restored.shards_written == state.shards_written


class TestConsolidationPhaseState:
    """Test ConsolidationPhaseState dataclass."""

    def test_default_values(self):
        """ConsolidationPhaseState should have sensible defaults."""
        state = ConsolidationPhaseState()
        assert state.completed is False
        assert state.partitions_total == 0
        assert state.partitions_completed == 0
        assert state.completed_partitions == []

    def test_to_dict_from_dict_roundtrip(self):
        """ConsolidationPhaseState should serialize/deserialize correctly."""
        state = ConsolidationPhaseState(
            completed=True,
            partitions_total=5,
            partitions_completed=5,
            completed_partitions=["p1", "p2", "p3", "p4", "p5"],
        )

        data = state.to_dict()
        restored = ConsolidationPhaseState.from_dict(data)

        assert restored.completed == state.completed
        assert restored.partitions_total == state.partitions_total
        assert restored.partitions_completed == state.partitions_completed
        assert restored.completed_partitions == state.completed_partitions


class TestJobManifest:
    """Test JobManifest dataclass."""

    @pytest.fixture
    def temp_dir(self):
        """Create a temporary directory for testing."""
        with tempfile.TemporaryDirectory() as tmpdir:
            yield tmpdir

    @pytest.fixture
    def storage(self, temp_dir):
        """Create a LocalStorage instance."""
        return LocalStorage(temp_dir)

    def test_create_new_manifest(self):
        """Creating a new JobManifest should set defaults correctly."""
        manifest = JobManifest(
            job_id="test-job-123",
            input_urls_count=1000,
        )

        assert manifest.job_id == "test-job-123"
        assert manifest.status == JobStatus.PENDING
        assert manifest.input_urls_count == 1000
        assert manifest.created_at is not None
        assert manifest.download_phase is not None
        assert manifest.consolidation_phase is not None

    def test_to_dict_from_dict_roundtrip(self):
        """JobManifest should serialize/deserialize correctly."""
        manifest = JobManifest(
            job_id="test-job-456",
            input_urls_count=5000,
            config_hash="abc123",
        )
        manifest.status = JobStatus.DOWNLOADING
        manifest.download_phase.batches_total = 10
        manifest.download_phase.batches_completed = 5

        data = manifest.to_dict()
        restored = JobManifest.from_dict(data)

        assert restored.job_id == manifest.job_id
        assert restored.status == manifest.status
        assert restored.input_urls_count == manifest.input_urls_count
        assert restored.config_hash == manifest.config_hash
        assert restored.download_phase.batches_total == 10
        assert restored.download_phase.batches_completed == 5

    def test_save_creates_manifest_file(self, storage, temp_dir):
        """save() should create manifest.json in the correct location."""
        manifest = JobManifest(
            job_id="save-test-job",
            input_urls_count=100,
        )

        manifest.save(storage, temp_dir)

        expected_path = Path(temp_dir) / "jobs" / "save-test-job" / "manifest.json"
        assert expected_path.exists()

        # Verify content
        with open(expected_path) as f:
            data = json.load(f)
        assert data["job_id"] == "save-test-job"

    def test_load_reads_manifest_file(self, storage, temp_dir):
        """load() should read an existing manifest."""
        # Create and save a manifest
        original = JobManifest(
            job_id="load-test-job",
            input_urls_count=200,
        )
        original.status = JobStatus.CONSOLIDATING
        original.save(storage, temp_dir)

        # Load it back
        loaded = JobManifest.load(storage, temp_dir, "load-test-job")

        assert loaded.job_id == "load-test-job"
        assert loaded.status == JobStatus.CONSOLIDATING
        assert loaded.input_urls_count == 200

    def test_load_raises_for_missing_manifest(self, storage, temp_dir):
        """load() should raise FileNotFoundError for missing manifests."""
        with pytest.raises(FileNotFoundError):
            JobManifest.load(storage, temp_dir, "nonexistent-job")

    def test_find_incomplete_returns_none_when_no_jobs(self, storage, temp_dir):
        """find_incomplete() should return None when no jobs directory exists."""
        result = JobManifest.find_incomplete(storage, temp_dir)
        assert result is None

    def test_find_incomplete_returns_none_when_all_complete(self, storage, temp_dir):
        """find_incomplete() should return None when all jobs are complete."""
        # Create a completed job
        completed = JobManifest(
            job_id="completed-job",
            input_urls_count=100,
        )
        completed.status = JobStatus.COMPLETED
        completed.save(storage, temp_dir)

        result = JobManifest.find_incomplete(storage, temp_dir)
        assert result is None

    def test_find_incomplete_finds_downloading_job(self, storage, temp_dir):
        """find_incomplete() should find jobs in DOWNLOADING status."""
        incomplete = JobManifest(
            job_id="incomplete-job",
            input_urls_count=100,
        )
        incomplete.status = JobStatus.DOWNLOADING
        incomplete.save(storage, temp_dir)

        result = JobManifest.find_incomplete(storage, temp_dir)
        assert result is not None
        assert result.job_id == "incomplete-job"

    def test_find_incomplete_finds_consolidating_job(self, storage, temp_dir):
        """find_incomplete() should find jobs in CONSOLIDATING status."""
        incomplete = JobManifest(
            job_id="consolidating-job",
            input_urls_count=100,
        )
        incomplete.status = JobStatus.CONSOLIDATING
        incomplete.save(storage, temp_dir)

        result = JobManifest.find_incomplete(storage, temp_dir)
        assert result is not None
        assert result.job_id == "consolidating-job"

    def test_manifest_path_property(self, temp_dir):
        """manifest_path() should return correct path."""
        manifest = JobManifest(job_id="path-test", input_urls_count=0)
        path = manifest.manifest_path(temp_dir)
        assert path == f"{temp_dir}/jobs/path-test/manifest.json"


class TestJobLogger:
    """Test JobLogger class."""

    @pytest.fixture
    def temp_dir(self):
        """Create a temporary directory for testing."""
        with tempfile.TemporaryDirectory() as tmpdir:
            yield tmpdir

    @pytest.fixture
    def storage(self, temp_dir):
        """Create a LocalStorage instance."""
        return LocalStorage(temp_dir)

    def test_logger_creates_log_file(self, storage, temp_dir):
        """JobLogger should create log file on first log."""
        logger = JobLogger(storage, temp_dir, "test-job")
        logger.log("INFO", "Test message")

        # Check that logs directory exists
        logs_dir = Path(temp_dir) / "jobs" / "logs"
        assert logs_dir.exists()

        # Check that a log file was created
        log_files = list(logs_dir.glob("*.txt"))
        assert len(log_files) == 1

    def test_log_writes_message(self, storage, temp_dir):
        """log() should write message to file."""
        logger = JobLogger(storage, temp_dir, "test-job")
        logger.log("INFO", "Hello world")

        # Read the log file
        logs_dir = Path(temp_dir) / "jobs" / "logs"
        log_file = next(logs_dir.glob("*.txt"))
        content = log_file.read_text()

        assert "Hello world" in content
        assert "INFO" in content
        assert "test-job" in content

    def test_log_with_context(self, storage, temp_dir):
        """log() should include context in message."""
        logger = JobLogger(storage, temp_dir, "test-job")
        logger.log("WARNING", "Something happened", url="http://example.com", count=42)

        logs_dir = Path(temp_dir) / "jobs" / "logs"
        log_file = next(logs_dir.glob("*.txt"))
        content = log_file.read_text()

        assert "Something happened" in content
        assert "url" in content
        assert "http://example.com" in content

    def test_log_phase_start(self, storage, temp_dir):
        """log_phase_start() should log phase beginning."""
        logger = JobLogger(storage, temp_dir, "test-job")
        logger.log_phase_start("download")

        logs_dir = Path(temp_dir) / "jobs" / "logs"
        log_file = next(logs_dir.glob("*.txt"))
        content = log_file.read_text()

        assert "download" in content.lower()
        assert "start" in content.lower() or "begin" in content.lower()

    def test_log_phase_complete(self, storage, temp_dir):
        """log_phase_complete() should log phase completion with stats."""
        logger = JobLogger(storage, temp_dir, "test-job")
        logger.log_phase_complete("download", {"urls_processed": 1000, "shards": 10})

        logs_dir = Path(temp_dir) / "jobs" / "logs"
        log_file = next(logs_dir.glob("*.txt"))
        content = log_file.read_text()

        assert "download" in content.lower()
        assert "1000" in content or "urls_processed" in content

    def test_log_error(self, storage, temp_dir):
        """log_error() should log error with context."""
        logger = JobLogger(storage, temp_dir, "test-job")
        logger.log_error("Something failed", url="http://bad.com")

        logs_dir = Path(temp_dir) / "jobs" / "logs"
        log_file = next(logs_dir.glob("*.txt"))
        content = log_file.read_text()

        assert "ERROR" in content
        assert "Something failed" in content

    def test_log_appends_to_existing_file(self, storage, temp_dir):
        """Multiple log() calls should append to same file."""
        logger = JobLogger(storage, temp_dir, "test-job")
        logger.log("INFO", "First message")
        logger.log("INFO", "Second message")

        logs_dir = Path(temp_dir) / "jobs" / "logs"
        log_file = next(logs_dir.glob("*.txt"))
        content = log_file.read_text()

        assert "First message" in content
        assert "Second message" in content


class TestPipelineResumeIntegration:
    """Test pipeline resume functionality with job tracking."""

    @pytest.fixture
    def temp_dir(self):
        """Create a temporary directory for testing."""
        with tempfile.TemporaryDirectory() as tmpdir:
            yield tmpdir

    @pytest.fixture
    def storage(self, temp_dir):
        """Create a LocalStorage instance."""
        return LocalStorage(temp_dir)

    def test_run_creates_job_manifest(self, temp_dir):
        """run() should create a job manifest."""
        from unittest.mock import patch

        import pandas as pd

        from earthcatalog.ingestion_pipeline import (
            LocalProcessor,
            ProcessingConfig,
            STACIngestionPipeline,
        )

        # Create config
        input_file = Path(temp_dir) / "input.parquet"
        output_catalog = Path(temp_dir) / "catalog"
        scratch_location = Path(temp_dir) / "scratch"

        df = pd.DataFrame({"url": ["http://example.com/item1.json"]})
        df.to_parquet(input_file, index=False)

        config = ProcessingConfig(
            input_file=str(input_file),
            output_catalog=str(output_catalog),
            scratch_location=str(scratch_location),
            max_workers=1,
        )

        processor = LocalProcessor(n_workers=1)
        pipeline = STACIngestionPipeline(config, processor)

        # Mock the download to return nothing (empty processing)
        with patch.object(pipeline, "_download_stac_item", return_value=None):
            pipeline.run(job_id="test-manifest-job")

        # Check that manifest was created
        manifest_path = output_catalog / "jobs" / "test-manifest-job" / "manifest.json"
        assert manifest_path.exists()

        # Load and verify
        storage = LocalStorage(str(output_catalog))
        manifest = JobManifest.load(storage, str(output_catalog), "test-manifest-job")
        assert manifest.status == JobStatus.COMPLETED

        processor.close()

    def test_run_marks_job_failed_on_exception(self, temp_dir):
        """run() should mark job as FAILED when exception occurs."""
        from unittest.mock import patch

        import pandas as pd

        from earthcatalog.ingestion_pipeline import (
            LocalProcessor,
            ProcessingConfig,
            STACIngestionPipeline,
        )

        # Create config
        input_file = Path(temp_dir) / "input.parquet"
        output_catalog = Path(temp_dir) / "catalog"
        scratch_location = Path(temp_dir) / "scratch"

        df = pd.DataFrame({"url": ["http://example.com/item1.json"]})
        df.to_parquet(input_file, index=False)

        config = ProcessingConfig(
            input_file=str(input_file),
            output_catalog=str(output_catalog),
            scratch_location=str(scratch_location),
            max_workers=1,
            distributed=True,  # Force distributed mode to test _process_urls_distributed
        )

        processor = LocalProcessor(n_workers=1)
        pipeline = STACIngestionPipeline(config, processor)

        # Mock to raise an exception
        with patch.object(pipeline, "_process_urls_distributed", side_effect=RuntimeError("Test error")):
            with pytest.raises(RuntimeError, match="Test error"):
                pipeline.run(job_id="failed-job")

        # Check that manifest shows FAILED
        storage = LocalStorage(str(output_catalog))
        manifest = JobManifest.load(storage, str(output_catalog), "failed-job")
        assert manifest.status == JobStatus.FAILED
        assert "Test error" in manifest.error

        processor.close()

    def test_resume_raises_when_no_incomplete_job(self, temp_dir):
        """run(resume=True) should raise ValueError if no incomplete job exists."""
        import pandas as pd

        from earthcatalog.ingestion_pipeline import (
            LocalProcessor,
            ProcessingConfig,
            STACIngestionPipeline,
        )

        input_file = Path(temp_dir) / "input.parquet"
        output_catalog = Path(temp_dir) / "catalog"
        scratch_location = Path(temp_dir) / "scratch"

        df = pd.DataFrame({"url": ["http://example.com/item.json"]})
        df.to_parquet(input_file, index=False)

        config = ProcessingConfig(
            input_file=str(input_file),
            output_catalog=str(output_catalog),
            scratch_location=str(scratch_location),
        )

        processor = LocalProcessor(n_workers=1)
        pipeline = STACIngestionPipeline(config, processor)

        with pytest.raises(ValueError, match="No incomplete job found"):
            pipeline.run(resume=True)

        processor.close()

    def test_resume_finds_and_uses_incomplete_job(self, temp_dir):
        """run(resume=True) should find and continue an incomplete job."""
        from unittest.mock import patch

        import pandas as pd

        from earthcatalog.ingestion_pipeline import (
            LocalProcessor,
            ProcessingConfig,
            STACIngestionPipeline,
        )

        # Create config
        input_file = Path(temp_dir) / "input.parquet"
        output_catalog = Path(temp_dir) / "catalog"
        scratch_location = Path(temp_dir) / "scratch"

        df = pd.DataFrame({"url": ["http://example.com/item1.json"]})
        df.to_parquet(input_file, index=False)

        config = ProcessingConfig(
            input_file=str(input_file),
            output_catalog=str(output_catalog),
            scratch_location=str(scratch_location),
            max_workers=1,
        )

        # Create an incomplete job manifest manually
        storage = LocalStorage(str(output_catalog))
        incomplete_manifest = JobManifest(
            job_id="incomplete-job-to-resume",
            input_urls_count=1,
        )
        incomplete_manifest.status = JobStatus.DOWNLOADING
        incomplete_manifest.download_phase.completed = False
        incomplete_manifest.save(storage, str(output_catalog))

        processor = LocalProcessor(n_workers=1)
        pipeline = STACIngestionPipeline(config, processor)

        # Mock the download to return nothing
        with patch.object(pipeline, "_download_stac_item", return_value=None):
            pipeline.run(resume=True)

        # Check that the same job was resumed and completed
        manifest = JobManifest.load(storage, str(output_catalog), "incomplete-job-to-resume")
        assert manifest.status == JobStatus.COMPLETED

        processor.close()

    def test_resume_skips_completed_download_phase(self, temp_dir):
        """run(resume=True) should skip download phase if already completed."""
        from unittest.mock import MagicMock, patch

        import pandas as pd

        from earthcatalog.ingestion_pipeline import (
            LocalProcessor,
            ProcessingConfig,
            STACIngestionPipeline,
        )

        input_file = Path(temp_dir) / "input.parquet"
        output_catalog = Path(temp_dir) / "catalog"
        scratch_location = Path(temp_dir) / "scratch"

        df = pd.DataFrame({"url": ["http://example.com/item1.json"]})
        df.to_parquet(input_file, index=False)

        config = ProcessingConfig(
            input_file=str(input_file),
            output_catalog=str(output_catalog),
            scratch_location=str(scratch_location),
            max_workers=1,
        )

        # Create manifest with download phase completed
        storage = LocalStorage(str(output_catalog))
        manifest = JobManifest(
            job_id="download-complete-job",
            input_urls_count=1,
        )
        manifest.status = JobStatus.CONSOLIDATING
        manifest.download_phase.completed = True
        manifest.download_phase.shards_written = []  # No shards
        manifest.save(storage, str(output_catalog))

        processor = LocalProcessor(n_workers=1)
        pipeline = STACIngestionPipeline(config, processor)

        # Mock _process_urls_distributed to track if it's called
        mock_process = MagicMock()
        with patch.object(pipeline, "_process_urls_distributed", mock_process):
            pipeline.run(resume=True)

        # Download phase should NOT have been called (it was already complete)
        mock_process.assert_not_called()

        # Job should be completed
        manifest = JobManifest.load(storage, str(output_catalog), "download-complete-job")
        assert manifest.status == JobStatus.COMPLETED

        processor.close()

    def test_consolidation_checkpoints_progress(self, storage, temp_dir):
        """Consolidation should checkpoint progress in manifest."""
        # This is tested indirectly through the _consolidate_shards behavior
        # Create a manifest to track
        manifest = JobManifest(
            job_id="checkpoint-test",
            input_urls_count=100,
        )
        manifest.status = JobStatus.CONSOLIDATING
        manifest.consolidation_phase.partitions_total = 20

        # Simulate adding completed partitions
        for i in range(15):
            manifest.consolidation_phase.completed_partitions.append(f"partition_{i}")
            manifest.consolidation_phase.partitions_completed += 1

        manifest.save(storage, temp_dir)

        # Load and verify checkpoint
        loaded = JobManifest.load(storage, temp_dir, "checkpoint-test")
        assert len(loaded.consolidation_phase.completed_partitions) == 15
        assert loaded.consolidation_phase.partitions_completed == 15
