import uuid
import tempfile
from unittest.mock import patch

from django.urls import reverse
from django.contrib.auth import get_user_model
from django.core.files.base import ContentFile
from django.test import TestCase, override_settings, tag

from waffle.models import Flag

from api.models import (
    Organization,
    OrganizationMembership,
    AgentCollaborator,
    BrowserUseAgent,
    PersistentAgent,
    BrowserUseAgentTask,
    TaskCredit,
    OrganizationInvite,
    ProxyServer,
    DedicatedProxyAllocation,
)
from django.utils import timezone
from constants.plans import PlanNamesChoices


User = get_user_model()


@override_settings(
    DATABASES={
        'default': {
            'ENGINE': 'django.db.backends.sqlite3',
            'NAME': ':memory:',
        }
    },
    SEGMENT_WRITE_KEY="",
    SEGMENT_WEB_WRITE_KEY="",
)

@tag('batch_console_context')
class ConsoleContextTests(TestCase):
    def setUp(self):
        # Enable organizations feature flag for all requests
        Flag.objects.update_or_create(name="organizations", defaults={"everyone": True})

        # Users
        self.owner = User.objects.create_user(username="owner", email="owner@example.com", password="pw")
        self.stranger = User.objects.create_user(username="stranger", email="stranger@example.com", password="pw")

        # Org and membership
        self.org = Organization.objects.create(
            name="Acme, Inc.", slug="acme", plan="free", created_by=self.owner
        )
        OrganizationMembership.objects.create(
            org=self.org,
            user=self.owner,
            role=OrganizationMembership.OrgRole.OWNER,
            status=OrganizationMembership.OrgStatus.ACTIVE,
        )
        billing = self.org.billing
        billing.purchased_seats = 3
        billing.subscription = PlanNamesChoices.ORG_TEAM.value
        billing.save(update_fields=["purchased_seats", "subscription"])
        owner_billing = self.owner.billing
        owner_billing.subscription = PlanNamesChoices.STARTUP.value
        owner_billing.save(update_fields=["subscription"])

        # Agents
        self.personal_browser = BrowserUseAgent.objects.create(user=self.owner, name="Personal Agent")
        self.personal_agent = PersistentAgent.objects.create(
            user=self.owner,
            organization=None,
            name="Personal PA",
            charter="",
            browser_use_agent=self.personal_browser,
        )

        self.org_browser = BrowserUseAgent.objects.create(user=self.owner, name="Org Agent")
        self.org_agent = PersistentAgent.objects.create(
            user=self.owner,
            organization=self.org,
            name="Org PA",
            charter="",
            browser_use_agent=self.org_browser,
        )

        # Ensure the organization has credits so org tasks can be created
        TaskCredit.objects.create(
            organization=self.org,
            credits=10,
            credits_used=0,
            granted_date=timezone.now(),
            expiration_date=timezone.now() + timezone.timedelta(days=30),
        )

        # Tasks: one personal, one org-owned, one agent-less
        self.personal_task = BrowserUseAgentTask.objects.create(user=self.owner, agent=self.personal_browser)
        self.org_task = BrowserUseAgentTask.objects.create(user=self.owner, agent=self.org_browser)
        self.agentless_task = BrowserUseAgentTask.objects.create(user=self.owner, agent=None)

        # Login owner by default
        assert self.client.login(username="owner", password="pw")

    def _set_personal_context(self):
        session = self.client.session
        session["context_type"] = "personal"
        session["context_id"] = str(self.owner.id)
        session["context_name"] = self.owner.get_full_name() or self.owner.username
        session.save()

    def _set_org_context(self):
        session = self.client.session
        session["context_type"] = "organization"
        session["context_id"] = str(self.org.id)
        session["context_name"] = self.org.name
        session.save()

    def test_tasks_view_personal_excludes_org_owned(self):
        self._set_personal_context()
        url = reverse("tasks")
        resp = self.client.get(url)
        self.assertEqual(resp.status_code, 200)
        tasks = list(resp.context["tasks"])  # paginated page object
        ids = {t.id for t in tasks}
        self.assertIn(self.personal_task.id, ids)
        self.assertIn(self.agentless_task.id, ids)
        self.assertNotIn(self.org_task.id, ids)

    def test_switch_context_invalid_org_override_format_returns_403(self):
        resp = self.client.get(
            reverse("switch_context"),
            HTTP_X_GOBII_CONTEXT_TYPE="organization",
            HTTP_X_GOBII_CONTEXT_ID="not-a-uuid",
        )
        self.assertEqual(resp.status_code, 403)
        self.assertEqual(resp.json().get("error"), "Invalid context override.")

    def test_switch_context_for_agent_returns_org_context_without_persisting_session(self):
        self._set_personal_context()

        resp = self.client.get(
            reverse("switch_context"),
            {"for_agent": str(self.org_agent.id)},
        )
        self.assertEqual(resp.status_code, 200)
        payload = resp.json()
        self.assertEqual(payload.get("context", {}).get("type"), "organization")
        self.assertEqual(payload.get("context", {}).get("id"), str(self.org.id))

        session = self.client.session
        self.assertEqual(session.get("context_type"), "personal")
        self.assertEqual(session.get("context_id"), str(self.owner.id))

    def test_switch_context_for_agent_overrides_stale_context_headers(self):
        self._set_personal_context()

        resp = self.client.get(
            reverse("switch_context"),
            {"for_agent": str(self.org_agent.id)},
            HTTP_X_GOBII_CONTEXT_TYPE="personal",
            HTTP_X_GOBII_CONTEXT_ID=str(self.owner.id),
        )
        self.assertEqual(resp.status_code, 200)
        payload = resp.json()
        self.assertEqual(payload.get("context", {}).get("type"), "organization")
        self.assertEqual(payload.get("context", {}).get("id"), str(self.org.id))

    def test_switch_context_for_agent_forbidden_without_access(self):
        self.client.logout()
        assert self.client.login(username="stranger", password="pw")
        session = self.client.session
        session["context_type"] = "personal"
        session["context_id"] = str(self.stranger.id)
        session["context_name"] = self.stranger.username
        session.save()

        resp = self.client.get(
            reverse("switch_context"),
            {"for_agent": str(self.org_agent.id)},
        )
        self.assertEqual(resp.status_code, 403)
        self.assertEqual(resp.json().get("error"), "Not permitted")

    def test_roster_for_deleted_agent_returns_remaining_agents(self):
        extra_browser = BrowserUseAgent.objects.create(user=self.owner, name="Org Agent Two")
        extra_org_agent = PersistentAgent.objects.create(
            user=self.owner,
            organization=self.org,
            name="Org PA Two",
            charter="",
            browser_use_agent=extra_browser,
        )
        self.org_agent.soft_delete()
        self._set_personal_context()

        resp = self.client.get(
            reverse("console_agent_roster"),
            {"for_agent": str(self.org_agent.id)},
        )
        self.assertEqual(resp.status_code, 200)
        payload = resp.json()
        self.assertEqual(payload.get("requested_agent_status"), "deleted")
        self.assertEqual(payload.get("context", {}).get("type"), "organization")
        self.assertEqual(payload.get("context", {}).get("id"), str(self.org.id))
        roster_ids = {entry["id"] for entry in payload.get("agents", [])}
        self.assertIn(str(extra_org_agent.id), roster_ids)
        self.assertNotIn(str(self.org_agent.id), roster_ids)

    def test_switch_context_for_deleted_agent_returns_org_context(self):
        self.org_agent.soft_delete()
        self._set_personal_context()

        resp = self.client.get(
            reverse("switch_context"),
            {"for_agent": str(self.org_agent.id)},
        )
        self.assertEqual(resp.status_code, 200)
        payload = resp.json()
        self.assertEqual(payload.get("context", {}).get("type"), "organization")
        self.assertEqual(payload.get("context", {}).get("id"), str(self.org.id))

        session = self.client.session
        self.assertEqual(session.get("context_type"), "personal")
        self.assertEqual(session.get("context_id"), str(self.owner.id))

    def test_switch_context_for_org_agent_allows_collaborator_without_membership(self):
        AgentCollaborator.objects.create(agent=self.org_agent, user=self.stranger)

        self.client.logout()
        assert self.client.login(username="stranger", password="pw")
        session = self.client.session
        session["context_type"] = "personal"
        session["context_id"] = str(self.stranger.id)
        session["context_name"] = self.stranger.username
        session.save()

        resp = self.client.get(
            reverse("switch_context"),
            {"for_agent": str(self.org_agent.id)},
        )
        self.assertEqual(resp.status_code, 200)
        payload = resp.json().get("context", {})
        self.assertEqual(payload.get("type"), "personal")
        self.assertEqual(payload.get("id"), str(self.stranger.id))

    def test_switch_context_for_personal_agent_allows_collaborator(self):
        AgentCollaborator.objects.create(agent=self.personal_agent, user=self.stranger)

        self.client.logout()
        assert self.client.login(username="stranger", password="pw")
        session = self.client.session
        session["context_type"] = "personal"
        session["context_id"] = str(self.stranger.id)
        session["context_name"] = self.stranger.username
        session.save()

        resp = self.client.get(
            reverse("switch_context"),
            {"for_agent": str(self.personal_agent.id)},
        )
        self.assertEqual(resp.status_code, 200)
        payload = resp.json().get("context", {})
        self.assertEqual(payload.get("type"), "personal")
        self.assertEqual(payload.get("id"), str(self.stranger.id))

    def test_tasks_view_org_requires_membership_and_shows_org_tasks(self):
        # As owner (member) — should see org tasks
        self._set_org_context()
        url = reverse("tasks")
        resp = self.client.get(url)
        self.assertEqual(resp.status_code, 200)
        tasks = list(resp.context["tasks"])  # paginated page object
        ids = {t.id for t in tasks}
        self.assertIn(self.org_task.id, ids)
        # Switch to stranger (no membership) — should be forbidden
        self.client.logout()
        assert self.client.login(username="stranger", password="pw")
        self._set_org_context()
        resp2 = self.client.get(url)
        self.assertEqual(resp2.status_code, 403)

    def test_agent_detail_scoping(self):
        self._set_personal_context()
        url = reverse("agent_detail", kwargs={"pk": self.org_agent.id})

        # Direct navigation to another context's agent should still render.
        resp = self.client.get(url)
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(resp.context["current_context"]["type"], "organization")
        self.assertEqual(resp.context["current_context"]["id"], str(self.org.id))
        session = self.client.session
        self.assertEqual(session.get("context_type"), "personal")
        self.assertEqual(session.get("context_id"), str(self.owner.id))

        # Explicit context override query should also render and not mutate session.
        resp_with_override = self.client.get(
            url,
            {
                "context_type": "organization",
                "context_id": str(self.org.id),
            },
        )
        self.assertEqual(resp_with_override.status_code, 200)
        self.assertEqual(resp_with_override.context["current_context"]["type"], "organization")
        self.assertEqual(resp_with_override.context["current_context"]["id"], str(self.org.id))
        session = self.client.session
        self.assertEqual(session.get("context_type"), "personal")
        self.assertEqual(session.get("context_id"), str(self.owner.id))

        # Org context with membership: should 200
        self._set_org_context()
        resp2 = self.client.get(url)
        self.assertEqual(resp2.status_code, 200)

    def test_agent_targeted_views_use_agent_owner_context_without_persisting_session(self):
        self._set_personal_context()

        chat_response = self.client.get(
            reverse("agent_chat_shell", kwargs={"pk": self.org_agent.id}),
        )
        self.assertEqual(chat_response.status_code, 200)
        self.assertEqual(chat_response.context["current_context"]["type"], "organization")
        self.assertEqual(chat_response.context["current_context"]["id"], str(self.org.id))

        files_response = self.client.get(
            reverse("agent_files", kwargs={"pk": self.org_agent.id}),
        )
        self.assertEqual(files_response.status_code, 200)
        self.assertEqual(files_response.context["current_context"]["type"], "organization")
        self.assertEqual(files_response.context["current_context"]["id"], str(self.org.id))

        secrets_response = self.client.get(
            reverse("agent_secrets", kwargs={"pk": self.org_agent.id}),
        )
        self.assertEqual(secrets_response.status_code, 200)
        self.assertEqual(secrets_response.context["current_context"]["type"], "organization")
        self.assertEqual(secrets_response.context["current_context"]["id"], str(self.org.id))

        email_response = self.client.get(
            reverse("agent_email_settings", kwargs={"pk": self.org_agent.id}),
        )
        self.assertEqual(email_response.status_code, 200)
        self.assertEqual(email_response.context["current_context"]["type"], "organization")
        self.assertEqual(email_response.context["current_context"]["id"], str(self.org.id))

        session = self.client.session
        self.assertEqual(session.get("context_type"), "personal")
        self.assertEqual(session.get("context_id"), str(self.owner.id))

    def test_agent_targeted_apis_allow_authorized_user_outside_current_context(self):
        self._set_personal_context()

        with tempfile.TemporaryDirectory() as tmp_media:
            with override_settings(MEDIA_ROOT=tmp_media, MEDIA_URL="/media/"):
                self.org_agent.avatar.save("avatar.png", ContentFile(b"avatar-bytes"), save=True)

                avatar_response = self.client.get(
                    reverse("agent_avatar", kwargs={"pk": self.org_agent.id}),
                )
                self.assertEqual(avatar_response.status_code, 200)

                files_response = self.client.get(
                    reverse("console_agent_fs_list", kwargs={"agent_id": self.org_agent.id}),
                )
                self.assertEqual(files_response.status_code, 200)

                timeline_response = self.client.get(
                    reverse("console_agent_timeline", kwargs={"agent_id": self.org_agent.id}),
                )
                self.assertEqual(timeline_response.status_code, 200)

                email_settings_response = self.client.get(
                    reverse("console_agent_email_settings", kwargs={"agent_id": self.org_agent.id}),
                )
                self.assertEqual(email_settings_response.status_code, 200)

        session = self.client.session
        self.assertEqual(session.get("context_type"), "personal")
        self.assertEqual(session.get("context_id"), str(self.owner.id))

    def test_org_detail_sets_console_context(self):
        # Visiting org detail should set session context to organization
        url = reverse("organization_detail", kwargs={"org_id": self.org.id})
        resp = self.client.get(url)
        self.assertEqual(resp.status_code, 200)
        session = self.client.session
        self.assertEqual(session.get("context_type"), "organization")
        self.assertEqual(session.get("context_id"), str(self.org.id))
        self.assertEqual(session.get("context_name"), self.org.name)

    def test_leaving_org_resets_context_to_personal(self):
        # Add a second owner so the original owner can leave
        another = User.objects.create_user(username="other", email="other@example.com", password="pw")
        OrganizationMembership.objects.create(
            org=self.org,
            user=another,
            role=OrganizationMembership.OrgRole.OWNER,
            status=OrganizationMembership.OrgStatus.ACTIVE,
        )
        self._set_org_context()
        leave_url = reverse("org_leave_org", kwargs={"org_id": self.org.id})
        resp = self.client.post(leave_url, follow=True)
        self.assertEqual(resp.status_code, 200)
        # Verify membership updated
        mem = OrganizationMembership.objects.get(org=self.org, user=self.owner)
        self.assertEqual(mem.status, OrganizationMembership.OrgStatus.REMOVED)
        # Session reset to personal
        session = self.client.session
        self.assertEqual(session.get("context_type"), "personal")
        self.assertEqual(session.get("context_id"), str(self.owner.id))

    def test_header_menu_reflects_context(self):
        # Organization context should show Organization link and hide Profile
        self._set_org_context()
        resp = self.client.get(reverse("console-home"))
        self.assertEqual(resp.status_code, 200)
        html = resp.content.decode()
        self.assertIn(str(self.org.id), html)
        self.assertIn("Organization", html)
        # Switch to personal context
        self._set_personal_context()
        resp2 = self.client.get(reverse("console-home"))
        self.assertEqual(resp2.status_code, 200)
        html2 = resp2.content.decode()
        self.assertIn("Profile", html2)

    def test_sidebar_nav_reflects_context(self):
        # Org context: sidebar should show Organization link
        self._set_org_context()
        resp = self.client.get(reverse("agents"))
        self.assertEqual(resp.status_code, 200)
        content = resp.content.decode()
        self.assertIn("Organization", content)
        # Personal context: sidebar should show Profile link
        self._set_personal_context()
        resp2 = self.client.get(reverse("agents"))
        self.assertEqual(resp2.status_code, 200)
        content2 = resp2.content.decode()
        self.assertIn("Profile", content2)

    def test_agent_detail_includes_dedicated_ip_counts(self):
        self._set_personal_context()
        proxy = ProxyServer.objects.create(
            name="Dedicated Proxy",
            proxy_type=ProxyServer.ProxyType.HTTP,
            host="dedicated.example.com",
            port=8080,
            username="dedicated",
            password="secret",
            static_ip="203.0.113.5",
            is_active=True,
            is_dedicated=True,
        )
        DedicatedProxyAllocation.objects.assign_to_owner(proxy, self.owner)

        url = reverse("agent_detail", kwargs={"pk": self.personal_agent.id})
        resp = self.client.get(url)
        self.assertEqual(resp.status_code, 200)
        props = resp.context.get("agent_detail_props") or {}
        dedicated_ips = props.get("dedicatedIps") or {}
        self.assertEqual(dedicated_ips.get("total"), 1)
        self.assertEqual(dedicated_ips.get("available"), 1)
        self.assertEqual(dedicated_ips.get("ownerType"), "user")
        self.assertEqual(dedicated_ips.get("selectedId"), None)
        options = dedicated_ips.get("options") or []
        self.assertEqual(len(options), 1)
        self.assertEqual(options[0]["label"], "203.0.113.5")
        self.assertEqual(options[0]["assignedNames"], [])

    def test_billing_query_switches_to_org_context(self):
        self._set_personal_context()
        billing_url = f"{reverse('billing')}?org_id={self.org.id}"
        resp = self.client.get(billing_url)
        self.assertEqual(resp.status_code, 200)
        session = self.client.session
        self.assertEqual(session.get('context_type'), 'organization')
        self.assertEqual(session.get('context_id'), str(self.org.id))
        self.assertEqual(session.get('context_name'), self.org.name)

    def test_agent_contact_view_shows_org_context_banner(self):
        self._set_org_context()
        session = self.client.session
        session["agent_charter"] = "help the organization"
        session.save()

        resp = self.client.get(reverse("agent_create_contact"))
        self.assertEqual(resp.status_code, 200)
        html = resp.content.decode()
        self.assertIn('id="agent-owner-selector-contact"', html)
        self.assertIn(self.org.name, html)
        self.assertNotIn('disabled aria-disabled="true"', html)

    def test_agent_contact_view_blocks_member_role(self):
        membership = OrganizationMembership.objects.get(org=self.org, user=self.owner)
        membership.role = OrganizationMembership.OrgRole.MEMBER
        membership.save(update_fields=["role"])

        self._set_org_context()
        session = self.client.session
        session["agent_charter"] = "help the organization"
        session.save()

        resp = self.client.get(reverse("agent_create_contact"))
        self.assertEqual(resp.status_code, 200)
        html = resp.content.decode()
        self.assertIn('id="agent-owner-selector-contact"', html)
        self.assertIn("You need to be an organization owner or admin", html)
        self.assertIn('disabled aria-disabled="true"', html)

    @patch("console.views.AgentService.has_agents_available")
    def test_agent_contact_view_allows_when_personal_capacity_available(self, mock_has_capacity):
        """Users can proceed when personal capacity exists even if org is full."""
        self._set_org_context()
        session = self.client.session
        session["agent_charter"] = "coordinate research across teams"
        session.save()

        def _side_effect(owner):
            if isinstance(owner, Organization):
                return False
            return True

        mock_has_capacity.side_effect = _side_effect

        resp = self.client.get(reverse("agent_create_contact"))
        self.assertEqual(resp.status_code, 200)
        self.assertGreaterEqual(mock_has_capacity.call_count, 2)

    def test_agent_contact_post_denied_for_member_role(self):
        membership = OrganizationMembership.objects.get(org=self.org, user=self.owner)
        membership.role = OrganizationMembership.OrgRole.MEMBER
        membership.save(update_fields=["role"])

        self._set_org_context()
        session = self.client.session
        session["agent_charter"] = "orchestrate research"
        session.save()

        payload = {
            "preferred_contact_method": "email",
            "contact_endpoint_email": "owner@example.com",
            "email_enabled": "on",
        }

        resp = self.client.post(reverse("agent_create_contact"), data=payload)
        self.assertEqual(resp.status_code, 200)
        html = resp.content.decode()
        self.assertIn("You need to be an organization owner or admin", html)
        # Ensure no additional org-owned agents were created
        org_agent_count = PersistentAgent.objects.filter(organization=self.org).count()
        self.assertEqual(org_agent_count, 1)

    def test_org_invite_accept_sets_context_and_membership(self):
        # Create invite for a new user
        invitee = User.objects.create_user(username="invitee", email="invitee@example.com", password="pw")
        inv = OrganizationInvite.objects.create(
            org=self.org,
            email=invitee.email,
            role=OrganizationMembership.OrgRole.MEMBER,
            token=uuid.uuid4().hex,
            expires_at=timezone.now() + timezone.timedelta(days=7),
            invited_by=self.owner,
        )
        # Login as invitee and accept
        self.client.logout()
        assert self.client.login(username="invitee", password="pw")
        url = reverse("org_invite_accept", kwargs={"token": inv.token})
        resp = self.client.get(url, follow=True)
        self.assertEqual(resp.status_code, 200)
        # Context set to org
        session = self.client.session
        self.assertEqual(session.get("context_type"), "organization")
        self.assertEqual(session.get("context_id"), str(self.org.id))
        # Membership created/active
        mem = OrganizationMembership.objects.get(org=self.org, user=invitee)
        self.assertEqual(mem.status, OrganizationMembership.OrgStatus.ACTIVE)

    def test_org_invite_reject_sets_context(self):
        invitee = User.objects.create_user(username="invitee2", email="invitee2@example.com", password="pw")
        inv = OrganizationInvite.objects.create(
            org=self.org,
            email=invitee.email,
            role=OrganizationMembership.OrgRole.MEMBER,
            token=uuid.uuid4().hex,
            expires_at=timezone.now() + timezone.timedelta(days=7),
            invited_by=self.owner,
        )
        self.client.logout()
        assert self.client.login(username="invitee2", password="pw")
        url = reverse("org_invite_reject", kwargs={"token": inv.token})
        resp = self.client.get(url, follow=True)
        self.assertEqual(resp.status_code, 200)
        # Context set to org
        session = self.client.session
        self.assertEqual(session.get("context_type"), "organization")
        self.assertEqual(session.get("context_id"), str(self.org.id))
        # No active membership created by rejection
        self.assertFalse(OrganizationMembership.objects.filter(org=self.org, user=invitee, status=OrganizationMembership.OrgStatus.ACTIVE).exists())
