from dataclasses import dataclass
from typing import Any, Iterable, cast
from pants.backend.python.target_types import ConsoleScript
from pants.backend.python.subsystems.python_tool_base import PythonToolBase
from pants.backend.python.target_types import (
    InterpreterConstraintsField,
    PythonResolveField,
    PythonTestsBatchCompatibilityTagField,
    PythonTestsExtraEnvVarsField,
    PythonTestSourceField,
    PythonTestsTimeoutField,
    PythonTestsXdistConcurrencyField,
    RuntimePackageDependenciesField,
    SkipPythonTestsField,
)
from pants.engine.process import FallibleProcessResult, ProcessExecutionFailure
from pants.backend.python.util_rules.pex import Pex, PexRequest, VenvPex, VenvPexProcess
from pants.backend.python.util_rules.pex_from_targets import RequirementsPexRequest
from pants.backend.python.util_rules.python_sources import (
    PythonSourceFiles,
    PythonSourceFilesRequest,
)
from pants.core.goals.test import TestFieldSet, TestRequest, TestResult, TestSubsystem
from pants.core.util_rules.environments import EnvironmentField
from pants.engine.process import Process, ProcessResultWithRetries, ProcessWithRetries
from pants.engine.rules import (
    Get,
    Rule,
    collect_rules,
    rule,  # type: ignore
)
from pants.engine.target import (
    COMMON_TARGET_FIELDS,
    Dependencies,
    Target,
    TransitiveTargets,
    TransitiveTargetsRequest,
)
from pants.option.global_options import GlobalOptions
from pants.option.option_types import SkipOption
from pants.util.logging import LogLevel

@dataclass(frozen=True)
class NoseTestFieldSet(TestFieldSet):
    required_fields = (PythonTestSourceField,)

    source: PythonTestSourceField
    interpreter_constraints: InterpreterConstraintsField
    timeout: PythonTestsTimeoutField
    runtime_package_dependencies: RuntimePackageDependenciesField
    extra_env_vars: PythonTestsExtraEnvVarsField
    resolve: PythonResolveField
    environment: EnvironmentField
    xdist_concurrency: PythonTestsXdistConcurrencyField
    batch_compatibility_tag: PythonTestsBatchCompatibilityTagField

    @classmethod
    def opt_out(cls, tgt: Target) -> bool:
        return tgt.get(SkipPythonTestsField).value


class NoseTest(Target):
    alias = "nose_test"
    help = "Run Nose (Python) tests."
    core_fields = (
        *COMMON_TARGET_FIELDS,
        Dependencies,
        PythonTestSourceField,
        PythonResolveField,
    )


class NoseTool(PythonToolBase):
    name = "Nose"
    options_scope = "nose"
    help = "The Nose test runner (and framework) for Python (https://nose.readthedocs.io/en/latest/)."

    register_interpreter_constraints = True
    default_lockfile_resource = ("experimental.nose", "nose.lock",)

    default_version = ["pynose==1.5.1"]
    default_main = ConsoleScript("nosetests")

    skip = SkipOption("test")

@dataclass(frozen=True)
class NoseRequest(TestRequest):
    tool_subsystem = NoseTool
    field_set_type = NoseTestFieldSet


@rule(desc="Run Python test(s) with Nose")  # type: ignore
async def run_nose_test(
    batch: NoseRequest.Batch[NoseTestFieldSet, Any],
    nose: NoseTool,
    test_subsystem: TestSubsystem,
    global_options: GlobalOptions,
) -> TestResult:
    
    transitive_targets = await Get(
        TransitiveTargets, TransitiveTargetsRequest((batch.single_element.address,))
    )

    nose_runner_pex = await Get(VenvPex, PexRequest, nose.to_pex_request())
    sources = await Get(
        PythonSourceFiles,
        PythonSourceFilesRequest(transitive_targets.closure, include_files=True),
    )

    results = await Get(
        FallibleProcessResult,
        VenvPexProcess(
            nose_runner_pex,
            description=f"Run Nose for {batch.single_element.address}",
            argv=sources.source_files.files,
            input_digest=sources.source_files.snapshot.digest,
            level=LogLevel.DEBUG,
        ),
    )

    # TODO: Just commenting this out, as I can't really do much with this, with a broken LSP

    # results = await Get(
    #     ProcessResultWithRetries,
    #     ProcessWithRetries(process, test_subsystem.attempts_default),
    # )

    return TestResult.from_batched_fallible_process_result(
        (results,),
        batch=batch,
        output_setting=test_subsystem.output,
        output_simplifier=global_options.output_simplifier(),
    )


def target_types() -> Iterable[type[Target]]:
    return [NoseTest]


def rules() -> Iterable[Rule]:
    return [
        *collect_rules(),
        *cast(Iterable[Rule], NoseRequest.rules()),  # type: ignore
        *NoseTool.rules(),
    ]