from datetime import datetime, timedelta
from decimal import Decimal

from django.test import TestCase, tag
from django.contrib.auth import get_user_model
from unittest.mock import MagicMock, patch

from django.utils import timezone

from api.models import Organization, TaskCredit, TaskCreditConfig
from constants.plans import PlanNames
from tasks.services import TaskCreditService
from util.constants.task_constants import TASKS_UNLIMITED
from constants.grant_types import GrantTypeChoices
from util.tool_costs import get_default_task_credit_cost, clear_tool_credit_cost_cache


User = get_user_model()


@tag("batch_task_credits")
class TaskCreditServiceCalculateAvailableTasksTests(TestCase):
    @patch("tasks.services.TaskCreditService.get_user_task_credits_used")
    @patch("tasks.services.TaskCreditService.get_tasks_entitled")
    @tag("batch_task_credits")
    def test_calculate_available_tasks_regular(self, mock_entitled, mock_used):
        user = User.objects.create(username="user1")
        mock_entitled.return_value = 10
        mock_used.return_value = 4

        available = TaskCreditService.calculate_available_tasks(user)

        self.assertEqual(available, 6)
        mock_entitled.assert_called_once_with(user)
        mock_used.assert_called_once_with(user, None)

    @patch("tasks.services.TaskCreditService.get_user_task_credits_used")
    @patch("tasks.services.TaskCreditService.get_tasks_entitled")
    def test_calculate_available_tasks_unlimited(self, mock_entitled, mock_used):
        user = User.objects.create(username="user2")
        mock_entitled.return_value = TASKS_UNLIMITED

        available = TaskCreditService.calculate_available_tasks(user)

        self.assertEqual(available, TASKS_UNLIMITED)
        mock_used.assert_not_called()


@tag("batch_task_credits")
class TaskCreditServiceGrantSubscriptionCreditsTests(TestCase):
    @patch("tasks.services.apps.get_model")
    def test_grant_subscription_credits_avoids_duplicate_invoice(self, mock_get_model):
        user = User.objects.create(username="user3")
        TaskCredit = MagicMock()
        mock_get_model.return_value = TaskCredit
        TaskCredit.objects.filter.return_value.first.return_value = MagicMock()

        granted = TaskCreditService.grant_subscription_credits(user, invoice_id="inv-1")

        self.assertEqual(granted, 0)
        TaskCredit.objects.create.assert_not_called()

    @patch("tasks.services.timezone")
    @patch("tasks.services.apps.get_model")
    @patch("tasks.services.get_active_subscription")
    @patch("tasks.services.get_user_plan")
    @tag("batch_task_credits")
    def test_grant_subscription_credits_sets_expiration_to_subscription_end(self, mock_plan, mock_subscription, mock_get_model, mock_timezone):
        user = User.objects.create(username="user4")
        TaskCredit = MagicMock()
        mock_get_model.return_value = TaskCredit
        TaskCredit.objects.filter.return_value.first.return_value = None

        plan = {"id": "startup", "monthly_task_credits": 5}
        mock_plan.return_value = plan

        sub = MagicMock()
        sub.current_period_end = datetime(2024, 1, 31)
        mock_subscription.return_value = sub

        mock_timezone.now.return_value = datetime(2024, 1, 1)

        TaskCreditService.grant_subscription_credits(user, invoice_id="inv-2")

        TaskCredit.objects.create.assert_called_once()
        args, kwargs = TaskCredit.objects.create.call_args
        self.assertEqual(kwargs["expiration_date"], sub.current_period_end)
        self.assertEqual(kwargs["credits"], plan["monthly_task_credits"])
        self.assertFalse(kwargs["free_trial_start"])

    @patch("tasks.services.timezone")
    @patch("tasks.services.apps.get_model")
    @patch("tasks.services.get_active_subscription")
    @patch("tasks.services.get_user_plan")
    def test_grant_subscription_credits_updates_existing_invoice_to_trial_start(
        self,
        mock_plan,
        mock_subscription,
        mock_get_model,
        mock_timezone,
    ):
        user = User.objects.create(username="user4b")
        TaskCredit = MagicMock()
        mock_get_model.return_value = TaskCredit

        existing_credit = MagicMock()
        existing_credit.plan = PlanNames.STARTUP
        existing_credit.credits = 5
        existing_credit.expiration_date = datetime(2024, 1, 31)
        existing_credit.granted_date = datetime(2023, 12, 31)
        existing_credit.free_trial_start = False
        TaskCredit.objects.filter.return_value.first.return_value = existing_credit

        plan = {"id": PlanNames.STARTUP, "monthly_task_credits": 5}
        mock_plan.return_value = plan

        sub = MagicMock()
        sub.current_period_end = datetime(2024, 1, 31)
        mock_subscription.return_value = sub

        mock_timezone.now.return_value = datetime(2024, 1, 1)

        granted = TaskCreditService.grant_subscription_credits(
            user,
            plan=plan,
            invoice_id="inv-trial",
            free_trial_start=True,
        )

        self.assertEqual(granted, 0)
        self.assertTrue(existing_credit.free_trial_start)
        self.assertEqual(existing_credit.granted_date, datetime(2023, 12, 31))
        self.assertEqual(existing_credit.expiration_date, datetime(2024, 1, 31))
        existing_credit.save.assert_called_once()
        save_kwargs = existing_credit.save.call_args.kwargs
        self.assertEqual(save_kwargs["update_fields"], ["free_trial_start"])


