role_based_system/user_management/models.py

from itertools import count
from django.contrib.auth.models import AbstractUser
from django.db import models
from django.utils import timezone
from django.core.exceptions import ValidationError
import uuid
import logging
from django.contrib.auth import get_user_model

logger = logging.getLogger(__name__)

class User(AbstractUser):
    """自定义用户模型"""
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    
    ROLE_CHOICES = (
        ('admin', '管理员'),
        ('leader', '组长'),
        ('member', '组员'),
    )

    name = models.CharField(max_length=150, verbose_name='真实姓名', default='未设置')
    role = models.CharField(max_length=20, choices=ROLE_CHOICES, verbose_name='角色', default='member')
    department = models.CharField(max_length=100, verbose_name='部门', default='未分配')
    group = models.CharField(max_length=100, null=True, blank=True, verbose_name='小组')
    
    class Meta:
        db_table = 'users'
        verbose_name = '用户'
        verbose_name_plural = '用户'

    def __str__(self):
        return f"{self.username}({self.name})"

    def can_manage_department(self):
        """检查是否可以管理部门"""
        return self.role in ['admin', 'leader']

    def can_manage_knowledge_base(self, knowledge_base):
        """检查是否可以管理知识库"""
        if self.role == 'admin':
            return knowledge_base.type != 'private'  # 管理员不能管理私人知识库
        if self.role == 'leader' and self.department == knowledge_base.department:
            return knowledge_base.type == 'member'  # 组长只能管理本部门的成员知识库
        return knowledge_base.user_id == str(self.id)  # 用户可以管理自己创建的知识库

    def has_access_permission(self, knowledge_base):
        """检查用户是否有权限访问知识库"""
        
        # 1. 如果是私人知识库
        if knowledge_base.type == 'private':
            # 创建者直接允许
            if str(knowledge_base.user_id) == str(self.id):
                return True
            # 其他人需要申请权限
            return Permission.objects.filter(
                resource_type='knowledge',
                resource_id=str(knowledge_base.id),
                applicant=self,
                status='approved',
                expires_at__gt=timezone.now()
            ).exists()
        
        # 2. 如果是管理级知识库
        if knowledge_base.type == 'admin':
            # 管理员直接允许
            if self.role == 'admin':
                return True
            # 其他人需要申请权限
            return Permission.objects.filter(
                resource_type='knowledge',
                resource_id=str(knowledge_base.id),
                applicant=self,
                status='approved'
            ).exists()
        
        # 3. 如果是部门级知识库
        if knowledge_base.type == 'leader':
            # 同部门的管理员和组长可以访问
            if self.department == knowledge_base.department:
                return self.role in ['admin', 'leader']
            return False
        
        # 4. 如果是成员级知识库
        if knowledge_base.type == 'member':
            # 同部门的所有人可以访问
            return self.department == knowledge_base.department

class Data(models.Model):
    """统一的数据模型"""
    DATA_TYPES = (
        ('admin', '管理员数据'),
        ('leader', '组长数据'),
        ('member', '组员数据'),
    )
    
    # 基本信息
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    name = models.CharField(max_length=100, help_text="数据名称")
    desc = models.TextField(blank=True, help_text="数据描述")
    type = models.CharField(max_length=10, choices=DATA_TYPES)
    meta = models.JSONField(default=dict, blank=True, help_text="元数据")
    
    # 关联信息
    user_id = models.UUIDField(help_text="创建者ID")
    department = models.CharField(max_length=50, blank=True)
    
    # 统计信息
    char_length = models.IntegerField(null=True, blank=True, help_text="字符长度")
    document_count = models.IntegerField(null=True, blank=True, help_text="文档数量")
    application_mapping_count = models.IntegerField(default=0, help_text="应用映射数量")
    
    # 时间信息
    create_time = models.DateTimeField(auto_now_add=True)
    update_time = models.DateTimeField(auto_now=True)

    class Meta:
        db_table = 'user_data'
        indexes = [
            models.Index(fields=['type', 'user_id']),
            models.Index(fields=['create_time']),
        ]

    def __str__(self):
        return f"{self.name} ({self.get_type_display()})"

    def publish(self):
        """发布数据"""
        if self.status == 'draft':
            self.status = 'published'
            self.published_at = timezone.now()
            self.save()
            logger.info(f"Data {self.id} published by {self.owner.username}")

    def archive(self):
        """归档数据"""
        if self.status == 'published':
            self.status = 'archived'
            self.archived_at = timezone.now()
            self.save()
            logger.info(f"Data {self.id} archived by {self.owner.username}")

