# Copyright 2015 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

import itertools
import time
import unittest
from dataclasses import dataclass, field
from textwrap import dedent
from typing import List, Optional, Tuple

import pytest

from pants.backend.python.target_types import PythonSourcesGeneratorTarget
from pants.base.build_environment import get_buildroot
from pants.base.specs import Specs
from pants.base.specs_parser import SpecsParser
from pants.engine.engine_aware import EngineAwareParameter, EngineAwareReturnType
from pants.engine.fs import (
    EMPTY_FILE_DIGEST,
    EMPTY_SNAPSHOT,
    CreateDigest,
    Digest,
    DigestContents,
    FileContent,
    Snapshot,
)
from pants.engine.internals.engine_testutil import (
    assert_equal_with_printing,
    remove_locations_from_traceback,
)
from pants.engine.internals.scheduler import ExecutionError, SchedulerSession
from pants.engine.internals.scheduler_test_base import SchedulerTestBase
from pants.engine.process import Process, ProcessCacheScope, ProcessResult
from pants.engine.rules import Get, MultiGet, rule
from pants.engine.streaming_workunit_handler import (
    StreamingWorkunitContext,
    StreamingWorkunitHandler,
    TargetInfo,
    WorkunitsCallback,
)
from pants.goal.run_tracker import RunTracker
from pants.testutil.option_util import create_options_bootstrapper
from pants.testutil.rule_runner import QueryRule, RuleRunner
from pants.util.logging import LogLevel


class A:
    pass


class B:
    pass


class C:
    pass


class D:
    pass


def fn_raises(x):
    raise Exception(f"An exception for {type(x).__name__}")


@rule
def nested_raise(x: B) -> A:  # type: ignore[return]
    fn_raises(x)


@dataclass(frozen=True)
class Fib:
    val: int


@rule(desc="Fibonacci", level=LogLevel.INFO)
async def fib(n: int) -> Fib:
    if n < 2:
        return Fib(n)
    x, y = tuple(await MultiGet([Get(Fib, int(n - 2)), Get(Fib, int(n - 1))]))
    return Fib(x.val + y.val)


@dataclass(frozen=True)
class MyInt:
    val: int


@dataclass(frozen=True)
class MyFloat:
    val: float


@rule
def upcast(n: MyInt) -> MyFloat:
    return MyFloat(float(n.val))


# This set of dummy types and the following `@rule`s are intended to test that workunits are
# being generated correctly and with the correct parent-child relationships.


class Input:
    pass


class Alpha:
    pass


class Beta:
    pass


class Gamma:
    pass


class Omega:
    pass


class Epsilon:
    pass


@rule(canonical_name="canonical_rule_one", desc="Rule number 1", level=LogLevel.INFO)
async def rule_one_function(i: Input) -> Beta:
    """This rule should be the first one executed by the engine, and thus have no parent."""
    a = Alpha()
    o = await Get(Omega, Alpha, a)
    b = await Get(Beta, Omega, o)
    time.sleep(1)
    return b


@rule(desc="Rule number 2", level=LogLevel.INFO)
async def rule_two(a: Alpha) -> Omega:
    """This rule should be invoked in the body of `rule_one` and therefore its workunit should be a
    child of `rule_one`'s workunit."""
    await Get(Gamma, Alpha, a)
    return Omega()


@rule(desc="Rule number 3", level=LogLevel.INFO)
async def rule_three(o: Omega) -> Beta:
    """This rule should be invoked in the body of `rule_one` and therefore its workunit should be a
    child of `rule_one`'s workunit."""
    return Beta()


@rule(desc="Rule number 4", level=LogLevel.INFO)
def rule_four(a: Alpha) -> Gamma:
    """This rule should be invoked in the body of `rule_two` and therefore its workunit should be a
    child of `rule_two`'s workunit."""
    return Gamma()


@rule(desc="Rule A", level=LogLevel.INFO)
async def rule_A(i: Input) -> Alpha:
    o = Omega()
    a = await Get(Alpha, Omega, o)
    return a


@rule
async def rule_B(o: Omega) -> Alpha:
    e = Epsilon()
    a = await Get(Alpha, Epsilon, e)
    return a