@tag("batch_task_credits")
class TaskCreditServiceConsumeCreditTests(TestCase):
    @patch("tasks.services.TaskCreditService.handle_task_threshold")
    @patch("tasks.services.apps.get_model")
    @tag("batch_task_credits")
    def test_consume_credit_without_additional_task(self, mock_get_model, mock_handle):
        user = User.objects.create(username="user5")
        TaskCredit = MagicMock()
        mock_get_model.return_value = TaskCredit
        credit = MagicMock(credits_used=0, credits=1)
        select = TaskCredit.objects.select_for_update.return_value
        filt = select.filter.return_value
        ordered = filt.order_by.return_value
        ordered.first.return_value = credit

        def refresh():
            credit.credits_used = 1
        credit.refresh_from_db.side_effect = refresh

        result = TaskCreditService.consume_credit(user)

        self.assertIs(result, credit)
        credit.save.assert_called_once()
        credit.refresh_from_db.assert_called_once()
        # No immediate Stripe usage reporting; handled by rollup task
        # (assertion intentionally removed)
        mock_handle.assert_called_once_with(user)

    @patch("tasks.services.apps.get_model")
    @patch("tasks.services.get_user_plan")
    @patch("tasks.services.BillingService.get_current_billing_period_for_user")
    @patch("tasks.services.TaskCreditService.handle_task_threshold")
    def test_consume_credit_with_additional_task(self, mock_handle, mock_period, mock_plan, mock_get_model):
        user = User.objects.create(username="user6")
        TaskCredit = MagicMock()
        mock_get_model.return_value = TaskCredit
        credit = MagicMock(credits_used=0)
        TaskCredit.objects.create.return_value = credit
        mock_plan.return_value = {"id": "startup"}
        mock_period.return_value = (datetime(2024, 1, 1), datetime(2024, 2, 1))

        def refresh():
            credit.credits_used = 1
        credit.refresh_from_db.side_effect = refresh

        result = TaskCreditService.consume_credit(user, additional_task=True)

        TaskCredit.objects.create.assert_called_once()
        self.assertIs(result, credit)
        # No immediate Stripe usage reporting; handled by rollup task
        mock_handle.assert_called_once_with(user)


