#!/usr/bin/env python3
"""
FileChannel regression tests.

These checks focus on append/rewrite semantics, polling fallback, and basic
high-volume file ingestion behavior.
"""

from __future__ import annotations

import json
import os
import shutil
import sys
import tempfile
import threading
import time
import importlib
from pathlib import Path

from test_support import install_sdk_dependency_stubs

install_sdk_dependency_stubs()

sys.path.insert(0, str(Path(__file__).parent.parent))

FileChannel = importlib.import_module("sandai.operator.channel").FileChannel


def _write_jsonl(path: Path, rows: list[dict]) -> None:
    with open(path, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")


def _append_jsonl(path: Path, rows: list[dict]) -> None:
    with open(path, "a", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")


def _replace_jsonl(path: Path, rows: list[dict]) -> None:
    temp_path = path.with_suffix(path.suffix + ".tmp")
    _write_jsonl(temp_path, rows)
    shutil.move(str(temp_path), str(path))


def _consume_tasks(channel: FileChannel, limit: int, timeout: float) -> list[dict]:
    results: list[dict] = []

    def consumer() -> None:
        while len(results) < limit:
            try:
                task = channel.pull(timeout=timeout)
            except Exception:
                break
            results.append(task)
            channel.push({"id": task["task_id"], "processed": True})

    thread = threading.Thread(target=consumer, daemon=True)
    thread.start()
    thread.join(timeout=max(3.0, timeout * limit + 1.0))
    return results


def test_file_append_detects_new_rows() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "append_input.jsonl"
        output_path = Path(tmpdir) / "append_output.jsonl"
        _write_jsonl(input_path, [{"task_id": "append-001", "data": "initial"}])

        channel = FileChannel(str(input_path), str(output_path), debug=False)
        results: list[dict] = []

        def consumer() -> None:
            while len(results) < 3:
                try:
                    task = channel.pull(timeout=3.0)
                except Exception:
                    break
                results.append(task)
                channel.push({"id": task["task_id"], "processed": True})

        thread = threading.Thread(target=consumer, daemon=True)
        thread.start()
        time.sleep(0.4)
        _append_jsonl(
            input_path,
            [
                {"task_id": "append-002", "data": "appended-1"},
                {"task_id": "append-003", "data": "appended-2"},
            ],
        )
        thread.join(timeout=5.0)
        channel.close()

        assert [task["task_id"] for task in results] == [
            "append-001",
            "append-002",
            "append-003",
        ], results


def test_file_replacement_detects_vim_style_save() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "replace_input.jsonl"
        output_path = Path(tmpdir) / "replace_output.jsonl"
        _write_jsonl(input_path, [{"task_id": "replace-001", "data": "initial"}])

        channel = FileChannel(str(input_path), str(output_path), debug=False)
        results: list[dict] = []

        def consumer() -> None:
            while len(results) < 3:
                try:
                    task = channel.pull(timeout=4.0)
                except Exception:
                    break
                results.append(task)
                channel.push({"id": task["task_id"], "processed": True})

        thread = threading.Thread(target=consumer, daemon=True)
        thread.start()
        time.sleep(0.4)
        _replace_jsonl(
            input_path,
            [
                {"task_id": "replace-001", "data": "initial"},
                {"task_id": "replace-002", "data": "rewrite-1"},
                {"task_id": "replace-003", "data": "rewrite-2"},
            ],
        )
        deadline = time.time() + 6.0
        while len(results) < 3 and time.time() < deadline:
            time.sleep(0.2)

        thread.join(timeout=1.0)
        channel.close()

        observed_ids = {task["task_id"] for task in results}
        assert "replace-002" in observed_ids, observed_ids
        assert "replace-003" in observed_ids, observed_ids


def test_polling_fallback_detects_new_rows_without_watcher() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "polling_input.jsonl"
        output_path = Path(tmpdir) / "polling_output.jsonl"
        _write_jsonl(input_path, [{"task_id": "polling-001", "data": "initial"}])

        channel = FileChannel(
            str(input_path),
            str(output_path),
            use_file_watcher=False,
            debug=False,
        )
        results: list[dict] = []

        def consumer() -> None:
            while len(results) < 2:
                try:
                    task = channel.pull(timeout=3.0)
                except Exception:
                    break
                results.append(task)
                channel.push({"id": task["task_id"], "processed": True})

        thread = threading.Thread(target=consumer, daemon=True)
        thread.start()
        time.sleep(1.0)
        _append_jsonl(input_path, [{"task_id": "polling-002", "data": "polled"}])
        thread.join(timeout=5.0)
        channel.close()

        assert [task["task_id"] for task in results] == ["polling-001", "polling-002"], results


def test_large_batch_append_drains_all_rows() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "batch_input.jsonl"
        output_path = Path(tmpdir) / "batch_output.jsonl"
        _write_jsonl(input_path, [{"task_id": "batch-001", "data": "initial"}])

        channel = FileChannel(str(input_path), str(output_path), debug=False)
        results: list[dict] = []

        def consumer() -> None:
            while len(results) < 10:
                try:
                    task = channel.pull(timeout=2.0)
                except Exception:
                    break
                results.append(task)
                channel.push({"id": task["task_id"], "processed": True})

        thread = threading.Thread(target=consumer, daemon=True)
        thread.start()
        time.sleep(0.4)
        _append_jsonl(
            input_path,
            [{"task_id": f"batch-{i:03d}", "data": f"batch-{i}"} for i in range(2, 11)],
        )
        thread.join(timeout=6.0)
        channel.close()

        assert len(results) == 10, len(results)
        assert results[-1]["task_id"] == "batch-010", results


def test_vim_style_rewrite_preserves_progress() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "vim_input.jsonl"
        output_path = Path(tmpdir) / "vim_output.jsonl"
        rows = [{"task_id": "vim-001", "data": "initial"}]
        _write_jsonl(input_path, rows)

        channel = FileChannel(str(input_path), str(output_path), debug=False)
        results: list[dict] = []

        def consumer() -> None:
            while len(results) < 3:
                try:
                    task = channel.pull(timeout=3.0)
                except Exception:
                    break
                results.append(task)
                channel.push({"id": task["task_id"], "processed": True})

        thread = threading.Thread(target=consumer, daemon=True)
        thread.start()
        time.sleep(0.4)

        rows.append({"task_id": "vim-002", "data": "rewrite-1"})
        _replace_jsonl(input_path, rows)
        deadline = time.time() + 5.0
        while len(results) < 2 and time.time() < deadline:
            time.sleep(0.2)

        thread.join(timeout=1.0)
        channel.close()

        observed_ids = {task["task_id"] for task in results}
        assert {"vim-001", "vim-002"}.issubset(observed_ids), observed_ids


def test_reappend_same_task_id_after_processing_is_not_globally_deduplicated() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "replay_input.jsonl"
        output_path = Path(tmpdir) / "replay_output.jsonl"
        _write_jsonl(input_path, [{"task_id": "replay-001", "data": "first"}])

        channel = FileChannel(str(input_path), str(output_path), debug=False)
        results: list[dict] = []

        def consumer() -> None:
            while len(results) < 2:
                try:
                    task = channel.pull(timeout=3.0)
                except Exception:
                    break
                results.append(task)
                channel.push({"task_id": task["task_id"], "processed": True})

        thread = threading.Thread(target=consumer, daemon=True)
        thread.start()
        time.sleep(0.6)
        _append_jsonl(input_path, [{"task_id": "replay-001", "data": "second"}])
        thread.join(timeout=5.0)
        channel.close()


def test_file_channel_shards_input_when_sdrun_env_is_present() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "shard_input.jsonl"
        output_path = Path(tmpdir) / "shard_output.jsonl"
        _write_jsonl(
            input_path,
            [{"task_id": f"shard-{i:03d}", "data": i} for i in range(6)],
        )

        original_env = {key: os.environ.get(key) for key in ("SDRUN_MODE", "SDRUN_WORLD_SIZE", "SDRUN_RANK")}
        os.environ["SDRUN_MODE"] = "true"
        os.environ["SDRUN_WORLD_SIZE"] = "2"
        os.environ["SDRUN_RANK"] = "1"
        try:
            channel = FileChannel(str(input_path), str(output_path), debug=False, job_mode=True)
            results = _consume_tasks(channel, limit=3, timeout=2.0)
            channel.close()
        finally:
            for key, value in original_env.items():
                if value is None:
                    os.environ.pop(key, None)
                else:
                    os.environ[key] = value

        assert [task["task_id"] for task in results] == [
            "shard-001",
            "shard-003",
            "shard-005",
        ], results


def test_file_channel_sdrun_output_path_uses_rank_suffix_before_extension() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "suffix_input.jsonl"
        output_path = Path(tmpdir) / "suffix_output.jsonl"
        _write_jsonl(
            input_path,
            [
                {"task_id": "suffix-000", "data": "skip-0"},
                {"task_id": "suffix-001", "data": "skip-1"},
                {"task_id": "suffix-002", "data": "take-2"},
            ],
        )

        original_env = {key: os.environ.get(key) for key in ("SDRUN_MODE", "SDRUN_WORLD_SIZE", "SDRUN_RANK")}
        os.environ["SDRUN_MODE"] = "true"
        os.environ["SDRUN_WORLD_SIZE"] = "4"
        os.environ["SDRUN_RANK"] = "2"
        try:
            channel = FileChannel(str(input_path), str(output_path), debug=False, job_mode=True)
            task = channel.pull(timeout=2.0)
            if task.get("marker_type") == "eof":
                raise AssertionError("expected a sharded task, got EOF marker")
            channel.push({"task_id": task["task_id"], "processed": True})
            time.sleep(0.3)
            sink_path = Path(channel.sink_jsonl_path)
            channel.close()
        finally:
            for key, value in original_env.items():
                if value is None:
                    os.environ.pop(key, None)
                else:
                    os.environ[key] = value

        assert sink_path.name == "suffix_output.2.jsonl", sink_path
        assert sink_path.exists(), sink_path
        rows = [json.loads(line) for line in sink_path.read_text(encoding="utf-8").splitlines() if line]
        assert rows == [{"task_id": "suffix-002", "processed": True}], rows


def main() -> int:
    tests = [
        ("append", test_file_append_detects_new_rows),
        ("replace", test_file_replacement_detects_vim_style_save),
        ("polling_fallback", test_polling_fallback_detects_new_rows_without_watcher),
        ("large_batch", test_large_batch_append_drains_all_rows),
        ("vim_rewrite_progress", test_vim_style_rewrite_preserves_progress),
        ("same_task_id_replay", test_reappend_same_task_id_after_processing_is_not_globally_deduplicated),
        ("sdrun_shards_input", test_file_channel_shards_input_when_sdrun_env_is_present),
        ("sdrun_output_suffix", test_file_channel_sdrun_output_path_uses_rank_suffix_before_extension),
    ]

    passed = 0
    failed = 0

    print("FileChannel regression tests")
    print("=" * 60)
    for name, test_func in tests:
        try:
            test_func()
            print(f"PASS {name}")
            passed += 1
        except Exception as exc:
            print(f"FAIL {name}: {exc}")
            import traceback

            traceback.print_exc()
            failed += 1
    print("=" * 60)
    print(f"passed={passed} failed={failed}")
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main())