class Notification(models.Model):
    """通知模型"""
    NOTIFICATION_TYPES = (
        ('permission_request', '权限申请'),
        ('permission_approved', '权限批准'),
        ('permission_rejected', '权限拒绝'),
        ('permission_expired', '权限过期'),
        ('system_notice', '系统通知'),
    )
    
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    type = models.CharField(max_length=20, choices=NOTIFICATION_TYPES)
    title = models.CharField(max_length=100)
    content = models.TextField()
    sender = models.ForeignKey('User', on_delete=models.CASCADE, related_name='sent_notifications')
    receiver = models.ForeignKey('User', on_delete=models.CASCADE, related_name='received_notifications')
    is_read = models.BooleanField(default=False)
    related_resource = models.CharField(max_length=100, blank=True)  # 相关资源ID
    created_at = models.DateTimeField(auto_now_add=True)
    
    class Meta:
        ordering = ['-created_at']
        indexes = [
            models.Index(fields=['receiver', '-created_at']),
            models.Index(fields=['type', 'is_read']),
        ]

    def __str__(self):
        return f"{self.get_type_display()} - {self.title}"

class Permission(models.Model):
    """权限申请模型"""
    STATUS_CHOICES = (
        ('pending', '待审批'),
        ('approved', '已批准'),
        ('rejected', '已拒绝'),
    )

    knowledge_base = models.ForeignKey(
        'KnowledgeBase',
        on_delete=models.CASCADE,
        related_name='permissions',
        verbose_name='知识库'
    )
    applicant = models.ForeignKey(
        'User',
        on_delete=models.CASCADE,
        related_name='permission_applications',
        verbose_name='申请人'
    )
    approver = models.ForeignKey(
        'User',
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name='permission_approvals',
        verbose_name='审批人'
    )
    
    permissions = models.JSONField(default=dict, verbose_name='权限配置')  # {"can_read": true, "can_edit": false, "can_delete": false}
    status = models.CharField(max_length=20, choices=STATUS_CHOICES, default='pending', verbose_name='状态')
    reason = models.TextField(verbose_name='申请原因')
    response_message = models.TextField(null=True, blank=True, verbose_name='审批意见')
    expires_at = models.DateTimeField(null=True, blank=True, verbose_name='过期时间')  # 修改为可空
    
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    updated_at = models.DateTimeField(auto_now=True, verbose_name='更新时间')

    class Meta:
        verbose_name = '权限申请'
        verbose_name_plural = '权限申请'
        ordering = ['-created_at']

    def __str__(self):
        return f"{self.applicant} 申请 {self.knowledge_base} 的权限"

    def clean(self):
        """验证权限申请的合法性"""
        try:
            # 检查是否是自己的知识库
            if str(self.knowledge_base.user_id) == str(self.applicant.id):
                raise ValidationError('不能申请访问自己的知识库')
        except KnowledgeBase.DoesNotExist:
            raise ValidationError('知识库不存在')

    def save(self, *args, **kwargs):
        self.clean()
        super().save(*args, **kwargs)

    def approve(self, approver, response_message=''):
        """批准权限申请"""
        if self.status == 'pending':
            self.status = 'approved'
            self.approver = approver
            self.response_message = response_message
            self.save()

    def reject(self, approver, response_message=''):
        """拒绝权限申请"""
        if self.status == 'pending':
            self.status = 'rejected'
            self.approver = approver
            self.response_message = response_message
            self.save()

    def check_expiration(self):
        """检查权限是否过期"""
        if self.status == 'approved' and self.expires_at and self.expires_at < timezone.now():
            self.status = 'expired'
            self.save()
            logger.info(f"Permission {self.id} expired")

    def send_notification(self, notification_type, title, content):
        """发送通知"""
        Notification.objects.create(
            type=notification_type,
            title=title,
            content=content,
            sender=self.applicant,
            receiver=self.approver if notification_type == 'permission_request' else self.applicant,
            related_resource=str(self.id)
        )

    def notify_approver(self):
        """通知审批人"""
        self.send_notification(
            'permission_request',
            f'新的权限申请 - {self.get_resource_type_display()}',
            f'用户 {self.applicant.username} 申请访问 {self.get_resource_type_display()} 的权限'
        )

    def notify_applicant(self, status):
        """通知申请人审批结果"""
        notification_type = 'permission_approved' if status == 'approved' else 'permission_rejected'
        title = f'权限申请{self.get_status_display()} - {self.get_resource_type_display()}'
        content = f'您申请的 {self.get_resource_type_display()} 权限已{self.get_status_display()}'
        self.send_notification(notification_type, title, content)

