"""Tests for axiom.infra.state — safe concurrent JSON state access.

Proves:
1. Basic read/write/read-modify-write works correctly
2. Atomic writes don't corrupt on partial failure
3. Concurrent processes can safely read-modify-write the same file
4. Lock semantics are correct (shared reads, exclusive writes)
5. StateLocation registry is complete and consistent
"""

from __future__ import annotations

import json
import multiprocessing
import os
import time
from pathlib import Path

import pytest

from axiom.infra.state import (
    STATE_LOCATIONS,
    LockedJsonFile,
    atomic_write,
    locked_read,
)

# ---------------------------------------------------------------------------
# Basic read/write
# ---------------------------------------------------------------------------


class TestLockedJsonFileBasic:
    """Basic LockedJsonFile operations."""

    def test_read_missing_file_returns_empty_dict(self, tmp_path: Path):
        path = tmp_path / "missing.json"
        with LockedJsonFile(path) as f:
            assert f.read() == {}

    def test_read_empty_file_returns_empty_dict(self, tmp_path: Path):
        path = tmp_path / "empty.json"
        path.write_text("")
        with LockedJsonFile(path) as f:
            assert f.read() == {}

    def test_read_valid_json(self, tmp_path: Path):
        path = tmp_path / "data.json"
        path.write_text('{"key": "value"}')
        with LockedJsonFile(path) as f:
            assert f.read() == {"key": "value"}

    def test_read_json_array(self, tmp_path: Path):
        path = tmp_path / "arr.json"
        path.write_text('[1, 2, 3]')
        with LockedJsonFile(path) as f:
            assert f.read() == [1, 2, 3]

    def test_read_corrupt_json_returns_empty_dict(self, tmp_path: Path):
        path = tmp_path / "bad.json"
        path.write_text("{not valid json")
        with LockedJsonFile(path) as f:
            assert f.read() == {}

    def test_write_requires_exclusive(self, tmp_path: Path):
        path = tmp_path / "data.json"
        with LockedJsonFile(path) as f:
            with pytest.raises(RuntimeError, match="exclusive"):
                f.write({"key": "value"})

    def test_write_creates_file(self, tmp_path: Path):
        path = tmp_path / "new.json"
        with LockedJsonFile(path, exclusive=True) as f:
            f.write({"hello": "world"})
        assert path.exists()
        assert json.loads(path.read_text()) == {"hello": "world"}

    def test_write_overwrites_existing(self, tmp_path: Path):
        path = tmp_path / "data.json"
        path.write_text('{"old": true}')
        with LockedJsonFile(path, exclusive=True) as f:
            f.write({"new": True})
        assert json.loads(path.read_text()) == {"new": True}

    def test_read_modify_write(self, tmp_path: Path):
        path = tmp_path / "data.json"
        path.write_text('{"counter": 0}')
        with LockedJsonFile(path, exclusive=True) as f:
            data = f.read()
            data["counter"] += 1
            f.write(data)
        assert json.loads(path.read_text()) == {"counter": 1}

    def test_context_manager_required(self, tmp_path: Path):
        f = LockedJsonFile(tmp_path / "data.json")
        with pytest.raises(RuntimeError, match="context manager"):
            f.read()

    def test_creates_parent_dirs(self, tmp_path: Path):
        path = tmp_path / "deep" / "nested" / "data.json"
        with LockedJsonFile(path, exclusive=True) as f:
            f.write({"created": True})
        assert path.exists()


# ---------------------------------------------------------------------------
# Convenience helpers
# ---------------------------------------------------------------------------


class TestConvenienceHelpers:
    def test_atomic_write(self, tmp_path: Path):
        path = tmp_path / "data.json"
        atomic_write(path, {"fast": True})
        assert json.loads(path.read_text()) == {"fast": True}

    def test_locked_read(self, tmp_path: Path):
        path = tmp_path / "data.json"
        path.write_text('{"x": 42}')
        assert locked_read(path) == {"x": 42}

    def test_locked_read_missing(self, tmp_path: Path):
        assert locked_read(tmp_path / "nope.json") == {}


