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

These tests cover task workdir creation, isolation, idempotent lookup, and
basic batch-style usage through the public task input/output types.
"""

from __future__ import annotations

import json
import sys
import tempfile
from pathlib import Path
from typing import Generator, List

from test_support import install_sdk_dependency_stubs

install_sdk_dependency_stubs()

from pydantic import BaseModel

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

from sandai.operator import InputArtifact, ProcessingContext, TaskInput, TaskOutput


class ContextOptions(BaseModel):
    test_param: str = "default"


class ContextResults(BaseModel):
    files_created: int = 0
    workdir_path: str = ""


def test_workdir_creation_uses_expected_path() -> None:
    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(Path(temp_dir))
        workdir = context.get_task_workdir("task-001")

        expected = Path(temp_dir) / "task-001" / "workdir"
        assert Path(workdir) == expected, (workdir, expected)
        assert expected.exists()
        assert expected.is_dir()


def test_task_workdirs_are_isolated() -> None:
    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(Path(temp_dir))
        task_ids = ["task-001", "task-002", "task-003"]
        workdirs = {task_id: Path(context.get_task_workdir(task_id)) for task_id in task_ids}

        for task_id, workdir in workdirs.items():
            marker = workdir / f"{task_id}.txt"
            marker.write_text(f"marker for {task_id}", encoding="utf-8")

        assert len(set(workdirs.values())) == len(task_ids)
        for task_id, workdir in workdirs.items():
            marker = workdir / f"{task_id}.txt"
            assert marker.read_text(encoding="utf-8") == f"marker for {task_id}"


def test_workdir_supports_nested_files() -> None:
    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(Path(temp_dir))
        workdir = Path(context.get_task_workdir("file-test"))
        nested = workdir / "subdir"
        nested.mkdir()

        created_files = [
            workdir / "test1.txt",
            workdir / "test2.json",
            nested / "nested.txt",
        ]
        for file_path in created_files:
            file_path.write_text(f"content for {file_path.name}", encoding="utf-8")

        assert all(file_path.exists() for file_path in created_files)
        assert (nested / "nested.txt").read_text(encoding="utf-8") == "content for nested.txt"


def test_get_task_workdir_is_idempotent() -> None:
    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(Path(temp_dir))
        first = context.get_task_workdir("repeat-task")
        second = context.get_task_workdir("repeat-task")
        third = context.get_task_workdir("repeat-task")

        assert first == second == third
        assert Path(first).exists()


def test_batch_style_usage_can_write_task_local_files() -> None:
    def mock_batch_processor(
        batch_inputs: List[TaskInput[ContextOptions]],
        operator_config: dict,
        context: ProcessingContext,
    ) -> Generator[TaskOutput[ContextResults], None, None]:
        for task_input in batch_inputs:
            workdir = Path(context.get_task_workdir(task_input.task_id))
            files_created = 0

            config_file = workdir / "config.json"
            config_file.write_text(
                json.dumps(
                    {
                        "task_id": task_input.task_id,
                        "test_param": task_input.options.test_param,
                    },
                    indent=2,
                ),
                encoding="utf-8",
            )
            files_created += 1

            log_file = workdir / "processing.log"
            log_file.write_text(f"processed {task_input.task_id}\n", encoding="utf-8")
            files_created += 1

            for index, artifact in enumerate(task_input.artifacts or []):
                temp_file = workdir / f"processed_{index}.tmp"
                temp_file.write_text(f"processed {artifact.uri}", encoding="utf-8")
                files_created += 1

            yield TaskOutput[ContextResults](
                task_id=task_input.task_id,
                results=ContextResults(
                    files_created=files_created,
                    workdir_path=str(workdir),
                ),
                status="success",
            )

    with tempfile.TemporaryDirectory() as temp_dir:
        context = ProcessingContext(Path(temp_dir))
        batch_inputs = [
            TaskInput[ContextOptions](
                task_id="batch-001",
                options=ContextOptions(test_param="param-1"),
                artifacts=[
                    InputArtifact(uri="file1.txt"),
                    InputArtifact(uri="file2.txt"),
                ],
            ),
            TaskInput[ContextOptions](
                task_id="batch-002",
                options=ContextOptions(test_param="param-2"),
                artifacts=[InputArtifact(uri="file3.txt")],
            ),
        ]

        results = list(mock_batch_processor(batch_inputs, {}, context))

        assert len(results) == 2
        for result in results:
            assert result.status == "success"
            assert result.results is not None
            assert result.results.files_created >= 3
            workdir = Path(result.results.workdir_path)
            assert workdir.exists()
            assert workdir.parent.exists()


def main() -> int:
    tests = [
        ("workdir_creation", test_workdir_creation_uses_expected_path),
        ("isolation", test_task_workdirs_are_isolated),
        ("nested_files", test_workdir_supports_nested_files),
        ("idempotent_lookup", test_get_task_workdir_is_idempotent),
        ("batch_usage", test_batch_style_usage_can_write_task_local_files),
    ]

    passed = 0
    failed = 0

    print("ProcessingContext 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())
