import os
import threading
import time
from concurrent.futures import TimeoutError
from typing import Any, List, Optional

import pytest
from attrs import define, field
from hypothesis import given, settings
from hypothesis import strategies as st
from hypothesis.stateful import RuleBasedStateMachine, precondition, rule

from cog.server.eventtypes import Done, Log, PredictionOutput, PredictionOutputType
from cog.server.exceptions import FatalWorkerException, InvalidStateException
from cog.server.worker import Worker

# Set a longer deadline on CI as the instances are a bit slower.
settings.register_profile("ci", max_examples=100, deadline=2000)
settings.register_profile("default", max_examples=10, deadline=1500)
settings.register_profile("slow", max_examples=10, deadline=2000)
settings.load_profile(os.getenv("HYPOTHESIS_PROFILE", "default"))

ST_NAMES = st.sampled_from(["John", "Barry", "Elspeth", "Hamid", "Ronnie", "Yasmeen"])

SETUP_FATAL_FIXTURES = [
    ("exc_in_setup", {}),
    ("exc_in_setup_and_predict", {}),
    ("exc_on_import", {}),
    ("exit_in_setup", {}),
    ("exit_on_import", {}),
    ("missing_predictor", {}),
    ("nonexistent_file", {}),  # this fixture doesn't even exist
]

PREDICTION_FATAL_FIXTURES = [
    ("exit_in_predict", {}),
    ("killed_in_predict", {}),
]

RUNNABLE_FIXTURES = [
    ("simple", {}),
    ("exc_in_predict", {}),
    ("missing_predict", {}),
]

OUTPUT_FIXTURES = [
    (
        "hello_world",
        {"name": ST_NAMES},
        lambda x: f"hello, {x['name']}",
    ),
    (
        "count_up",
        {"upto": st.integers(min_value=0, max_value=100)},
        lambda x: list(range(x["upto"])),
    ),
    ("complex_output", {}, lambda _: {"number": 42, "text": "meaning of life"}),
]

SETUP_LOGS_FIXTURES = [
    (
        "logging",
        (
            "writing some stuff from C at import time\n"
            "writing to stdout at import time\n"
            "setting up predictor\n"
        ),
        "writing to stderr at import time\n",
    )
]

PREDICT_LOGS_FIXTURES = [
    (
        "logging",
        {},
        ("writing from C\n" "writing with print\n"),
        ("WARNING:root:writing log message\n" "writing to stderr\n"),
    )
]


@define
class Result:
    stdout_lines: List[str] = field(factory=list)
    stderr_lines: List[str] = field(factory=list)
    heartbeat_count: int = 0
    output_type: Optional[PredictionOutputType] = None
    output: Any = None
    done: Optional[Done] = None
    exception: Optional[Exception] = None

    @property
    def stdout(self):
        return "".join(self.stdout_lines)

    @property
    def stderr(self):
        return "".join(self.stderr_lines)


def _handle_event(result, event):
    if isinstance(event, Log) and event.source == "stdout":
        result.stdout_lines.append(event.message)
    elif isinstance(event, Log) and event.source == "stderr":
        result.stderr_lines.append(event.message)
    elif isinstance(event, Done):
        assert not result.done
        result.done = event
    elif isinstance(event, PredictionOutput):
        assert result.output_type, "Should get output type before any output"
        if result.output_type.multi:
            result.output.append(event.payload)
        else:
            assert (
                result.output is None
            ), "Should not get multiple outputs for output type single"
            result.output = event.payload
    elif isinstance(event, PredictionOutputType):
        assert result.output_type is None, "Should not get multiple output type events"
        result.output_type = event
        if result.output_type.multi:
            result.output = []
    else:
        pytest.fail(f"saw unexpected event: {event}")


def _process(worker, work, swallow_exceptions=False):
    """
    Helper function to collect events generated by Worker during tests.
    """
    result = Result()
    subid = worker.subscribe(lambda event: _handle_event(result, event))
    try:
        work().result()
    except Exception as exc:
        result.exception = exc
        if not swallow_exceptions:
            raise
    finally:
        worker.unsubscribe(subid)
    return result


def _fixture_path(name):
    test_dir = os.path.dirname(os.path.realpath(__file__))
    return os.path.join(test_dir, f"fixtures/{name}.py") + ":Predictor"


@pytest.mark.parametrize("name,payloads", SETUP_FATAL_FIXTURES)
def test_fatalworkerexception_from_setup_failures(name, payloads):
    """
    Any failure during setup is fatal and should raise FatalWorkerException.
    """
    w = Worker(predictor_ref=_fixture_path(name), tee_output=False)

    with pytest.raises(FatalWorkerException):
        _process(w, w.setup)

    w.terminate()


@pytest.mark.parametrize("name,payloads", PREDICTION_FATAL_FIXTURES)
@given(data=st.data())
def test_fatalworkerexception_from_irrecoverable_failures(data, name, payloads):
    """
    Certain kinds of failure during predict (crashes, unexpected exits) are
    irrecoverable and should raise FatalWorkerException.
    """
    w = Worker(predictor_ref=_fixture_path(name), tee_output=False)

    result = _process(w, w.setup)
    assert not result.done.error

    with pytest.raises(FatalWorkerException):
        _process(w, lambda: w.predict(data.draw(st.fixed_dictionaries(payloads))))

    with pytest.raises(InvalidStateException):
        _process(w, lambda: w.predict(data.draw(st.fixed_dictionaries(payloads))))

    w.terminate()


@pytest.mark.parametrize("name,payloads", RUNNABLE_FIXTURES)
@given(data=st.data())
def test_no_exceptions_from_recoverable_failures(data, name, payloads):
    """
    Well-behaved predictors, or those that only throw exceptions, should not
    raise.
    """
    w = Worker(predictor_ref=_fixture_path(name), tee_output=False)

    try:
        result = _process(w, w.setup)
        assert not result.done.error

        for _ in range(5):
            _process(w, lambda: w.predict(data.draw(st.fixed_dictionaries(payloads))))
    finally:
        w.terminate()


def test_stream_redirector_race_condition():
    """
    StreamRedirector and _ChildWorker are using the same _events pipe to send
    data. When there are multiple threads trying to write to the same pipe, it
    can cause data corruption by race condition. The data corruption will cause
    pipe receiver to raise an exception due to unpickling error.
    """
    w = Worker(
        predictor_ref=_fixture_path("stream_redirector_race_condition"),
        tee_output=False,
    )

    try:
        result = _process(w, w.setup)
        assert not result.done.error

        for _ in range(5):
            result = _process(w, lambda: w.predict({}))
            assert not result.done.error
    finally:
        w.terminate()


@pytest.mark.parametrize("name,payloads,output_generator", OUTPUT_FIXTURES)
@given(data=st.data())
def test_output(data, name, payloads, output_generator):
    """
    We should get the outputs we expect from predictors that generate output.

    Note that most of the validation work here is actually done in _process.
    """
    w = Worker(predictor_ref=_fixture_path(name), tee_output=False)

    try:
        result = _process(w, w.setup)
        assert not result.done.error

        payload = data.draw(st.fixed_dictionaries(payloads))
        expected_output = output_generator(payload)

        result = _process(w, lambda: w.predict(payload))

        assert result.output == expected_output
    finally:
        w.terminate()


@pytest.mark.parametrize("name,expected_stdout,expected_stderr", SETUP_LOGS_FIXTURES)
def test_setup_logging(name, expected_stdout, expected_stderr):
    """
    We should get the logs we expect from predictors that generate logs during
    setup.
    """
    w = Worker(predictor_ref=_fixture_path(name), tee_output=False)

    try:
        result = _process(w, w.setup)
        assert not result.done.error

        assert result.stdout == expected_stdout
        assert result.stderr == expected_stderr
    finally:
        w.terminate()


@pytest.mark.parametrize(
    "name,payloads,expected_stdout,expected_stderr", PREDICT_LOGS_FIXTURES
)
def test_predict_logging(name, payloads, expected_stdout, expected_stderr):
    """
    We should get the logs we expect from predictors that generate logs during
    predict.
    """
    w = Worker(predictor_ref=_fixture_path(name), tee_output=False)

    try:
        result = _process(w, w.setup)
        assert not result.done.error

        result = _process(w, lambda: w.predict({}))

        assert result.stdout == expected_stdout
        assert result.stderr == expected_stderr
    finally:
        w.terminate()


def test_cancel_is_safe():
    """
    Calls to cancel at any time should not result in unexpected things
    happening or the cancelation of unexpected predictions.
    """

    w = Worker(predictor_ref=_fixture_path("sleep"), tee_output=True)

    try:
        for _ in range(50):
            w.cancel()

        _process(w, w.setup)

        for _ in range(50):
            w.cancel()

        result1 = _process(
            w, lambda: w.predict({"sleep": 0.5}), swallow_exceptions=True
        )

        for _ in range(50):
            w.cancel()

        result2 = _process(
            w, lambda: w.predict({"sleep": 0.1}), swallow_exceptions=True
        )

        assert not result1.exception
        assert not result1.done.canceled
        assert not result2.exception
        assert not result2.done.canceled
        assert result2.output == "done in 0.1 seconds"
    finally:
        w.terminate()


def test_cancel_idempotency():
    """
    Multiple calls to cancel within the same prediction, while not necessary or
    recommended, should still only result in a single cancelled prediction, and
    should not affect subsequent predictions.
    """
    w = Worker(predictor_ref=_fixture_path("sleep"), tee_output=True)

    def cancel_a_bunch(_):
        for _ in range(100):
            w.cancel()

    try:
        _process(w, w.setup)

        fut = w.predict({"sleep": 0.5})
        # We call cancel a WHOLE BUNCH to make sure that we don't propagate any
        # of those cancelations to subsequent predictions, regardless of the
        # internal implementation of exceptions raised inside signal handlers.
        for _ in range(5):
            time.sleep(0.05)
            for _ in range(100):
                w.cancel()
        result = fut.result()
        assert result.canceled

        result2 = _process(w, lambda: w.predict({"sleep": 0.1}))

        assert not result2.done.canceled
        assert result2.output == "done in 0.1 seconds"
    finally:
        w.terminate()


