import sys
from collections.abc import Iterator
from contextlib import ExitStack, contextmanager

import dagster as dg
from dagster._core.instance import DagsterInstance
from dagster._core.remote_origin import ManagedGrpcPythonEnvCodeLocationOrigin
from dagster._core.remote_representation.code_location import GrpcServerCodeLocation
from dagster._core.remote_representation.handle import JobHandle, RepositoryHandle
from dagster._core.types.loadable_target_origin import LoadableTargetOrigin
from dagster._core.workspace.context import WorkspaceProcessContext, WorkspaceRequestContext
from dagster._core.workspace.load_target import PythonFileTarget
from dagster_shared.record import as_dict


@contextmanager
def get_bar_workspace(instance: DagsterInstance) -> Iterator[WorkspaceRequestContext]:
    with WorkspaceProcessContext(
        instance,
        PythonFileTarget(
            python_file=dg.file_relative_path(__file__, "api_tests_repo.py"),
            attribute="bar_repo",
            working_directory=None,
            location_name="bar_code_location",
        ),
    ) as workspace_process_context:
        yield workspace_process_context.create_request_context()


@contextmanager
def get_workspace(
    instance: DagsterInstance,
    python_file: str,
    attribute: str | None,
    location_name: str,
) -> Iterator[WorkspaceRequestContext]:
    with WorkspaceProcessContext(
        instance,
        PythonFileTarget(
            python_file=python_file,
            attribute=attribute,
            working_directory=None,
            location_name=location_name,
        ),
    ) as workspace_process_context:
        yield workspace_process_context.create_request_context()


@contextmanager
def get_bar_repo_code_location(
    instance: dg.DagsterInstance | None = None,
) -> Iterator[GrpcServerCodeLocation]:
    with ExitStack() as stack:
        if not instance:
            instance = stack.enter_context(dg.instance_for_test())

        loadable_target_origin = LoadableTargetOrigin(
            executable_path=sys.executable,
            python_file=dg.file_relative_path(__file__, "api_tests_repo.py"),
            attribute="bar_repo",
        )
        location_name = "bar_code_location"

        origin = ManagedGrpcPythonEnvCodeLocationOrigin(loadable_target_origin, location_name)

        with origin.create_single_location(instance) as location:
            yield location


@contextmanager
def get_code_location(
    python_file: str,
    attribute: str,
    location_name: str,
    instance: dg.DagsterInstance | None = None,
) -> Iterator[GrpcServerCodeLocation]:
    with ExitStack() as stack:
        if not instance:
            instance = stack.enter_context(dg.instance_for_test())

        loadable_target_origin = LoadableTargetOrigin(
            executable_path=sys.executable,
            python_file=python_file,
            attribute=attribute,
        )
        origin = ManagedGrpcPythonEnvCodeLocationOrigin(loadable_target_origin, location_name)

        with origin.create_single_location(instance) as location:
            yield location


@contextmanager
def get_bar_repo_handle(
    instance: dg.DagsterInstance | None = None,
) -> Iterator[RepositoryHandle]:
    with ExitStack() as stack:
        if not instance:
            instance = stack.enter_context(dg.instance_for_test())

        with get_bar_repo_code_location(instance) as location:
            yield location.get_repository("bar_repo").handle


@contextmanager
def get_foo_job_handle(instance: dg.DagsterInstance | None = None) -> Iterator[JobHandle]:
    with ExitStack() as stack:
        if not instance:
            instance = stack.enter_context(dg.instance_for_test())

        with get_bar_repo_handle(instance) as repo_handle:
            yield JobHandle("foo", repo_handle)


def with_invalid_origin(args):
    """Bypass record type checking to create an invalid object for testing deserialization errors."""
    return args.__class__.__nt_new__(
        args.__class__,
        **{**as_dict(args), "repository_origin": "INVALID"},
    )
