from app.models.user import User
from auth.deps import get_current_user_id
from fastapi import Request
from strawberry.fastapi import BaseContext
from strawberry.dataloader import DataLoader
from typing import List, Optional
from collections import defaultdict
from app.db import SessionLocal
from app.models import WorkspaceMember
from typing import AsyncGenerator

# 定义批量查询函数：接收一组 workspace ID 列表，返回相同长度的结果列表
async def load_members_by_workspace_ids(workspace_ids: List[str]) -> List[List[WorkspaceMember]]:
    db = SessionLocal()
    try:
        # 用一条 IN 语句查出所有相关成员
        members = db.query(WorkspaceMember).filter(WorkspaceMember.workspace_id.in_(workspace_ids)).all()
        
        # 将结果归类到对应的 workspace ID
        member_map = defaultdict(list)
        for member in members:
            member_map[member.workspace_id].append(member)
        
        # 严格按照传入的 workspace ID 顺序返回结果列表
        return [member_map[w_id] for w_id in workspace_ids]
    finally:
        db.close()

class CustomContext(BaseContext):
    def __init__(self, request: Request):
        super().__init__(request)
        self.db = SessionLocal()
        self.request = request
        self._current_user: Optional[User] = None
        
        self.members_loader = DataLoader(load_fn=load_members_by_workspace_ids)

    async def get_current_user(self) -> User | None:
        """在 Context 级别负责把 user_id 变成真正的数据库 User 对象"""
        if self._current_user is not None:
            return self._current_user
            
        try:
            # 借用 auth 包的纯工具解出 ID
            user_id = get_current_user_id(self.request)
            # 在 Context 的数据库连接中查询完整用户
            user = self.db.query(User).filter(User.id == user_id).first()
            self._current_user = user
            return user
        except Exception:
            return None


async def get_context(request: Request) -> AsyncGenerator[CustomContext, None]:
    context = CustomContext(request)
    try:
        yield context
    finally:
        context.db.close()  # 确保请求结束时关闭数据库连接