@tag("batch_task_credits")
class TaskCreditServiceCheckAndConsumeCreditForOwnerTests(TestCase):
    def setUp(self):
        clear_tool_credit_cost_cache()
        TaskCreditConfig.objects.update_or_create(
            singleton_id=1,
            defaults={"default_task_cost": Decimal("0.40")},
        )

    def test_org_consumes_additional_task_when_allowed(self):
        owner = User.objects.create(username="org_owner3")
        org = Organization.objects.create(name="Org3", slug="org3", created_by=owner)
        billing = org.billing
        billing.subscription = PlanNames.ORG_TEAM
        billing.purchased_seats = 1
        billing.max_extra_tasks = 5
        billing.save(update_fields=["subscription", "purchased_seats", "max_extra_tasks"])

        result = TaskCreditService.check_and_consume_credit_for_owner(org)

        self.assertTrue(result["success"])
        credit = result["credit"]
        self.assertIsNotNone(credit)
        credit.refresh_from_db()
        self.assertTrue(credit.additional_task)
        self.assertEqual(credit.organization, org)
        default_cost = float(get_default_task_credit_cost())
        self.assertEqual(float(credit.credits), default_cost)
        self.assertEqual(float(credit.credits_used), default_cost)

    def test_org_consumption_fails_without_overage(self):
        owner = User.objects.create(username="org_owner4")
        org = Organization.objects.create(name="Org4", slug="org4", created_by=owner)
        billing = org.billing
        billing.subscription = PlanNames.ORG_TEAM
        billing.purchased_seats = 1
        billing.max_extra_tasks = 0
        billing.save(update_fields=["subscription", "purchased_seats", "max_extra_tasks"])

        result = TaskCreditService.check_and_consume_credit_for_owner(org)

        self.assertFalse(result["success"])
        self.assertIn("no remaining task credits", result["error_message"].lower())
        self.assertIsNone(result["credit"])

    def test_org_additional_tasks_blocked_without_paid_seats(self):
        owner = User.objects.create(username="org_owner5")
        org = Organization.objects.create(name="Org5", slug="org5", created_by=owner)
        billing = org.billing
        billing.subscription = PlanNames.ORG_TEAM
        billing.purchased_seats = 0
        billing.max_extra_tasks = TASKS_UNLIMITED
        billing.save(update_fields=["subscription", "purchased_seats", "max_extra_tasks"])

        result = TaskCreditService.check_and_consume_credit_for_owner(org)

        self.assertFalse(result["success"])
        self.assertIsNone(result["credit"])
        self.assertIn("no remaining task credits", result["error_message"].lower())

    def test_refund_consumed_credit_prefers_linked_credit(self):
        user = User.objects.create(username="refund_user")
        now = timezone.now()
        credit = TaskCredit.objects.create(
            user=user,
            credits=Decimal("10.000"),
            credits_used=Decimal("5.000"),
            granted_date=now - timedelta(days=1),
            expiration_date=now + timedelta(days=30),
            grant_type=GrantTypeChoices.COMPENSATION,
        )

        result = TaskCreditService.refund_consumed_credit_for_owner(
            user,
            amount=Decimal("2.000"),
            preferred_credit=credit,
        )

        self.assertTrue(result["success"])
        self.assertEqual(result["refunded"], Decimal("2.000"))
        credit.refresh_from_db()
        self.assertEqual(credit.credits_used, Decimal("3.000"))

    def test_refund_consumed_credit_clamps_without_negative_usage(self):
        user = User.objects.create(username="refund_clamp_user")
        now = timezone.now()
        preferred_credit = TaskCredit.objects.create(
            user=user,
            credits=Decimal("10.000"),
            credits_used=Decimal("1.000"),
            granted_date=now - timedelta(days=2),
            expiration_date=now + timedelta(days=20),
            grant_type=GrantTypeChoices.COMPENSATION,
        )
        fallback_credit = TaskCredit.objects.create(
            user=user,
            credits=Decimal("10.000"),
            credits_used=Decimal("2.000"),
            granted_date=now - timedelta(days=1),
            expiration_date=now + timedelta(days=30),
            grant_type=GrantTypeChoices.COMPENSATION,
        )

        result = TaskCreditService.refund_consumed_credit_for_owner(
            user,
            amount=Decimal("5.000"),
            preferred_credit=preferred_credit,
        )

        self.assertFalse(result["success"])
        self.assertEqual(result["refunded"], Decimal("3.000"))
        self.assertEqual(result["remaining"], Decimal("2.000"))
        preferred_credit.refresh_from_db()
        fallback_credit.refresh_from_db()
        self.assertEqual(preferred_credit.credits_used, Decimal("0.000"))
        self.assertEqual(fallback_credit.credits_used, Decimal("0.000"))

    def test_refund_consumed_credit_skips_expired_fallback_blocks(self):
        user = User.objects.create(username="refund_expired_fallback_user")
        now = timezone.now()
        preferred_credit = TaskCredit.objects.create(
            user=user,
            credits=Decimal("10.000"),
            credits_used=Decimal("1.000"),
            granted_date=now - timedelta(days=5),
            expiration_date=now + timedelta(days=10),
            grant_type=GrantTypeChoices.COMPENSATION,
        )
        expired_credit = TaskCredit.objects.create(
            user=user,
            credits=Decimal("10.000"),
            credits_used=Decimal("2.000"),
            granted_date=now - timedelta(days=30),
            expiration_date=now - timedelta(days=1),
            grant_type=GrantTypeChoices.COMPENSATION,
        )
        active_fallback_credit = TaskCredit.objects.create(
            user=user,
            credits=Decimal("10.000"),
            credits_used=Decimal("1.000"),
            granted_date=now - timedelta(days=2),
            expiration_date=now + timedelta(days=20),
            grant_type=GrantTypeChoices.COMPENSATION,
        )

        result = TaskCreditService.refund_consumed_credit_for_owner(
            user,
            amount=Decimal("3.000"),
            preferred_credit=preferred_credit,
        )

        self.assertFalse(result["success"])
        self.assertEqual(result["refunded"], Decimal("2.000"))
        self.assertEqual(result["remaining"], Decimal("1.000"))
        preferred_credit.refresh_from_db()
        expired_credit.refresh_from_db()
        active_fallback_credit.refresh_from_db()
        self.assertEqual(preferred_credit.credits_used, Decimal("0.000"))
        self.assertEqual(expired_credit.credits_used, Decimal("2.000"))
        self.assertEqual(active_fallback_credit.credits_used, Decimal("0.000"))

    def test_refund_consumed_credit_fallback_prefers_latest_expiration(self):
        user = User.objects.create(username="refund_latest_fallback_user")
        now = timezone.now()
        early_expiring_credit = TaskCredit.objects.create(
            user=user,
            credits=Decimal("10.000"),
            credits_used=Decimal("1.000"),
            granted_date=now - timedelta(days=3),
            expiration_date=now + timedelta(days=10),
            grant_type=GrantTypeChoices.COMPENSATION,
        )
        late_expiring_credit = TaskCredit.objects.create(
            user=user,
            credits=Decimal("10.000"),
            credits_used=Decimal("1.000"),
            granted_date=now - timedelta(days=2),
            expiration_date=now + timedelta(days=30),
            grant_type=GrantTypeChoices.COMPENSATION,
        )

        result = TaskCreditService.refund_consumed_credit_for_owner(
            user,
            amount=Decimal("1.000"),
        )

        self.assertTrue(result["success"])
        early_expiring_credit.refresh_from_db()
        late_expiring_credit.refresh_from_db()
        self.assertEqual(early_expiring_credit.credits_used, Decimal("1.000"))
        self.assertEqual(late_expiring_credit.credits_used, Decimal("0.000"))


