from django.test import TestCase
from rest_framework.test import APIClient
from django.contrib.auth import get_user_model
from rest_framework import status
from django.urls import reverse
from apps.companies.models import Company
from apps.permissions.models import Role, Permission, UserCompanyRole

User = get_user_model()


class MultipleCompaniesTestCase(TestCase):
    def setUp(self):
        self.client = APIClient()
        
        # Create test companies
        self.company_a = Company.objects.create(
            name='Company A',
            phone='1234567890',
            is_active=True
        )
        self.company_b = Company.objects.create(
            name='Company B',
            phone='0987654321',
            is_active=True
        )
        
        # Create test user
        self.user = User.objects.create_user(
            username='testuser',
            email='test@example.com',
            password='testpass123',
            user_type='user'
        )
        
        # Add user to both companies
        self.user.companies.add(self.company_a, self.company_b)
        self.user.active_company = self.company_a
        self.user.save()
        
        # Create roles
        self.role_a = Role.objects.create(
            name='Manager',
            company=self.company_a,
            description='Manager in Company A'
        )
        self.role_b = Role.objects.create(
            name='Employee',
            company=self.company_b,
            description='Employee in Company B'
        )
        
        # Create permissions
        self.permission = Permission.objects.create(
            name='View Users',
            codename='view_users'
        )
        
        # Assign permissions to roles
        self.role_a.permissions.add(self.permission)
        self.role_b.permissions.add(self.permission)
        
        # Create UserCompanyRole relationships
        UserCompanyRole.objects.create(
            user=self.user,
            company=self.company_a,
            role=self.role_a
        )
        UserCompanyRole.objects.create(
            user=self.user,
            company=self.company_b,
            role=self.role_b
        )

    def test_user_belongs_to_multiple_companies(self):
        """Test that user belongs to multiple companies"""
        self.assertTrue(self.user.belongs_to_company(self.company_a))
        self.assertTrue(self.user.belongs_to_company(self.company_b))
        self.assertEqual(self.user.companies.count(), 2)

    def test_user_has_active_company(self):
        """Test that user has an active company"""
        self.assertEqual(self.user.active_company, self.company_a)
        self.assertEqual(self.user.get_active_company_name(), 'Company A')

    def test_user_can_switch_companies(self):
        """Test that user can switch between companies"""
        # Switch to company B
        success = self.user.switch_company(self.company_b)
        self.assertTrue(success)
        self.assertEqual(self.user.active_company, self.company_b)

    def test_user_has_different_roles_in_different_companies(self):
        """Test that user has different roles in different companies"""
        role_a = self.user.get_role_in_company(self.company_a)
        role_b = self.user.get_role_in_company(self.company_b)
        
        self.assertEqual(role_a.name, 'Manager')
        self.assertEqual(role_b.name, 'Employee')

    def test_user_permissions_are_company_specific(self):
        """Test that user permissions are specific to active company"""
        # In Company A context
        self.user.active_company = self.company_a
        self.user.save()
        self.assertTrue(self.user.has_company_permission('view_users'))
        
        # In Company B context
        self.user.active_company = self.company_b
        self.user.save()
        self.assertTrue(self.user.has_company_permission('view_users'))

    def test_list_user_companies_api(self):
        """Test the API endpoint for listing user companies"""
        self.client.force_authenticate(user=self.user)
        url = reverse('user-companies')
        response = self.client.get(url)
        
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(len(response.data['companies']), 2)
        self.assertEqual(response.data['active_company']['id'], self.company_a.id)

    def test_switch_company_api(self):
        """Test the API endpoint for switching companies"""
        self.client.force_authenticate(user=self.user)
        url = reverse('switch-company')
        data = {'company_id': self.company_b.id}
        response = self.client.post(url, data)
        
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(response.data['company']['id'], self.company_b.id)
        
        # Verify user's active company changed
        self.user.refresh_from_db()
        self.assertEqual(self.user.active_company, self.company_b)