from app.models.user import User
from app.auth.deps import get_current_user
from fastapi import Request
from strawberry.fastapi import BaseContext
from strawberry.dataloader import DataLoader
from typing import List
from collections import defaultdict
from app.db import SessionLocal
from app.models import WorkspaceMember
from app.graphql.types import WorkspaceMemberType
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

    async def get_current_user(self) -> User | None:
        # 复用现有的 auth 逻辑
        try:
            return await get_current_user(self.request)
        except Exception:
            return None


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