class ChatHistory(models.Model):
    """聊天历史记录"""
    ROLE_CHOICES = [
        ('user', '用户'),
        ('assistant', 'AI助手'),
        ('system', '系统')
    ]

    user = models.ForeignKey(User, on_delete=models.CASCADE)
    # 保留与主知识库的关联
    knowledge_base = models.ForeignKey('KnowledgeBase', on_delete=models.CASCADE)
    # 用于标识知识库组合的对话
    conversation_id = models.CharField(max_length=100, db_index=True)
    parent_id = models.CharField(max_length=100, null=True, blank=True)
    role = models.CharField(max_length=20, choices=ROLE_CHOICES)
    content = models.TextField()
    tokens = models.IntegerField(default=0, help_text="消息token数")
    # 扩展metadata字段，用于存储知识库组合信息
    metadata = models.JSONField(default=dict, blank=True, help_text="""
    {
        'model_id': 'xxx',
        'dataset_id_list': ['id1', 'id2', ...],
        'dataset_external_id_list': ['ext1', 'ext2', ...],
        'primary_knowledge_base': 'id1'
    }
    """)
    created_at = models.DateTimeField(auto_now_add=True)
    is_deleted = models.BooleanField(default=False)

    class Meta:
        db_table = 'chat_history'
        ordering = ['created_at']
        indexes = [
            models.Index(fields=['conversation_id', 'created_at']),
            models.Index(fields=['user', 'created_at']),
            # 添加新的索引以支持知识库组合查询
            models.Index(fields=['conversation_id', 'is_deleted']),
        ]

    def __str__(self):
        return f"{self.user.username} - {self.knowledge_base.name} - {self.created_at}"

    @classmethod
    def get_conversation(cls, conversation_id):
        """获取完整对话历史"""
        return cls.objects.filter(
            conversation_id=conversation_id,
            is_deleted=False
        ).order_by('created_at')

    @classmethod
    def get_conversations_by_knowledge_bases(cls, dataset_ids, user):
        """根据知识库组合获取对话历史"""
        # 对知识库ID列表排序以确保一致性
        sorted_kb_ids = sorted(dataset_ids)
        conversation_id = str(uuid.uuid5(
            uuid.NAMESPACE_DNS, 
            '-'.join(sorted_kb_ids)
        ))
        
        return cls.objects.filter(
            conversation_id=conversation_id,
            user=user,
            is_deleted=False
        ).order_by('created_at')

    @classmethod
    def get_knowledge_base_combinations(cls, user):
        """获取用户的所有知识库组合"""
        return cls.objects.filter(
            user=user,
            is_deleted=False
        ).values('conversation_id').annotate(
            last_message=max('created_at'),
            message_count=count('id')
        ).values(
            'conversation_id',
            'last_message',
            'message_count',
            'metadata'
        ).order_by('-last_message')

    def get_knowledge_bases(self):
        """获取此消息关联的所有知识库"""
        if self.metadata and 'dataset_id_list' in self.metadata:
            return KnowledgeBase.objects.filter(
                id__in=self.metadata['dataset_id_list']
            )
        return KnowledgeBase.objects.filter(id=self.knowledge_base.id)

    def soft_delete(self):
        """软删除消息"""
        self.is_deleted = True
        self.save()

    def to_dict(self):
        """转换为字典格式"""
        return {
            'id': str(self.id),
            'conversation_id': self.conversation_id,
            'role': self.role,
            'content': self.content,
            'created_at': self.created_at.strftime('%Y-%m-%d %H:%M:%S'),
            'metadata': self.metadata,
            'knowledge_bases': [
                {
                    'id': str(kb.id),
                    'name': kb.name,
                    'type': kb.type
                } for kb in self.get_knowledge_bases()
            ]
        }

