import asyncio
import os
import random
import string
from typing import Callable, Dict, Optional

import pytest
import pytest_asyncio

from e2b import (
    AsyncCommandHandle,
    AsyncSandbox,
    AsyncTemplate,
    AsyncVolume,
    CommandExitException,
    CommandHandle,
    LogEntry,
    Sandbox,
    Template,
    TemplateClass,
    Volume,
)


@pytest.fixture
def test_api_key() -> str:
    """Placeholder API key with a valid format for tests that don't hit the API."""
    return "e2b_" + "0" * 40


@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    outcome = yield
    rep = outcome.get_result()
    if rep.when == "call":
        item._test_failed = rep.failed


@pytest.fixture()
def sandbox_test_id():
    return f"test_{_generate_random_string()}"


@pytest.fixture()
def template():
    return "base"


@pytest.fixture()
def sandbox_factory(request, template, sandbox_test_id):
    def factory(*, template_name: str = template, **kwargs):
        metadata = kwargs.setdefault("metadata", dict())
        metadata.setdefault("sandbox_test_id", sandbox_test_id)

        sandbox = Sandbox.create(template_name, **kwargs)

        def finalizer():
            if getattr(request.node, "_test_failed", False):
                print(f"\n[TEST FAILED] Sandbox ID: {sandbox.sandbox_id}")
            sandbox.kill()

        request.addfinalizer(finalizer)

        return sandbox

    return factory


@pytest.fixture()
def sandbox(sandbox_factory):
    return sandbox_factory()


@pytest_asyncio.fixture
async def async_sandbox_factory(request, template, sandbox_test_id):
    sandboxes: list = []

    async def factory(*, template_name: str = template, **kwargs):
        metadata = kwargs.setdefault("metadata", dict())
        metadata.setdefault("sandbox_test_id", sandbox_test_id)

        sandbox = await AsyncSandbox.create(template_name, **kwargs)
        sandboxes.append(sandbox)
        return sandbox

    yield factory

    if getattr(request.node, "_test_failed", False):
        for sandbox in sandboxes:
            print(f"\n[TEST FAILED] Sandbox ID: {sandbox.sandbox_id}")

    results = await asyncio.gather(
        *(sandbox.kill() for sandbox in sandboxes), return_exceptions=True
    )
    for sandbox, result in zip(sandboxes, results):
        if isinstance(result, BaseException):
            print(f"\n[TEARDOWN FAILED] Sandbox ID: {sandbox.sandbox_id}: {result!r}")


@pytest_asyncio.fixture
async def async_sandbox(async_sandbox_factory):
    return await async_sandbox_factory()


@pytest.fixture
def build():
    def _build(
        template: TemplateClass,
        name: Optional[str] = None,
        skip_cache: bool = False,
        on_build_logs: Optional[Callable[[LogEntry], None]] = None,
    ):
        build_name = name or f"e2b-test-{_generate_random_string()}"
        build_info: Dict[str, Optional[str]] = {"template_id": None, "build_id": None}

        def capture_logs(log: LogEntry):
            import re

            if "Template created with ID:" in log.message:
                match = re.search(
                    r"Template created with ID: ([^,]+), Build ID: (.+)", log.message
                )
                if match:
                    build_info["template_id"] = match.group(1)
                    build_info["build_id"] = match.group(2)
            if on_build_logs:
                on_build_logs(log)

        try:
            return Template.build(
                template,
                build_name,
                cpu_count=1,
                memory_mb=1024,
                skip_cache=skip_cache,
                on_build_logs=capture_logs,
            )
        except Exception as e:
            print(
                f"\n[BUILD FAILED] name={build_name}, "
                f"template_id={build_info['template_id']}, "
                f"build_id={build_info['build_id']}, error={e}"
            )
            raise

    return _build


@pytest_asyncio.fixture
def async_build():
    async def _async_build(
        template: TemplateClass,
        name: Optional[str] = None,
        skip_cache: bool = False,
        on_build_logs: Optional[Callable[[LogEntry], None]] = None,
    ):
        build_name = name or f"e2b-test-{_generate_random_string()}"
        build_info: Dict[str, Optional[str]] = {"template_id": None, "build_id": None}

        def capture_logs(log: LogEntry):
            import re

            if "Template created with ID:" in log.message:
                match = re.search(
                    r"Template created with ID: ([^,]+), Build ID: (.+)", log.message
                )
                if match:
                    build_info["template_id"] = match.group(1)
                    build_info["build_id"] = match.group(2)
            if on_build_logs:
                on_build_logs(log)

        try:
            return await AsyncTemplate.build(
                template,
                build_name,
                cpu_count=1,
                memory_mb=1024,
                skip_cache=skip_cache,
                on_build_logs=capture_logs,
            )
        except Exception as e:
            print(
                f"\n[BUILD FAILED] name={build_name}, "
                f"template_id={build_info['template_id']}, "
                f"build_id={build_info['build_id']}, error={e}"
            )
            raise

    return _async_build


@pytest.fixture
def debug():
    return os.getenv("E2B_DEBUG") is not None


@pytest.fixture(autouse=True)
def skip_by_debug(request, debug):
    if request.node.get_closest_marker("skip_debug"):
        if debug:
            pytest.skip("skipped because E2B_DEBUG is set")


class Helpers:
    @staticmethod
    def catch_cmd_exit_error_in_background(cmd: AsyncCommandHandle):
        disabled = False

        async def wait_for_exit():
            try:
                await cmd.wait()
            except CommandExitException as e:
                if not disabled:
                    assert False, (
                        f"command failed with exit code {e.exit_code}: {e.stderr}"
                    )

        asyncio.create_task(wait_for_exit())

        def disable():
            nonlocal disabled
            disabled = True

        return disable

    @staticmethod
    def check_cmd_exit_error(cmd: CommandHandle):
        try:
            cmd.wait()
        except CommandExitException as e:
            assert False, f"command failed with exit code {e.exit_code}: {e.stderr}"
        except Exception as e:
            raise e


@pytest.fixture
def helpers():
    return Helpers


def _generate_random_string(length: int = 8) -> str:
    return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))


def _skip_unless_volume_tests_enabled():
    if os.getenv("ENABLE_VOLUME_TESTS") is None:
        pytest.skip("skipped because ENABLE_VOLUME_TESTS is not set")


@pytest.fixture
def volume(request):
    _skip_unless_volume_tests_enabled()
    vol = Volume.create(f"test-vol-{_generate_random_string()}")

    def finalizer():
        if getattr(request.node, "_test_failed", False):
            print(f"\n[TEST FAILED] Volume ID: {vol.volume_id}")
        try:
            Volume.destroy(vol.volume_id)
        except Exception:
            pass

    request.addfinalizer(finalizer)
    return vol


@pytest_asyncio.fixture
async def async_volume(request):
    _skip_unless_volume_tests_enabled()
    vol = await AsyncVolume.create(f"test-vol-{_generate_random_string()}")
    try:
        yield vol
    finally:
        if getattr(request.node, "_test_failed", False):
            print(f"\n[TEST FAILED] Volume ID: {vol.volume_id}")
        try:
            await AsyncVolume.destroy(vol.volume_id)
        except Exception:
            pass