def test_cancel_multiple_predictions():
    """
    Multiple predictions cancelled in a row shouldn't be a problem. This test
    is mainly ensuring that the _allow_cancel latch in Worker is correctly
    reset every time a prediction starts.
    """

    w = Worker(predictor_ref=_fixture_path("sleep"), tee_output=True)

    try:
        _process(w, w.setup)

        dones: list[Done] = []
        for _ in range(5):
            fut = w.predict({"sleep": 0.1})
            time.sleep(0.01)
            w.cancel()
            dones.append(fut.result())
        assert all(d.canceled for d in dones)

        done_future = w.predict({"sleep": 0})
        assert not done_future.result().canceled
    finally:
        w.terminate()


def test_graceful_shutdown():
    """
    On shutdown, the worker should finish running the current prediction, and
    then exit.
    """

    w = Worker(predictor_ref=_fixture_path("sleep"), tee_output=False)
    saw_first_event = threading.Event()

    try:
        _process(w, w.setup)

        # When we see the first event, we'll start the shutdown process.
        w.subscribe(lambda event: saw_first_event.set())

        fut = w.predict({"sleep": 1})

        saw_first_event.wait(timeout=1)
        w.shutdown(timeout=2)

        assert fut.result() == Done()
    finally:
        w.terminate()


class WorkerState(RuleBasedStateMachine):
    """
    This is a Hypothesis-driven rule-based state machine test. It is intended
    to ensure that any sequence of calls to the public API of Worker leaves the
    instance in an expected state.

    In short: any call should either throw InvalidStateException or should do
    what the caller asked.

    See https://hypothesis.readthedocs.io/en/latest/stateful.html for more on
    stateful testing with Hypothesis.
    """

    def __init__(self):
        super().__init__()

        self.events = []

        self.predict_canceled = False
        self.predict_payload = None
        self.predict_result = None
        self.setup_result = None

        self.worker = Worker(_fixture_path("steps"), tee_output=False)
        self.worker.subscribe(self.events.append)

    @rule(sleep=st.floats(min_value=0, max_value=0.1))
    def wait(self, sleep):
        time.sleep(sleep)

    @rule()
    def setup(self):
        try:
            self.setup_result = self.worker.setup()
        except InvalidStateException:
            pass

    @precondition(lambda x: x.setup_result)
    @rule(timeout=st.floats(min_value=0, max_value=0.1))
    def await_setup_complete(self, timeout):
        try:
            res = self.setup_result.result(timeout=timeout)
        except TimeoutError:
            pass
        else:
            assert isinstance(res, Done)
            self._check_events()

    # For now, don't run another prediction until we've read the result. This
    # is solely a limitation of the tests: predictions don't (yet) have
    # identifiers, so we can't distinguish between them.
    @precondition(lambda x: not x.predict_result)
    @rule(name=ST_NAMES, steps=st.integers(min_value=0, max_value=5))
    def predict(self, name, steps):
        try:
            payload = {"name": name, "steps": steps}
            self.predict_result = self.worker.predict(payload)
            self.predict_payload = payload
        except InvalidStateException:
            pass

    @precondition(lambda x: x.predict_result)
    @rule(timeout=st.floats(min_value=0, max_value=0.1))
    def await_predict_complete(self, timeout):
        try:
            res = self.predict_result.result(timeout=timeout)
        except TimeoutError:
            pass
        else:
            assert isinstance(res, Done)
            self._check_events()

    @precondition(lambda x: x.predict_result)
    @rule()
    def cancel(self):
        self.worker.cancel()
        self.predict_canceled = True

    def teardown(self):
        self.worker.shutdown()
        # self.worker.terminate()

    def _check_events(self):
        if self.setup_result and self.setup_result.done():
            self.setup_result = None
            self._check_setup_events()

        if self.predict_result and self.predict_result.done():
            canceled = self.predict_canceled
            payload = self.predict_payload
            self.predict_canceled = False
            self.predict_payload = None
            self.predict_result = None
            self._check_predict_events(payload, canceled)

    def _check_setup_events(self):
        result = self._consume_result()
        assert result.stdout == "did setup\n"
        assert result.stderr == ""
        assert result.done == Done()

    def _check_predict_events(self, payload, canceled=False):
        result = self._consume_result()

        if canceled:
            # Requesting cancelation does not guarantee that the prediction is
            # canceled. It may complete before the cancelation is processed.
            assert result.done == Done() or result.done == Done(canceled=True)

            # If it was canceled, we can't make any other assertions.
            return

        expected_stdout = ["START\n"]
        for i in range(payload["steps"]):
            expected_stdout.append(f"STEP {i+1}\n")
        expected_stdout.append("END\n")

        assert result.stdout == "".join(expected_stdout)
        assert result.stderr == ""
        assert result.output == f"NAME={payload['name']}"
        assert result.done == Done()

    def _consume_result(self):
        print(self.events)
        r = Result()
        while self.events:
            event = self.events.pop(0)
            _handle_event(r, event)
            if isinstance(event, Done):
                break
        return r


TestWorkerState = WorkerState.TestCase
