import time

from sqlalchemy.ext.asyncio import AsyncSession

from app.models import Agent, AgentTeam, CollaborationRun
from app.services.llm import LLMService


class CollaborationService:
    def __init__(self, db: AsyncSession):
        self.db = db
        self.llm = LLMService(db)

    async def run(self, team: AgentTeam, task: str) -> CollaborationRun:
        started = time.time()
        run = CollaborationRun(team_id=team.id, task=task, status="running", messages=[], final_answer="")
        self.db.add(run)
        await self.db.flush()

        messages = []
        context = task
        for member in team.members or []:
            agent_id = member.get("agent_id")
            role = member.get("role", "member")
            agent = await self.db.get(Agent, agent_id) if agent_id else None
            if not agent:
                content = f"{role}: 未配置有效智能体，跳过。"
            else:
                content = await self.llm.complete(agent, f"你的协作角色是 {role}。请基于上下文完成任务：{context}")
            messages.append({"role": role, "agent_id": agent_id, "agent_name": agent.name if agent else "", "content": content})
            context = f"{context}\n\n[{role}] {content}"

        run.messages = messages
        run.final_answer = messages[-1]["content"] if messages else "团队暂无成员，无法执行协作任务。"
        run.status = "completed"
        run.elapsed_time = int((time.time() - started) * 1000)
        await self.db.commit()
        await self.db.refresh(run)
        return run