# ---------------------------------------------------------------------------
# Atomic write safety
# ---------------------------------------------------------------------------


class TestAtomicWriteSafety:
    """Prove that atomic writes don't leave corrupt files."""

    def test_original_intact_on_write_error(self, tmp_path: Path):
        """If write fails mid-stream, original file is untouched."""
        path = tmp_path / "data.json"
        path.write_text('{"original": true}')

        class Unserializable:
            pass

        with pytest.raises(TypeError):
            with LockedJsonFile(path, exclusive=True) as f:
                f.write({"bad": Unserializable()})

        # Original file should be intact
        assert json.loads(path.read_text()) == {"original": True}

    def test_no_temp_files_left_on_success(self, tmp_path: Path):
        path = tmp_path / "data.json"
        with LockedJsonFile(path, exclusive=True) as f:
            f.write({"clean": True})
        # No .tmp files should remain
        tmp_files = list(tmp_path.glob("*.tmp"))
        assert tmp_files == []

    def test_no_temp_files_left_on_failure(self, tmp_path: Path):
        path = tmp_path / "data.json"
        path.write_text("{}")

        class Boom:
            pass

        with pytest.raises(TypeError):
            with LockedJsonFile(path, exclusive=True) as f:
                f.write({"bad": Boom()})

        tmp_files = list(tmp_path.glob("*.tmp"))
        assert tmp_files == []

    def test_lock_file_created(self, tmp_path: Path):
        """Lock sidecar file is created alongside data file."""
        path = tmp_path / "data.json"
        with LockedJsonFile(path, exclusive=True) as f:
            f.write({"x": 1})
            # Lock file should exist while we hold the lock
            assert path.with_suffix(".json.lock").exists()


# ---------------------------------------------------------------------------
# Concurrent access — the critical test
# ---------------------------------------------------------------------------

def _worker_increment(path_str: str, iterations: int, results_path_str: str):
    """Worker process: read-modify-write a counter N times."""
    path = Path(path_str)
    errors = 0
    for _ in range(iterations):
        try:
            with LockedJsonFile(path, exclusive=True) as f:
                data = f.read()
                data["counter"] = data.get("counter", 0) + 1
                f.write(data)
        except Exception:
            errors += 1
    # Write results
    Path(results_path_str).write_text(json.dumps({"errors": errors}))


def _worker_reader(path_str: str, result_str: str):
    """Worker process: read with shared lock, hold briefly."""
    path = Path(path_str)
    with LockedJsonFile(path) as f:
        data = f.read()
        time.sleep(0.1)  # Hold the shared lock briefly
    Path(result_str).write_text(json.dumps(data))