@tag("batch_task_credits")
class TaskCreditServiceGetTasksEntitledTests(TestCase):
    @patch("tasks.services.get_user_extra_task_limit")
    @patch("tasks.services.apps.get_model")
    @patch("tasks.services.get_user_plan")
    @tag("batch_task_credits")
    def test_get_tasks_entitled_sums_granted_and_extra(self, mock_plan, mock_get_model, mock_extra):
        user = User.objects.create(username="user7")
        mock_plan.return_value = {"id": "PRO", "monthly_task_credits": 5}
        mock_extra.return_value = 2
        TaskCredit = MagicMock()
        mock_get_model.return_value = TaskCredit
        TaskCredit.objects.filter.return_value.aggregate.return_value = {"total_granted": 10}

        result = TaskCreditService.get_tasks_entitled(user)

        self.assertEqual(result, 12)

    @patch("tasks.services.get_user_extra_task_limit")
    @patch("tasks.services.get_user_plan")
    def test_get_tasks_entitled_unlimited_extra(self, mock_plan, mock_extra):
        user = User.objects.create(username="user8")
        mock_plan.return_value = {"id": "startup"}
        mock_extra.return_value = TASKS_UNLIMITED

        result = TaskCreditService.get_tasks_entitled(user)

        self.assertEqual(result, TASKS_UNLIMITED)

    def test_get_tasks_entitled_for_organization_uses_seats(self):
        owner = User.objects.create(username="org_owner")
        org = Organization.objects.create(name="Org", slug="org", created_by=owner)
        billing = org.billing
        billing.subscription = PlanNames.ORG_TEAM
        billing.purchased_seats = 3
        billing.max_extra_tasks = 0
        billing.save(update_fields=["subscription", "purchased_seats", "max_extra_tasks"])

        result = TaskCreditService.get_tasks_entitled_for_owner(org)

        # Team plan grants 500 per seat => 1500 total
        self.assertEqual(result, 1500)

    def test_get_tasks_entitled_for_organization_with_unlimited_extra(self):
        owner = User.objects.create(username="org_owner2")
        org = Organization.objects.create(name="Org2", slug="org2", created_by=owner)
        billing = org.billing
        billing.subscription = PlanNames.ORG_TEAM
        billing.purchased_seats = 1
        billing.max_extra_tasks = TASKS_UNLIMITED
        billing.save(update_fields=["subscription", "purchased_seats", "max_extra_tasks"])

        result = TaskCreditService.get_tasks_entitled_for_owner(org)

        self.assertEqual(result, TASKS_UNLIMITED)


