"""
tests/test_batch_worker.py — Tests for batch ingestion worker and API (#1326).

Tests:
    1. Upload .md file → job created
    2. Job status shows progress
    3. Upload .jsonl file → Mem0 memories extracted
    4. Mem0JsonlSource deduplication
    5. Idempotency: same content twice → 0 new claims on second run
    6. Cancel job works
    7. List jobs returns correct data
    8. Start job transitions status
    9. Results endpoint returns detailed info
    10. Migration v23 creates batch_jobs table
"""
from __future__ import annotations

import json
import os
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

import aiosqlite

from cortex.batch_ingest import Chunk, Mem0JsonlSource


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def tmp_dir(tmp_path):
    """Create a temp dir for test data."""
    return tmp_path


@pytest.fixture
def sample_md_content():
    return "# Test Document\n\nThis is a test document for batch ingestion.\n\nIt has multiple paragraphs to test chunking behavior.\n"


@pytest.fixture
def sample_jsonl_content():
    records = [
        {"id": "mem-001", "memory": "Andrew prefers Python for backend work", "category": "preference"},
        {"id": "mem-002", "memory": "Eva was built on March 1, 2026", "category": "fact"},
        {"id": "mem-003", "memory": "The Electric Sheep project uses FastAPI", "category": "technical"},
    ]
    return "\n".join(json.dumps(r) for r in records) + "\n"


@pytest.fixture
def sample_jsonl_with_dupes():
    records = [
        {"id": "mem-001", "memory": "Andrew prefers Python", "category": "preference"},
        {"id": "mem-001", "memory": "Andrew prefers Python (duplicate)", "category": "preference"},
        {"id": "mem-002", "memory": "Eva uses FastAPI", "category": "technical"},
    ]
    return "\n".join(json.dumps(r) for r in records) + "\n"


# ---------------------------------------------------------------------------
# Test: Mem0JsonlSource
# ---------------------------------------------------------------------------

class TestMem0JsonlSource:
    def test_basic_extraction(self, tmp_dir, sample_jsonl_content):
        """JSONL file → correct chunks extracted."""
        jsonl_path = tmp_dir / "memories.jsonl"
        jsonl_path.write_text(sample_jsonl_content)

        source = Mem0JsonlSource(str(jsonl_path))
        chunks = list(source.iter_chunks())

        assert len(chunks) == 3
        assert chunks[0].content == "Andrew prefers Python for backend work"
        assert chunks[0].source_type == "mem0_jsonl"
        assert chunks[0].metadata["mem0_id"] == "mem-001"
        assert chunks[1].content == "Eva was built on March 1, 2026"
        assert chunks[2].content == "The Electric Sheep project uses FastAPI"

    def test_deduplication_by_id(self, tmp_dir, sample_jsonl_with_dupes):
        """Duplicate mem0 IDs → only first occurrence kept."""
        jsonl_path = tmp_dir / "dupes.jsonl"
        jsonl_path.write_text(sample_jsonl_with_dupes)

        source = Mem0JsonlSource(str(jsonl_path))
        chunks = list(source.iter_chunks())

        assert len(chunks) == 2
        assert chunks[0].content == "Andrew prefers Python"
        assert chunks[1].content == "Eva uses FastAPI"

    def test_empty_memories_skipped(self, tmp_dir):
        """Lines with empty memory field are skipped."""
        content = json.dumps({"id": "1", "memory": ""}) + "\n"
        content += json.dumps({"id": "2", "memory": "  "}) + "\n"
        content += json.dumps({"id": "3", "memory": "Real memory"}) + "\n"

        jsonl_path = tmp_dir / "sparse.jsonl"
        jsonl_path.write_text(content)

        source = Mem0JsonlSource(str(jsonl_path))
        chunks = list(source.iter_chunks())

        assert len(chunks) == 1
        assert chunks[0].content == "Real memory"

    def test_invalid_json_lines_skipped(self, tmp_dir):
        """Invalid JSON lines logged and skipped."""
        content = '{"id": "1", "memory": "Valid"}\n'
        content += "not valid json\n"
        content += '{"id": "2", "memory": "Also valid"}\n'

        jsonl_path = tmp_dir / "mixed.jsonl"
        jsonl_path.write_text(content)

        source = Mem0JsonlSource(str(jsonl_path))
        chunks = list(source.iter_chunks())

        assert len(chunks) == 2

    def test_file_not_found(self, tmp_dir):
        """Missing file raises FileNotFoundError."""
        source = Mem0JsonlSource(str(tmp_dir / "nonexistent.jsonl"))
        with pytest.raises(FileNotFoundError):
            list(source.iter_chunks())

    def test_no_id_field(self, tmp_dir):
        """Records without id field still work (no dedup)."""
        content = json.dumps({"memory": "First"}) + "\n"
        content += json.dumps({"memory": "Second"}) + "\n"
        content += json.dumps({"memory": "First"}) + "\n"  # same content, no id

        jsonl_path = tmp_dir / "noid.jsonl"
        jsonl_path.write_text(content)

        source = Mem0JsonlSource(str(jsonl_path))
        chunks = list(source.iter_chunks())

        # All 3 kept since no id-based dedup
        assert len(chunks) == 3