class UserProfile(models.Model):
    """用户档案模型"""
    user = models.OneToOneField(User, on_delete=models.CASCADE, related_name='profile')
    department = models.CharField(max_length=100, blank=True, help_text="部门")
    group = models.CharField(max_length=100, blank=True, help_text="小组")
    
    class Meta:
        db_table = 'user_profiles'

    def __str__(self):
        return f"{self.user.username}'s profile"

class KnowledgeBasePermission(models.Model):
    """知识库权限模型 - 实现知识库和用户的多对多关系"""
    STATUS_CHOICES = (
        ('active', '生效中'),
        ('expired', '已过期'),
        ('revoked', '已撤销'),
    )

    knowledge_base = models.ForeignKey(
        'KnowledgeBase', 
        on_delete=models.CASCADE,
        related_name='user_permissions',
        verbose_name='知识库'
    )
    user = models.ForeignKey(
        'User', 
        on_delete=models.CASCADE,
        related_name='knowledge_base_permissions',
        verbose_name='用户'
    )
    
    # 基础权限
    can_read = models.BooleanField(default=False, verbose_name='查看权限')
    can_edit = models.BooleanField(default=False, verbose_name='修改权限')
    can_delete = models.BooleanField(default=False, verbose_name='删除权限')
    
    # 权限状态
    status = models.CharField(
        max_length=10, 
        choices=STATUS_CHOICES, 
        default='active',
        verbose_name='状态'
    )
    
    # 授权信息
    granted_by = models.ForeignKey(
        'User',
        on_delete=models.SET_NULL,
        null=True,
        related_name='granted_permissions',
        verbose_name='授权人'
    )
    granted_at = models.DateTimeField(auto_now_add=True, verbose_name='授权时间')
    expires_at = models.DateTimeField(null=True, blank=True, verbose_name='过期时间')

    class Meta:
        app_label = 'user_management'
        db_table = 'knowledge_base_permissions'
        unique_together = ['knowledge_base', 'user']
        indexes = [
            models.Index(fields=['knowledge_base', 'user', 'status']),
        ]
        verbose_name = '知识库权限'
        verbose_name_plural = '知识库权限'

    def __str__(self):
        return f"{self.user.username} - {self.knowledge_base.name}"

    def is_valid(self):
        """检查权限是否有效"""
        if self.status != 'active':
            return False
        if self.expires_at and self.expires_at < timezone.now():
            self.status = 'expired'
            self.save()
            return False
        return True