@tag("batch_task_credits")
class TaskCreditServiceCalculateUsedPctTests(TestCase):
    @patch("tasks.services.TaskCreditService.get_tasks_entitled")
    @patch("tasks.services.TaskCreditService.get_user_total_tasks_used")
    def test_calculate_used_pct_normal_and_capped(self, mock_used, mock_entitled):
        user = User.objects.create(username="user9")
        mock_used.return_value = 5
        mock_entitled.return_value = 10
        self.assertEqual(TaskCreditService.calculate_used_pct(user), 50.0)

    @patch("tasks.services.TaskCreditService.get_user_total_tasks_used")
    @patch("tasks.services.TaskCreditService.get_tasks_entitled")
    def test_calculate_used_pct_capped_at_100(self, mock_used, mock_entitled):
        user = User.objects.create(username="user9")
        mock_entitled.return_value = 20
        mock_used.return_value = 20
        self.assertEqual(TaskCreditService.calculate_used_pct(user), 100.0)

    @patch("tasks.services.TaskCreditService.get_user_total_tasks_used")
    @patch("tasks.services.TaskCreditService.get_tasks_entitled")
    def test_calculate_used_pct_with_zero_entitled(self, mock_used, mock_entitled):
        user = User.objects.create(username="user9")
        mock_entitled.return_value = 0
        mock_used.return_value = 5
        self.assertEqual(TaskCreditService.calculate_used_pct(user), 0.0)