@rule(desc="Rule C", level=LogLevel.INFO)
def rule_C(e: Epsilon) -> Alpha:
    return Alpha()


class EngineTest(unittest.TestCase, SchedulerTestBase):

    assert_equal_with_printing = assert_equal_with_printing

    def scheduler(self, rules, include_trace_on_error):
        return self.mk_scheduler(rules=rules, include_trace_on_error=include_trace_on_error)

    def test_recursive_multi_get(self):
        # Tests that a rule that "uses itself" multiple times per invoke works.
        rules = [fib, QueryRule(Fib, (int,))]
        (fib_10,) = self.mk_scheduler(rules=rules).product_request(Fib, subjects=[10])
        self.assertEqual(55, fib_10.val)

    def test_no_include_trace_error_raises_boring_error(self):
        rules = [nested_raise, QueryRule(A, (B,))]
        scheduler = self.scheduler(rules, include_trace_on_error=False)
        with self.assertRaises(ExecutionError) as cm:
            list(scheduler.product_request(A, subjects=[(B())]))
        self.assert_equal_with_printing(
            "1 Exception encountered:\n\n  Exception: An exception for B\n", str(cm.exception)
        )

    def test_no_include_trace_error_multiple_paths_raises_executionerror(self):
        rules = [nested_raise, QueryRule(A, (B,))]
        scheduler = self.scheduler(rules, include_trace_on_error=False)
        with self.assertRaises(ExecutionError) as cm:
            list(scheduler.product_request(A, subjects=[B(), B()]))
        self.assert_equal_with_printing(
            dedent(
                """
                2 Exceptions encountered:

                  Exception: An exception for B
                  Exception: An exception for B
                """
            ).lstrip(),
            str(cm.exception),
        )

    def test_include_trace_error_raises_error_with_trace(self):
        rules = [nested_raise, QueryRule(A, (B,))]
        scheduler = self.scheduler(rules, include_trace_on_error=True)
        with self.assertRaises(ExecutionError) as cm:
            list(scheduler.product_request(A, subjects=[(B())]))
        self.assert_equal_with_printing(
            dedent(
                """
                1 Exception encountered:

                Engine traceback:
                  in select
                  in pants.engine.internals.engine_test.nested_raise
                Traceback (most recent call last):
                  File LOCATION-INFO, in nested_raise
                    fn_raises(x)
                  File LOCATION-INFO, in fn_raises
                    raise Exception(f"An exception for {type(x).__name__}")
                Exception: An exception for B
                """
            ).lstrip(),
            remove_locations_from_traceback(str(cm.exception)),
        )

    def test_nonexistent_root(self) -> None:
        rules = [QueryRule(A, [B])]
        # No rules are available to compute A.
        with self.assertRaises(ValueError) as cm:
            self.scheduler(rules, include_trace_on_error=False)
        assert (
            "No installed rules return the type A, and it was not provided by potential callers of "
        ) in str(cm.exception)

    def test_missing_query_rule(self) -> None:
        # Even if we register the rule to go from MyInt -> MyFloat, we must register a QueryRule
        # for the graph to work when making a synchronous call via `Scheduler.product_request`.
        scheduler = self.mk_scheduler(rules=[upcast], include_trace_on_error=False)
        with self.assertRaises(Exception) as cm:
            scheduler.product_request(MyFloat, subjects=[MyInt(0)])
        assert (
            "No installed QueryRules return the type MyFloat. Try registering QueryRule(MyFloat "
            "for MyInt)."
        ) in str(cm.exception)


@dataclass
class WorkunitTracker(WorkunitsCallback):
    """This class records every non-empty batch of started and completed workunits received from the
    engine."""

    finished_workunit_chunks: List[List[dict]] = field(default_factory=list)
    started_workunit_chunks: List[List[dict]] = field(default_factory=list)
    finished: bool = False

    @property
    def can_finish_async(self) -> bool:
        return False

    def __call__(self, **kwargs) -> None:
        if kwargs["finished"] is True:
            self.finished = True

        started_workunits = kwargs.get("started_workunits")
        if started_workunits:
            self.started_workunit_chunks.append(started_workunits)

        completed_workunits = kwargs.get("completed_workunits")
        if completed_workunits: