import logging
from collections.abc import Callable
from contextlib import contextmanager
from typing import Any

import pytest
from botocore.stub import Stubber
from generalresearch.currency import USDCent
from mypy_boto3_mturk import MTurkClient
from mypy_boto3_mturk.type_defs import (
    GetAssignmentResponseTypeDef,
)

from jb.flow.assignment_tasks import process_assignment_submitted
from jb.managers.amt import (
    APPROVAL_MESSAGE,
    NO_WORK_APPROVAL_MESSAGE,
    REJECT_MESSAGE_BADDIE,
    REJECT_MESSAGE_NO_WORK,
    REJECT_MESSAGE_UNKNOWN_ASSIGNMENT,
    AMTManager,
)
from jb.managers.assignment import AssignmentManager
from jb.managers.bonus import BonusManager
from jb.managers.hit import HitManager
from jb.models.assignment import Assignment, AssignmentStub
from jb.models.definitions import AssignmentStatus
from jb.models.event import MTurkEvent
from jb.models.hit import Hit


@contextmanager
def amt_stub_context(amt_client: MTurkClient, responses: list[dict[str, Any]]):

    # ty chatgpt for this
    with Stubber(amt_client) as stub:
        for r in responses:
            stub.add_response(
                r["operation"],
                r.get("response", {}),
                r.get("expected_params", {}),
            )
        yield stub


class TestHITTasks:

    def test_fake_get_assignment(
        self,
        amt_client: MTurkClient,
        amtm: AMTManager,
        amt_assignment_id: str,
        amt_worker_id: str,
        assignment_response: GetAssignmentResponseTypeDef,
    ):
        # Testing just that this boto stubber works (we fake a response
        # using the real boto client)
        fake_response = assignment_response.copy()

        with Stubber(amt_client) as stub:
            expected_params = {"AssignmentId": amt_assignment_id}
            stub.add_response("get_assignment", fake_response, expected_params)

            assignment = amtm.get_assignment_if_exists(
                amt_assignment_id=amt_assignment_id
            )
            assert assignment is not None
            assert assignment.amt_assignment_id == amt_assignment_id
            assert assignment.amt_worker_id == amt_worker_id

            # Optionally, ensure all queued responses were used:
            stub.assert_no_pending_responses()