# ---------------------------------------------------------------------------
# Test: Migration v23
# ---------------------------------------------------------------------------

class TestMigrationV23:
    @pytest.mark.asyncio
    async def test_batch_jobs_table_created(self, tmp_dir):
        """Migration v23 creates batch_jobs table with correct schema."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            # Create cortex_config table (needed by migration runner)
            await db.execute(
                "CREATE TABLE IF NOT EXISTS cortex_config "
                "(key TEXT PRIMARY KEY, value TEXT, updated_at TEXT)"
            )
            await db.commit()

            # Apply v23 migration directly
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)

            # Verify table exists and has correct columns
            async with db.execute("PRAGMA table_info(batch_jobs)") as cur:
                columns = {row[1] for row in await cur.fetchall()}

            expected_columns = {
                "id", "status", "source_type", "filename",
                "total_chunks", "processed_chunks", "claims_stored",
                "errors_json", "config_json", "created_at", "updated_at",
                "completed_at", "file_results_json",
            }
            assert expected_columns.issubset(columns)

            # Verify we can insert a row
            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "INSERT INTO batch_jobs (id, source_type, config_json, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?)",
                ("test-job", "file", '{"test": true}', now, now),
            )
            await db.commit()

            async with db.execute("SELECT * FROM batch_jobs WHERE id = 'test-job'") as cur:
                row = await cur.fetchone()
                assert row is not None


# ---------------------------------------------------------------------------
# Test: Batch Worker
# ---------------------------------------------------------------------------

class TestBatchWorker:
    @pytest.mark.asyncio
    async def test_reset_running_jobs(self, tmp_dir):
        """On startup, running jobs are reset to pending."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)

            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "INSERT INTO batch_jobs (id, status, source_type, config_json, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                ("stale-job", "running", "file", '{}', now, now),
            )
            await db.commit()

        from cortex.batch_worker import BatchWorker
        worker = BatchWorker(db_path=db_path, plugin=MagicMock())
        await worker._reset_running_jobs()

        async with aiosqlite.connect(db_path) as db:
            async with db.execute("SELECT status FROM batch_jobs WHERE id = 'stale-job'") as cur:
                row = await cur.fetchone()
                assert row[0] == "pending"

    @pytest.mark.asyncio
    async def test_pick_next_job(self, tmp_dir):
        """Worker picks oldest pending job and marks it running."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)

            now = datetime.now(timezone.utc).isoformat()
            # Insert two pending jobs
            await db.execute(
                "INSERT INTO batch_jobs (id, status, source_type, config_json, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                ("job-1", "pending", "file", '{}', "2026-01-01T00:00:00Z", now),
            )
            await db.execute(
                "INSERT INTO batch_jobs (id, status, source_type, config_json, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                ("job-2", "pending", "file", '{}', "2026-01-02T00:00:00Z", now),
            )
            await db.commit()

        from cortex.batch_worker import BatchWorker
        worker = BatchWorker(db_path=db_path, plugin=MagicMock())
        job = await worker._pick_next_job()

        assert job is not None
        assert job["id"] == "job-1"  # Oldest first

        # Verify it's now 'running'
        async with aiosqlite.connect(db_path) as db:
            async with db.execute("SELECT status FROM batch_jobs WHERE id = 'job-1'") as cur:
                row = await cur.fetchone()
                assert row[0] == "running"

    @pytest.mark.asyncio
    async def test_no_pending_jobs(self, tmp_dir):
        """Worker returns None when no pending jobs."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)
            await db.commit()

        from cortex.batch_worker import BatchWorker
        worker = BatchWorker(db_path=db_path, plugin=MagicMock())
        job = await worker._pick_next_job()
        assert job is None

    @pytest.mark.asyncio
    async def test_complete_job(self, tmp_dir):
        """Complete job marks status and timestamps."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)

            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "INSERT INTO batch_jobs (id, status, source_type, config_json, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                ("job-done", "running", "file", '{}', now, now),
            )
            await db.commit()

        from cortex.batch_worker import BatchWorker
        worker = BatchWorker(db_path=db_path, plugin=MagicMock())
        await worker._complete_job("job-done", 10, 5, [])

        async with aiosqlite.connect(db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute("SELECT * FROM batch_jobs WHERE id = 'job-done'") as cur:
                row = dict(await cur.fetchone())
                assert row["status"] == "completed"
                assert row["processed_chunks"] == 10
                assert row["claims_stored"] == 5
                assert row["completed_at"] is not None

    @pytest.mark.asyncio
    async def test_fail_job(self, tmp_dir):
        """Fail job sets status and error."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)

            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "INSERT INTO batch_jobs (id, status, source_type, config_json, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                ("job-fail", "running", "file", '{}', now, now),
            )
            await db.commit()

        from cortex.batch_worker import BatchWorker
        worker = BatchWorker(db_path=db_path, plugin=MagicMock())
        await worker._fail_job("job-fail", "Something went wrong")

        async with aiosqlite.connect(db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute("SELECT * FROM batch_jobs WHERE id = 'job-fail'") as cur:
                row = dict(await cur.fetchone())
                assert row["status"] == "failed"
                errors = json.loads(row["errors_json"])
                assert len(errors) == 1
                assert "Something went wrong" in errors[0]["error"]

    @pytest.mark.asyncio
    async def test_build_chunks_file(self, tmp_dir):
        """Build chunks from a .md file source."""
        db_path = str(tmp_dir / "test.db")
        job_id = "test-build"

        # Create upload dir with file
        upload_dir = tmp_dir / "batch_uploads" / job_id
        upload_dir.mkdir(parents=True)
        md_file = upload_dir / "test.md"
        md_file.write_text("# Hello\n\nThis is test content for chunking.\n")

        from cortex.batch_worker import BatchWorker
        worker = BatchWorker(db_path=db_path, plugin=MagicMock())

        with patch.dict(os.environ, {"CORTEX_DATA_DIR": str(tmp_dir)}):
            chunks = worker._build_chunks(
                job_id, "file", {"filename": "test.md"}
            )

        assert chunks is not None
        assert len(chunks) >= 1
        assert "Hello" in chunks[0].content or "test content" in chunks[0].content

    @pytest.mark.asyncio
    async def test_build_chunks_mem0(self, tmp_dir, sample_jsonl_content):
        """Build chunks from a .jsonl Mem0 source."""
        job_id = "test-mem0"

        upload_dir = tmp_dir / "batch_uploads" / job_id
        upload_dir.mkdir(parents=True)
        jsonl_file = upload_dir / "memories.jsonl"
        jsonl_file.write_text(sample_jsonl_content)

        from cortex.batch_worker import BatchWorker
        worker = BatchWorker(db_path="", plugin=MagicMock())

        with patch.dict(os.environ, {"CORTEX_DATA_DIR": str(tmp_dir)}):
            chunks = worker._build_chunks(
                job_id, "mem0_jsonl", {"filename": "memories.jsonl"}
            )

        assert chunks is not None
        assert len(chunks) == 3


# ---------------------------------------------------------------------------
# Test: Batch Routes (unit-level, no full app)
# ---------------------------------------------------------------------------

class TestBatchRoutes:
    @pytest.mark.asyncio
    async def test_cancel_nonexistent_job(self, tmp_dir):
        """Cancelling a non-existent job raises 404."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)
            await db.commit()

        # Test the cancel logic directly
        async with aiosqlite.connect(db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute(
                "SELECT * FROM batch_jobs WHERE id = ?", ("nonexistent",)
            ) as cur:
                row = await cur.fetchone()
                assert row is None  # 404 case

    @pytest.mark.asyncio
    async def test_cancel_completed_job(self, tmp_dir):
        """Cancelling a completed job is rejected (409)."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)

            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "INSERT INTO batch_jobs (id, status, source_type, config_json, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                ("done-job", "completed", "file", '{}', now, now),
            )
            await db.commit()

        # Verify completed status prevents cancellation
        async with aiosqlite.connect(db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute(
                "SELECT * FROM batch_jobs WHERE id = ?", ("done-job",)
            ) as cur:
                row = dict(await cur.fetchone())
                assert row["status"] in ("completed", "cancelled")

    @pytest.mark.asyncio
    async def test_list_jobs_with_filter(self, tmp_dir):
        """List jobs with status filter returns correct subset."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)

            now = datetime.now(timezone.utc).isoformat()
            for i, s in enumerate(["pending", "completed", "pending", "failed"]):
                await db.execute(
                    "INSERT INTO batch_jobs (id, status, source_type, config_json, created_at, updated_at) "
                    "VALUES (?, ?, ?, ?, ?, ?)",
                    (f"job-{i}", s, "file", '{}', now, now),
                )
            await db.commit()

        async with aiosqlite.connect(db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute(
                "SELECT * FROM batch_jobs WHERE status = ? ORDER BY created_at DESC",
                ("pending",),
            ) as cur:
                rows = await cur.fetchall()
                assert len(rows) == 2

    @pytest.mark.asyncio
    async def test_idempotency_store_prevents_reprocessing(self, tmp_dir):
        """Same chunk hash processed twice → second time is skipped."""
        from cortex.batch_ingest import IdempotencyStore

        store_path = str(tmp_dir / "idem.json")
        store = IdempotencyStore(state_path=store_path)

        chunk = Chunk(
            content="Test content",
            session_id="test",
            chunk_index=0,
            source_type="file",
        )

        assert not store.is_processed(chunk.chunk_hash)
        store.mark_processed(chunk.chunk_hash)
        assert store.is_processed(chunk.chunk_hash)

        # New store instance with same path should load state
        store.save()
        store2 = IdempotencyStore(state_path=store_path)
        assert store2.is_processed(chunk.chunk_hash)

    @pytest.mark.asyncio
    async def test_update_progress(self, tmp_dir):
        """Worker updates progress correctly."""
        db_path = str(tmp_dir / "test.db")

        async with aiosqlite.connect(db_path) as db:
            from cortex.storage.migrations.v23_batch_jobs import SQL_UP
            await db.executescript(SQL_UP)

            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "INSERT INTO batch_jobs (id, status, source_type, config_json, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                ("progress-job", "running", "file", '{}', now, now),
            )
            await db.commit()

        from cortex.batch_worker import BatchWorker
        worker = BatchWorker(db_path=db_path, plugin=MagicMock())
        await worker._update_progress("progress-job", 5, 3, [{"error": "test"}])

        async with aiosqlite.connect(db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute("SELECT * FROM batch_jobs WHERE id = 'progress-job'") as cur:
                row = dict(await cur.fetchone())
                assert row["processed_chunks"] == 5
                assert row["claims_stored"] == 3
                errors = json.loads(row["errors_json"])
                assert len(errors) == 1
