# =====================================================
# Apckeyl Framework
# Version 1.0
# test_compute_module.py
# =====================================================

"""
Tests for Apckeyl Compute Module Contract.

Version 1.0
"""

from compute_module import (
    ComputeModule,
    ExampleComputeModule,
)


# =====================================================
# Test 1 — Base Module Exists
# =====================================================

def test_base_module_exists():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[
            "test_task"
        ],
    )

    assert module is not None


# =====================================================
# Test 2 — Module Information
# =====================================================

def test_module_info():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[
            "test_task"
        ],
    )

    info = module.get_info()

    assert info["module_id"] == (
        "test_module"
    )

    assert info["module_name"] == (
        "Test Module"
    )

    assert info["module_version"] == (
        "1.0"
    )

    assert info["supported_tasks"] == [
        "test_task"
    ]


# =====================================================
# Test 3 — Supported Task
# =====================================================

def test_supported_task():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[
            "image_upscale"
        ],
    )

    assert module.supports_task(
        "image_upscale"
    )


# =====================================================
# Test 4 — Unsupported Task
# =====================================================

def test_unsupported_task():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[
            "image_upscale"
        ],
    )

    assert not module.supports_task(
        "image_generation"
    )


# =====================================================
# Test 5 — Health Check
# =====================================================

def test_health_check():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[],
    )

    health = module.health_check()

    assert health["module_id"] == (
        "test_module"
    )

    assert health["status"] == (
        "healthy"
    )


# =====================================================
# Test 6 — None Task Rejected
# =====================================================

def test_none_task_rejected():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[
            "test_task"
        ],
    )

    try:

        module.execute(None)

    except ValueError:

        return

    raise AssertionError(
        "None task was not rejected"
    )


# =====================================================
# Test 7 — Invalid Task Type Rejected
# =====================================================

def test_invalid_task_type_rejected():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[
            "test_task"
        ],
    )

    try:

        module.execute(
            "invalid_task"
        )

    except TypeError:

        return

    raise AssertionError(
        "Invalid task type was not rejected"
    )


# =====================================================
# Test 8 — Missing Task ID Rejected
# =====================================================

def test_missing_task_id_rejected():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[
            "test_task"
        ],
    )

    task = {

        "task_type": "test_task",

    }

    try:

        module.execute(task)

    except ValueError:

        return

    raise AssertionError(
        "Task without task_id was not rejected"
    )


# =====================================================
# Test 9 — Unsupported Task Rejected
# =====================================================

def test_unsupported_task_rejected():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[
            "test_task"
        ],
    )

    task = {

        "task_id": "task_001",

        "task_type": (
            "unsupported_task"
        ),

    }

    try:

        module.execute(task)

    except ValueError:

        return

    raise AssertionError(
        "Unsupported task was not rejected"
    )


# =====================================================
# Test 10 — Base Execute Is Abstract
# =====================================================

def test_base_execute_not_implemented():

    module = ComputeModule(
        module_id="test_module",
        module_name="Test Module",
        module_version="1.0",
        supported_tasks=[
            "test_task"
        ],
    )

    task = {

        "task_id": "task_002",

        "task_type": "test_task",

    }

    try:

        module.execute(task)

    except NotImplementedError:

        return

    raise AssertionError(
        "Base execute() was not protected"
    )


# =====================================================
# Test 11 — Example Module Exists
# =====================================================

def test_example_module_exists():

    module = ExampleComputeModule()

    assert module is not None


# =====================================================
# Test 12 — Example Module Information
# =====================================================

def test_example_module_info():

    module = ExampleComputeModule()

    info = module.get_info()

    assert info["module_id"] == (
        "example"
    )

    assert info["module_name"] == (
        "Example Compute Module"
    )

    assert info["module_version"] == (
        "1.0"
    )

    assert "example_task" in (
        info["supported_tasks"]
    )


# =====================================================
# Test 13 — Example Module Health
# =====================================================

def test_example_module_health():

    module = ExampleComputeModule()

    health = module.health_check()

    assert health["module_id"] == (
        "example"
    )

    assert health["status"] == (
        "healthy"
    )


# =====================================================
# Test 14 — Example Module Execution
# =====================================================

def test_example_module_execution():

    module = ExampleComputeModule()

    task = {

        "task_id": "example_001",

        "task_type": "example_task",

        "payload": {

            "message": "hello",

        },
    }

    result = module.execute(
        task
    )

    assert result["task_id"] == (
        "example_001"
    )

    assert result["status"] == (
        "completed"
    )

    assert "result" in result


# =====================================================
# Test Runner
# =====================================================

def main():

    tests = [

        test_base_module_exists,

        test_module_info,

        test_supported_task,

        test_unsupported_task,

        test_health_check,

        test_none_task_rejected,

        test_invalid_task_type_rejected,

        test_missing_task_id_rejected,

        test_unsupported_task_rejected,

        test_base_execute_not_implemented,

        test_example_module_exists,

        test_example_module_info,

        test_example_module_health,

        test_example_module_execution,

    ]

    passed = 0

    failed = 0

    print(
        "===================================="
    )

    print(
        "Apckeyl Compute Module Test"
    )

    print(
        "Version 1.0"
    )

    print(
        "===================================="
    )

    for test in tests:

        try:

            test()

            print(
                f"PASS: {test.__name__}"
            )

            passed += 1

        except Exception as error:

            print(
                f"FAIL: {test.__name__}"
            )

            print(
                f"ERROR: {error}"
            )

            failed += 1

    print(
        "===================================="
    )

    print(
        f"PASSED: {passed}"
    )

    print(
        f"FAILED: {failed}"
    )

    print(
        "===================================="
    )

    if failed > 0:

        raise RuntimeError(
            "Compute Module Test FAILED"
        )

    print(
        "COMPUTE MODULE TEST: PASSED"
    )


# =====================================================
# Entry Point
# =====================================================

if __name__ == "__main__":

    main()