@tag("batch_task_credits")
class TaskCreditServiceGrantOrgSubscriptionCreditsTests(TestCase):
    def setUp(self):
        owner = User.objects.create(username="org_owner_seed")
        self.org = Organization.objects.create(name="SeedOrg", slug="seedorg", created_by=owner)
        billing = self.org.billing
        billing.subscription = PlanNames.ORG_TEAM
        billing.purchased_seats = 2
        billing.save(update_fields=["subscription", "purchased_seats"])

    def test_grant_subscription_credits_for_organization_creates_credit(self):
        TaskCreditService.grant_subscription_credits_for_organization(
            self.org,
            seats=2,
            invoice_id="inv-org-1",
        )

        credits = TaskCredit.objects.filter(organization=self.org, stripe_invoice_id="inv-org-1")
        self.assertEqual(credits.count(), 1)
        self.assertGreater(float(credits.first().credits), 0)

    def test_grant_subscription_credits_for_organization_idempotent(self):
        TaskCreditService.grant_subscription_credits_for_organization(
            self.org,
            seats=1,
            invoice_id="inv-org-dup",
        )
        duplicate = TaskCreditService.grant_subscription_credits_for_organization(
            self.org,
            seats=1,
            invoice_id="inv-org-dup",
        )

        self.assertEqual(duplicate, 0)
        self.assertEqual(
            TaskCredit.objects.filter(organization=self.org, stripe_invoice_id="inv-org-dup").count(),
            1,
        )

    def test_grant_subscription_credits_for_organization_replaces_current_cycle(self):
        grant_date = timezone.now() - timedelta(days=1)
        expiration = timezone.now() + timedelta(days=29)
        existing = TaskCredit.objects.create(
            organization=self.org,
            credits=Decimal("200.0"),
            credits_used=Decimal("150.0"),
            granted_date=grant_date,
            expiration_date=expiration,
            plan=PlanNames.ORG_TEAM,
            grant_type=GrantTypeChoices.PLAN,
            additional_task=False,
        )

        returned = TaskCreditService.grant_subscription_credits_for_organization(
            self.org,
            seats=self.org.billing.purchased_seats,
            invoice_id="inv-org-renew",
            grant_date=timezone.now(),
            expiration_date=timezone.now() + timedelta(days=30),
            replace_current=True,
        )

        self.assertEqual(returned, 1000)
        # Should update existing block instead of creating a new one
        credits = TaskCredit.objects.filter(organization=self.org)
        self.assertEqual(credits.count(), 1)
        refreshed = TaskCredit.objects.get(pk=existing.pk)
        self.assertEqual(refreshed.stripe_invoice_id, "inv-org-renew")
        self.assertEqual(refreshed.credits, Decimal("1000"))
        self.assertEqual(refreshed.credits_used, Decimal("0"))

@tag("batch_task_credits")
class TaskCreditServiceHandleThresholdTests(TestCase):
    @patch("tasks.services.apps.get_model")
    @patch("tasks.services.Analytics.publish_threshold_event")
    def test_handle_task_threshold_triggers_notifications(self, mock_publish, mock_get_model):
        user = User.objects.create(username="user10")
        Usage = MagicMock()
        mock_get_model.return_value = Usage

        # First call crossing 75%
        Usage.objects.get_or_create.return_value = (MagicMock(), True)
        with patch.object(TaskCreditService, "get_tasks_entitled", return_value=100), \
             patch.object(TaskCreditService, "get_user_total_tasks_used", return_value=80):
            TaskCreditService.handle_task_threshold(user)
        mock_publish.assert_called_once()

        # Second call crossing 90%, 75% already sent
        mock_publish.reset_mock()
        Usage.objects.get_or_create.side_effect = [
            (MagicMock(), False),  # 75% already exists
            (MagicMock(), True),   # 90% newly crossed
        ]
        with patch.object(TaskCreditService, "get_tasks_entitled", return_value=100), \
             patch.object(TaskCreditService, "get_user_total_tasks_used", return_value=95):
            TaskCreditService.handle_task_threshold(user)
        mock_publish.assert_called_once()

    @patch("tasks.services.apps.get_model")
    @patch("tasks.services.Analytics.publish_threshold_event")
    def test_handle_task_threshold_treats_small_remainder_as_full_usage(self, mock_publish, mock_get_model):
        user = User.objects.create(username="user11")
        Usage = MagicMock()
        mock_get_model.return_value = Usage

        Usage.objects.get_or_create.side_effect = [
            (MagicMock(), False),  # 75% already recorded this month
            (MagicMock(), False),  # 90% already recorded this month
            (MagicMock(), True),   # 100% should fire when remainder within tolerance
        ]

        with patch.object(TaskCreditService, "get_tasks_entitled", return_value=Decimal("100")), \
             patch.object(TaskCreditService, "get_user_total_tasks_used", return_value=Decimal("99.96")):
            TaskCreditService.handle_task_threshold(user)

        mock_publish.assert_called_once()
        _, threshold, pct, *_ = mock_publish.call_args[0]
        self.assertEqual(threshold, 100)
        self.assertEqual(pct, 100)