class TestProcessAssignmentSubmitted:

    def test_no_assignment_in_db(
        self,
        amtm: AMTManager,
        hm: HitManager,
        am: AssignmentManager,
        bm: BonusManager,
        amt_client: MTurkClient,
        hit_record: Hit,
        mturk_event: MTurkEvent,
        amt_assignment_id: str,
        caplog: pytest.LogCaptureFixture,
        rejected_assignment_stubs: Callable[..., list[dict[str, Any]]],
    ):

        # These records are auto cleaned up, so we need to explicitly create
        # a HIT record in the DB so the process_assignment_submitted task
        # doesn't error when we try to process the Request
        _ = hit_record

        # An assignment is submitted. The hit exists in the DB. The amt
        # assignment id is valid, but the assignment stub is not in our
        # db. Reject it and write the assignment to the db.

        amt_stubs = rejected_assignment_stubs(
            reject_reason=REJECT_MESSAGE_UNKNOWN_ASSIGNMENT
        )

        with amt_stub_context(amt_client, amt_stubs) as stub, caplog.at_level(
            logging.WARNING
        ):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        assert f"No assignment found in DB: {amt_assignment_id}" in caplog.text
        assert "Rejected assignment doesn't exist in DB. Creating ... " in caplog.text
        assert "Rejected assignment: " in caplog.text
        stub.assert_no_pending_responses()

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Rejected
        assert ass.requester_feedback == REJECT_MESSAGE_UNKNOWN_ASSIGNMENT

    def test_assignment_in_db_user_doesnt_exist(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        amt_client: MTurkClient,
        mturk_event: MTurkEvent,
        amt_assignment_id: str,
        assignment_stub_record: AssignmentStub,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        rejected_assignment_stubs: Callable[..., list[dict[str, Any]]],
    ):
        # An assignment is submitted. The hit and AssignmentStub exist in the
        # DB. We think we're going to approve the Assignment, but the
        # user-profile / check blocked call on THL shows the user doesn't
        # exist (same thing would happen if the user does exist and is
        # blocked). So we reject.

        # We need this to make the assignment stub in the db
        _ = assignment_stub_record

        amt_stubs = rejected_assignment_stubs(reject_reason=REJECT_MESSAGE_BADDIE)

        mock_thl_responses(user_blocked=True)

        with amt_stub_context(amt_client, amt_stubs) as stub, caplog.at_level(
            logging.WARNING
        ):
            process_assignment_submitted(
                amtm=amtm, am=am, hm=hm, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        assert f"No assignment found in DB: {amt_assignment_id}" not in caplog.text
        assert "blocked or not exists" in caplog.text
        assert "Rejected assignment: " in caplog.text

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Rejected
        assert ass.requester_feedback == REJECT_MESSAGE_BADDIE

    def test_no_work_w_warning(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        mturk_event: MTurkEvent,
        amt_client: MTurkClient,
        amt_assignment_id: str,
        assignment_stub_record: AssignmentStub,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        approved_assignment_stubs: Callable[..., list[dict[str, Any]]],
        assignment_response_approved_no_tsid: GetAssignmentResponseTypeDef,
        assignment_response_no_tsid: GetAssignmentResponseTypeDef,
    ):
        # An Assignment is submitted. The hit and AssignmentStub exist in
        # the DB. The assignment has no tsid.
        # We APPROVE this assignment b/c we are very nice and give users a
        #   couple chances, with an explanation, before rejecting.

        # We need this to make the assignment stub in the db
        _ = assignment_stub_record

        # Simulate that the AMT.get_assignment call returns the assignment,
        #   but the answers XML has no tsid.
        amt_stubs = approved_assignment_stubs(
            feedback=NO_WORK_APPROVAL_MESSAGE,
            override_response=assignment_response_no_tsid,
            override_approve_response=assignment_response_approved_no_tsid,
        )

        mock_thl_responses(user_blocked=False)

        with amt_stub_context(amt_client, amt_stubs) as stub, caplog.at_level(
            logging.WARNING
        ):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        assert f"No assignment found in DB: {amt_assignment_id}" not in caplog.text
        assert "Assignment submitted with no tsid" in caplog.text
        assert "Approved assignment: " in caplog.text

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Approved
        assert ass.requester_feedback == NO_WORK_APPROVAL_MESSAGE
        assert am.missing_tsid_count(amt_worker_id=ass.amt_worker_id) == 1

    def test_no_work_no_warning(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        mturk_event: MTurkEvent,
        amt_client: MTurkClient,
        amt_assignment_id: str,
        assignment_stub_record: AssignmentStub,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        rejected_assignment_stubs: Callable[..., list[dict[str, Any]]],
        assignment_response_factory_rejected_no_tsid: Callable[
            ..., GetAssignmentResponseTypeDef
        ],
        assignment_response_no_tsid: GetAssignmentResponseTypeDef,
        assignment_record_factory: Callable[..., Assignment],
        hit_record: Hit,  # used in assignment_factory
        amt_worker_id: str,
    ):
        # An assignment is submitted. The hit and assignment stub exist in the DB.
        # The assignment has no tsid.

        # Going to create and submit 3 assignments w no work
        #  (all on the same hit, which we don't do in JB for real,
        #   but doesn't matter here)
        _a1 = assignment_record_factory(
            hit_id=hit_record.id, amt_worker_id=amt_worker_id
        )
        _a2 = assignment_record_factory(
            hit_id=hit_record.id, amt_worker_id=amt_worker_id
        )
        _a3 = assignment_record_factory(
            hit_id=hit_record.id, amt_worker_id=amt_worker_id
        )

        assert am.missing_tsid_count(amt_worker_id=amt_worker_id) == 3
        # So now, we'll reject, b/c they've already gotten 3 warnings

        _ = assignment_stub_record  # we need this to make the assignment stub in the db

        # Simulate that the AMT.get_assignment call returns the assignment, but the answers xml
        #   has no tsid.
        amt_stubs = rejected_assignment_stubs(
            reject_reason=REJECT_MESSAGE_NO_WORK,
            override_response=assignment_response_no_tsid,
            override_reject_response=assignment_response_factory_rejected_no_tsid(
                REJECT_MESSAGE_NO_WORK
            ),
        )

        mock_thl_responses(user_blocked=False)

        with amt_stub_context(amt_client, amt_stubs) as stub, caplog.at_level(
            logging.WARNING
        ):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        assert f"No assignment found in DB: {amt_assignment_id}" not in caplog.text
        assert "Assignment submitted with no tsid" in caplog.text
        assert "Rejected assignment: " in caplog.text

        # It will exist in the db since we can validate the model.
        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Rejected
        assert ass.requester_feedback == REJECT_MESSAGE_NO_WORK

    def test_assignment_submitted_no_bonus(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        amt_client: MTurkClient,
        mturk_event: MTurkEvent,
        amt_assignment_id: str,
        assignment_stub_record: Assignment,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        approved_assignment_stubs: Callable[..., list[dict[str, Any]]],
    ):

        _ = assignment_stub_record  # we need this to make the assignment stub in the db

        # The "send bonus" stuff will still run, even if the user didn't get
        # a complete, because all we do is check the user's wallet balance (if
        # an assignment is approved) and they may have money in their wallet
        # from a prev event or bribe
        #
        # So mock the wallet balance as 1cent, so no bonus will be triggered

        mock_thl_responses(status_complete=False, wallet_redeemable_amount=1)
        with amt_stub_context(
            amt_client, approved_assignment_stubs()
        ) as stub, caplog.at_level(logging.WARNING):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Approved
        assert ass.requester_feedback == APPROVAL_MESSAGE

    def test_assignment_submitted_w_bonus(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        amt_client: MTurkClient,
        mturk_event: MTurkEvent,
        amt_assignment_id: str,
        assignment_stub_record: Assignment,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        approved_assignment_stubs_w_bonus: list[dict[str, Any]],
    ):
        _ = assignment_stub_record  # we need this to make the assignment stub in the db
        mock_thl_responses(status_complete=True, wallet_redeemable_amount=10)

        with amt_stub_context(
            amt_client, approved_assignment_stubs_w_bonus
        ) as stub, caplog.at_level(logging.WARNING):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Approved
        assert ass.requester_feedback == APPROVAL_MESSAGE

        bonus = bm.filter(amt_assignment_id=amt_assignment_id)[0]
        assert bonus.amount == USDCent(7)