class KnowledgeBase(models.Model):
    """知识库模型"""
    KNOWLEDGE_BASE_TYPES = [
        ('admin', '管理级知识库'),
        ('leader', '部门级知识库'),
        ('member', '成员级知识库'),
        ('private', '私有知识库'),
        ('secret', '公司级别私密知识库'),
    ]

    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    user_id = models.UUIDField(verbose_name='创建者ID')
    name = models.CharField(max_length=100, unique=True, verbose_name='知识库名称')
    desc = models.TextField(verbose_name='知识库描述', null=True, blank=True)
    type = models.CharField(
        max_length=20, 
        choices=KNOWLEDGE_BASE_TYPES,
        default='private',
        verbose_name='知识库类型'
    )
    department = models.CharField(max_length=50, null=True, blank=True)
    group = models.CharField(max_length=50, null=True, blank=True)
    documents = models.JSONField(default=list)
    char_length = models.IntegerField(default=0)
    document_count = models.IntegerField(default=0)
    external_id = models.UUIDField(null=True, blank=True)
    create_time = models.DateTimeField(auto_now_add=True)
    update_time = models.DateTimeField(auto_now=True)

    def is_owner(self, user):
        """检查用户是否是所有者（通过权限表检查）"""
        return str(user.id) == str(self.user_id) or KnowledgeBasePermission.objects.filter(
            knowledge_base=self,
            user=user,
            can_read=True,
            can_edit=True,
            can_delete=True,
            status='active'
        ).exists()

    def get_owners(self):
        """获取所有所有者（包括创建者和具有完整权限的用户）"""
        from .models import User
        # 获取创建者
        owners = [self.user_id]
        # 获取具有完整权限的用户
        permission_owners = KnowledgeBasePermission.objects.filter(
            knowledge_base=self,
            can_read=True,
            can_edit=True,
            can_delete=True,
            status='active'
        ).values_list('user_id', flat=True)
        owners.extend(permission_owners)
        return User.objects.filter(id__in=set(owners))

    def calculate_stats(self):
        """计算文档统计信息"""
        total_chars = 0
        doc_count = 0
        
        if self.documents:
            doc_count = len(self.documents)
            for doc in self.documents:
                if 'paragraphs' in doc:
                    for para in doc['paragraphs']:
                        if 'content' in para:
                            total_chars += len(para['content'])
                        if 'title' in para:
                            total_chars += len(para['title'])
        
        return doc_count, total_chars

    def save(self, *args, **kwargs):
        """重写保存方法"""
        # 只在创建时计算统计信息
        if not self.pk:  # 如果是新实例
            doc_count, char_count = self.calculate_stats()
            self.document_count = doc_count
            self.char_length = char_count
        
        # 直接调用Model的save方法
        models.Model.save(self, *args, **kwargs)

    @classmethod
    def update_external_id(cls, instance_id, external_id):
        """更新外部ID的静态方法"""
        cls.objects.filter(id=instance_id).update(external_id=external_id)

    class Meta:
        db_table = 'knowledge_bases'
        indexes = [
            models.Index(fields=['type']),
            models.Index(fields=['department']),
            models.Index(fields=['group'])
        ]

    def __str__(self):
        return f"{self.name} ({self.get_type_display()})"

    def clean(self):
        """验证知识库类型与创建者权限是否匹配"""
        try:
            user = User.objects.get(id=self.user_id)
            if user.role == 'member' and self.type != 'private':
                raise ValidationError('组员只能创建私人知识库')
            if user.role == 'leader' and self.type not in ['member', 'private']:
                raise ValidationError('组长只能创建成员级或私人知识库')
        except User.DoesNotExist:
            raise ValidationError('创建者不存在')

    def to_response_dict(self):
        """转换为API响应格式"""
        return {
            "create_time": self.create_time.isoformat(),
            "update_time": self.update_time.isoformat(),
            "id": str(self.id),
            "name": self.name,
            "desc": self.desc,
            "type": self.type,
            "meta": self.meta,
            "user_id": str(self.user_id),
            "embedding_mode_id": str(self.embedding_mode_id),
            "char_length": self.char_length,
            "application_mapping_count": self.application_mapping_count,
            "document_count": self.document_count
        }

    def to_external_format(self):
        """转换为外部接口格式"""
        return {
            "name": self.name,
            "desc": self.desc,
            "documents": self.documents
        }

    @classmethod
    def from_external_format(cls, data, user_id, embedding_mode_id):
        """从外部接口格式创建实例"""
        return cls(
            name=data["name"],
            desc=data["desc"],
            documents=data["documents"],
            user_id=user_id,
            embedding_mode_id=embedding_mode_id,
            document_count=len(data["documents"]) if data.get("documents") else 0
        )