class TestConcurrentAccess:
    """Prove that LockedJsonFile prevents data corruption under concurrency.

    These tests fork real processes to exercise fcntl.flock.
    """

    @pytest.mark.skipif(
        os.name == "nt",
        reason="fcntl not available on Windows",
    )
    def test_concurrent_increments_no_lost_updates(self, tmp_path: Path):
        """N processes each increment a counter M times.

        Without locking, final count < N*M due to lost updates.
        With locking, final count == N*M exactly.
        """
        num_workers = 4
        iterations_per_worker = 50
        expected_total = num_workers * iterations_per_worker

        data_path = tmp_path / "counter.json"
        data_path.write_text('{"counter": 0}')

        # Create result files for each worker
        result_paths = [tmp_path / f"result_{i}.json" for i in range(num_workers)]

        processes = []
        for i in range(num_workers):
            p = multiprocessing.Process(
                target=_worker_increment,
                args=(str(data_path), iterations_per_worker, str(result_paths[i])),
            )
            processes.append(p)

        for p in processes:
            p.start()
        for p in processes:
            p.join(timeout=30)

        # Verify no worker errors
        for rp in result_paths:
            assert rp.exists(), f"Worker result missing: {rp}"
            result = json.loads(rp.read_text())
            assert result["errors"] == 0, f"Worker had errors: {result}"

        # THE CRITICAL ASSERTION: no lost updates
        final_data = json.loads(data_path.read_text())
        assert final_data["counter"] == expected_total, (
            f"Lost updates! Expected {expected_total}, got {final_data['counter']}. "
            f"This means {expected_total - final_data['counter']} increments were lost "
            f"due to concurrent access."
        )

    @pytest.mark.skipif(
        os.name == "nt",
        reason="fcntl not available on Windows",
    )
    def test_concurrent_reads_dont_block_each_other(self, tmp_path: Path):
        """Multiple shared-lock readers can proceed concurrently."""
        data_path = tmp_path / "shared.json"
        data_path.write_text('{"value": 42}')

        result_paths = [tmp_path / f"read_{i}.json" for i in range(3)]
        processes = []
        for i in range(3):
            p = multiprocessing.Process(
                target=_worker_reader,
                args=(str(data_path), str(result_paths[i])),
            )
            processes.append(p)

        start = time.monotonic()
        for p in processes:
            p.start()
        for p in processes:
            p.join(timeout=10)
        elapsed = time.monotonic() - start

        # All readers should complete; with shared locks they run concurrently
        for rp in result_paths:
            assert rp.exists()
            assert json.loads(rp.read_text()) == {"value": 42}

        # If they blocked each other serially, it'd take ~0.3s+. Concurrent ~0.1s+.
        # Allow generous margin but catch serial behavior.
        assert elapsed < 1.0, f"Readers appear to have serialized: {elapsed:.2f}s"

    @pytest.mark.skipif(
        os.name == "nt",
        reason="fcntl not available on Windows",
    )
    def test_valid_json_after_concurrent_writes(self, tmp_path: Path):
        """File always contains valid JSON, never a partial write."""
        data_path = tmp_path / "validity.json"
        data_path.write_text('{"counter": 0}')

        num_workers = 4
        iterations = 25

        result_paths = [tmp_path / f"res_{i}.json" for i in range(num_workers)]
        processes = []
        for i in range(num_workers):
            p = multiprocessing.Process(
                target=_worker_increment,
                args=(str(data_path), iterations, str(result_paths[i])),
            )
            processes.append(p)

        for p in processes:
            p.start()
        for p in processes:
            p.join(timeout=30)

        # File must be valid JSON
        content = data_path.read_text()
        data = json.loads(content)  # Should not raise
        assert isinstance(data, dict)
        assert "counter" in data


# ---------------------------------------------------------------------------
# StateLocation registry
# ---------------------------------------------------------------------------


class TestStateLocationRegistry:
    """Verify STATE_LOCATIONS registry is well-formed."""

    def test_all_locations_have_required_fields(self):
        for loc in STATE_LOCATIONS:
            assert loc.path, f"Empty path in {loc}"
            assert loc.category in (
                "runtime", "config", "documents", "corrections", "sessions", "credentials",
            ), f"Invalid category '{loc.category}' for {loc.path}"
            assert loc.description, f"Empty description for {loc.path}"
            assert loc.sensitivity in (
                "low", "medium", "high", "critical",
            ), f"Invalid sensitivity '{loc.sensitivity}' for {loc.path}"

    def test_no_duplicate_paths(self):
        paths = [loc.path for loc in STATE_LOCATIONS]
        assert len(paths) == len(set(paths)), (
            f"Duplicate paths: {[p for p in paths if paths.count(p) > 1]}"
        )

    def test_retention_keys_are_valid(self):
        """Locations with retention_key should reference known policy names."""
        known_keys = {"raw_voice", "raw_signals", "transcripts", "sessions", "drafts"}
        for loc in STATE_LOCATIONS:
            if loc.retention_key is not None:
                assert loc.retention_key in known_keys, (
                    f"Unknown retention_key '{loc.retention_key}' for {loc.path}"
                )

    def test_config_locations_have_no_retention(self):
        """Config and corrections are indefinite — no retention key."""
        for loc in STATE_LOCATIONS:
            if loc.category in ("config", "corrections"):
                assert loc.retention_key is None, (
                    f"Config/corrections location should not have retention: {loc.path}"
                )

    def test_minimum_location_count(self):
        """Sanity check: we should have at least 10 known locations."""
        assert len(STATE_LOCATIONS) >= 10