class FeishuCreator(models.Model):
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    record_id = models.CharField(max_length=100, unique=True, verbose_name='飞书记录ID')
    
    # 对接人信息
    contact_person = models.CharField(max_length=50, blank=True, verbose_name='对接人')
    
    # 基本账号信息
    handle = models.TextField(blank=True, verbose_name='Handle')
    tiktok_url = models.TextField(blank=True, verbose_name='链接')
    fans_count = models.CharField(max_length=50, blank=True, verbose_name='粉丝数')
    gmv = models.CharField(max_length=100, blank=True, verbose_name='GMV')
    
    # 联系方式
    email = models.EmailField(blank=True, verbose_name='邮箱')
    phone = models.CharField(max_length=50, blank=True, verbose_name='手机号|WhatsApp')
    
    # 账号属性和报价
    account_type = models.CharField(max_length=50, blank=True, verbose_name='账号属性')
    price_quote = models.TextField(blank=True, verbose_name='报价')
    response_speed = models.CharField(max_length=50, blank=True, verbose_name='回复速度')
    
    # 合作相关
    cooperation_intention = models.CharField(max_length=50, blank=True, verbose_name='合作意向')
    payment_method = models.CharField(max_length=50, blank=True, verbose_name='支付方式')
    payment_account = models.CharField(max_length=100, blank=True, verbose_name='收款账号')
    address = models.TextField(blank=True, verbose_name='收件地址')
    has_ooin = models.CharField(max_length=10, blank=True, verbose_name='签约OOIN?')
    
    # 渠道和进度
    source = models.CharField(max_length=100, blank=True, verbose_name='渠道来源')
    contact_status = models.CharField(max_length=50, blank=True, verbose_name='建联进度')
    cooperation_brands = models.JSONField(default=list, blank=True, verbose_name='合作品牌')
    
    # 品类信息
    system_categories = models.CharField(max_length=100, blank=True, verbose_name='系统展示的带货品类')
    actual_categories = models.CharField(max_length=100, blank=True, verbose_name='实际高播放量带货品类')
    human_categories = models.CharField(max_length=100, blank=True, verbose_name='达人标想要货品类')
    
    # 其他信息
    creator_base = models.CharField(max_length=100, blank=True, verbose_name='达人base')
    notes = models.TextField(blank=True, verbose_name='父记录')
    
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    updated_at = models.DateTimeField(auto_now=True, verbose_name='更新时间')

    class Meta:
        db_table = 'feishu_creators'
        verbose_name = '创作者数据'
        verbose_name_plural = '创作者数据'

class KnowledgeBaseDocument(models.Model):
    """知识库文档关联模型"""
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    knowledge_base = models.ForeignKey(
        'KnowledgeBase', 
        on_delete=models.CASCADE,
        related_name='kb_documents',
        verbose_name='知识库'
    )
    document_id = models.CharField(max_length=100, verbose_name='文档ID')
    document_name = models.CharField(max_length=255, verbose_name='文档名称')
    external_id = models.CharField(max_length=100, verbose_name='外部文档ID')
    status = models.CharField(
        max_length=20, 
        default='active',
        choices=[
            ('active', '有效'),
            ('deleted', '已删除')
        ],
        verbose_name='状态'
    )
    create_time = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    update_time = models.DateTimeField(auto_now=True, verbose_name='更新时间')

    class Meta:
        db_table = 'knowledge_base_documents'
        unique_together = ['knowledge_base', 'document_id']
        indexes = [
            models.Index(fields=['knowledge_base', 'status']),
            models.Index(fields=['document_id']),
            models.Index(fields=['external_id'])
        ]
        verbose_name = '知识库文档'
        verbose_name_plural = '知识库文档'

    def __str__(self):
        return f"{self.knowledge_base.name} - {self.document_name}"