JARVIS/backend/app/routers/agent.py

from datetime import datetime
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import get_db
from app.models.agent import Agent
from app.models.conversation import Conversation
from app.models.skill import Skill
from app.models.user import User
from app.routers.auth import get_current_user
from app.schemas.agent import (
    AgentConfigOut,
    AgentConfigUpdate,
    AgentCreate,
    AgentOut,
    AgentStats,
    AgentVisibilityEvidenceOut,
    AgentVisibilityEventsResponse,
    AgentVisibilityEventOut,
    AgentVisibilityTaskSummaryOut,
    AgentVisibilityThreadMessageOut,
    AgentVisibilityThreadOut,
    AgentVisibilityTopologyNodeOut,
    AgentVisibilityTopologyOut,
    AgentVisibilityVerifierOut,
)
from app.services.agent_service import _extract_continuity_snapshot

router = APIRouter(prefix="/api/agents", tags=["Agent"])

_agent_call_counts: dict[str, int] = {}
_agent_current_tasks: dict[str, str | None] = {}
_agent_statuses: dict[str, str] = {}

DEFAULT_AGENT_ROLES = ["master", "schedule_planner", "executor", "librarian", "analyst"]
SUB_COMMANDERS_BY_ROLE = {
    "schedule_planner": ["schedule_analysis", "schedule_planning"],
    "executor": ["executor_tasks", "executor_forum"],
    "librarian": ["librarian_retrieval", "librarian_graph"],
    "analyst": ["analyst_progress", "analyst_insights"],
}
ALLOWED_AGENT_ROLES = set(DEFAULT_AGENT_ROLES) | {
    role
    for sub_roles in SUB_COMMANDERS_BY_ROLE.values()
    for role in sub_roles
}


def _parse_visibility_datetime(value: str | None) -> datetime | None:
    if value is None:
        return None
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError as exc:
        raise HTTPException(status_code=400, detail="时间参数必须是 ISO 8601 格式") from exc


async def _get_visibility_state(
    conversation_id: str,
    *,
    current_user: User,
    db: AsyncSession,
) -> dict[str, Any]:
    result = await db.execute(
        select(Conversation).where(
            Conversation.id == conversation_id,
            Conversation.user_id == current_user.id,
        )
    )
    conversation = result.scalar_one_or_none()
    if conversation is None:
        raise HTTPException(status_code=404, detail="对话不存在")
    snapshot = _extract_continuity_snapshot(conversation.agent_state)
    if snapshot is None:
        raise HTTPException(status_code=404, detail="当前会话暂无可视化运行时数据")
    return snapshot


def _coerce_event_payload(event: dict[str, Any]) -> AgentVisibilityEventOut:
    return AgentVisibilityEventOut.model_validate(event)


def _filter_events(
    events: list[dict[str, Any]],
    *,
    agent_id: str | None,
    thread_id: str | None,
    event_type: str | None,
    started_after: datetime | None,
    ended_before: datetime | None,
) -> list[dict[str, Any]]:
    filtered: list[dict[str, Any]] = []
    for event in events:
        if agent_id and event.get("agent_id") != agent_id:
            continue
        if thread_id and event.get("thread_id") != thread_id:
            continue
        if event_type and event.get("event_type") != event_type:
            continue
        timestamp_raw = event.get("timestamp")
        timestamp = None
        if isinstance(timestamp_raw, str):
            try:
                timestamp = datetime.fromisoformat(timestamp_raw.replace("Z", "+00:00"))
            except ValueError:
                timestamp = None
        if started_after and timestamp and timestamp < started_after:
            continue
        if ended_before and timestamp and timestamp > ended_before:
            continue
        filtered.append(event)
    return filtered


def _summarize_tasks(tasks: list[dict[str, Any]], task_results: list[dict[str, Any]]) -> list[AgentVisibilityTaskSummaryOut]:
    result_by_task_id = {item.get("task_id"): item for item in task_results}
    summaries: list[AgentVisibilityTaskSummaryOut] = []
    for task in tasks:
        task_id = str(task.get("task_id") or "")
        result = result_by_task_id.get(task_id) or {}
        evidence = result.get("evidence") or task.get("evidence") or []
        summaries.append(
            AgentVisibilityTaskSummaryOut(
                task_id=task_id,
                role=task.get("role"),
                owner_agent_id=task.get("owner_agent_id") or result.get("owner_agent_id"),
                status=result.get("status") or task.get("status"),
                summary=result.get("summary") or task.get("result_summary"),
                evidence_count=len(evidence),
            )
        )
    return summaries


def _build_topology_nodes(
    state: dict[str, Any],
    tasks: list[dict[str, Any]],
    task_results: list[dict[str, Any]],
) -> list[AgentVisibilityTopologyNodeOut]:
    task_counts: dict[str, int] = {}
    completed_counts: dict[str, int] = {}
    for task in tasks:
        owner = str(task.get("owner_agent_id") or "")
        if owner:
            task_counts[owner] = task_counts.get(owner, 0) + 1
    for result in task_results:
        owner = str(result.get("owner_agent_id") or "")
        if owner and result.get("status") == "completed":
            completed_counts[owner] = completed_counts.get(owner, 0) + 1

    root_agent_id = str(state.get("root_agent_id") or state.get("agent_id") or "") or None
    current_agent = str(state.get("current_agent") or "") or None
    nodes: dict[str, AgentVisibilityTopologyNodeOut] = {}
    if root_agent_id:
        nodes[root_agent_id] = AgentVisibilityTopologyNodeOut(
            agent_id=root_agent_id,
            role=root_agent_id.split("-")[0],
            parent_agent_id=None,
            source="root",
            task_count=task_counts.get(root_agent_id, 0),
            completed_task_count=completed_counts.get(root_agent_id, 0),
        )
    for agent_id in state.get("spawned_agent_ids") or []:
        agent_id = str(agent_id)
        nodes[agent_id] = AgentVisibilityTopologyNodeOut(
            agent_id=agent_id,
            role=agent_id.split("-")[0],
            parent_agent_id=root_agent_id,
            source="spawned",
            task_count=task_counts.get(agent_id, 0),
            completed_task_count=completed_counts.get(agent_id, 0),
        )
    if current_agent and current_agent not in nodes:
        nodes[current_agent] = AgentVisibilityTopologyNodeOut(
            agent_id=current_agent,
            role=current_agent.split("-")[0],
            parent_agent_id=None if current_agent == root_agent_id else root_agent_id,
            source="current",
            task_count=task_counts.get(current_agent, 0),
            completed_task_count=completed_counts.get(current_agent, 0),
        )
    return list(nodes.values())


def record_agent_call(agent_id: str):
    _agent_call_counts[agent_id] = _agent_call_counts.get(agent_id, 0) + 1


def set_agent_task(agent_id: str, task: str | None):
    _agent_current_tasks[agent_id] = task
    _agent_statuses[agent_id] = "active" if task else "idle"


def set_agent_status(agent_id: str, status: str):
    _agent_statuses[agent_id] = status


def _build_agent_stats(agent_id: str) -> AgentStats:
    return AgentStats(
        agent_id=agent_id,
        call_count=_agent_call_counts.get(agent_id, 0),
        current_task=_agent_current_tasks.get(agent_id),
        status=_agent_statuses.get(agent_id, "idle"),
    )


@router.get("", response_model=list[AgentOut])
async def list_agents(
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    result = await db.execute(
        select(Agent).where(Agent.is_active == True).order_by(Agent.role)
    )
    return result.scalars().all()


@router.get("/stats", response_model=list[AgentStats])
async def get_agent_stats(
    current_user: User = Depends(get_current_user),
):
    return [_build_agent_stats(role) for role in DEFAULT_AGENT_ROLES]


@router.get("/stats/hierarchy")
async def get_agent_hierarchy_stats(
    current_user: User = Depends(get_current_user),
):
    main_agents = []
    for role in DEFAULT_AGENT_ROLES:
        if role == "master":
            continue
        node = _build_agent_stats(role).model_dump()
        node["sub_commanders"] = [
            _build_agent_stats(sub_id).model_dump()
            for sub_id in SUB_COMMANDERS_BY_ROLE.get(role, [])
        ]
        main_agents.append(node)
    return {"main_agents": main_agents}


@router.get("/config/{agent_id}", response_model=AgentConfigOut)
async def get_agent_config(
    agent_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    result = await db.execute(select(Agent).where(Agent.role == agent_id))
    agent = result.scalar_one_or_none()

    if not agent:
        from app.agents.prompts import MASTER_SYSTEM_PROMPT, SCHEDULE_PLANNER_SYSTEM_PROMPT, EXECUTOR_SYSTEM_PROMPT, LIBRARIAN_SYSTEM_PROMPT, ANALYST_SYSTEM_PROMPT
        defaults = {
            "master": ("JARVIS", "主控制核心", MASTER_SYSTEM_PROMPT),
            "schedule_planner": ("SCHEDULE PLANNER", "日程规划师", SCHEDULE_PLANNER_SYSTEM_PROMPT),
            "executor": ("EXECUTOR", "执行专家", EXECUTOR_SYSTEM_PROMPT),
            "librarian": ("LIBRARIAN", "知识管理员", LIBRARIAN_SYSTEM_PROMPT),
            "analyst": ("ANALYST", "数据分析师", ANALYST_SYSTEM_PROMPT),
        }
        if agent_id not in defaults:
            raise HTTPException(status_code=404, detail="Agent 不存在")
        name, desc, prompt = defaults[agent_id]
        return AgentConfigOut(
            id=agent_id,
            name=name,
            role=agent_id,
            description=desc,
            system_prompt=prompt,
            enabled=True,
            is_active=True,
            selected_skill_ids=[],
        )
    return AgentConfigOut(
        id=agent.role,
        name=agent.name,
        role=agent.role,
        description=agent.description,
        system_prompt=agent.system_prompt,
        enabled=agent.is_active,
        is_active=agent.is_active,
        selected_skill_ids=agent.selected_skill_ids or [],
    )


@router.put("/config/{agent_id}", response_model=AgentConfigOut)
async def update_agent_config(
    agent_id: str,
    data: AgentConfigUpdate,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    result = await db.execute(select(Agent).where(Agent.role == agent_id))
    agent = result.scalar_one_or_none()

    if not agent:
        raise HTTPException(status_code=404, detail="Agent 不存在")

    if data.name is not None:
        agent.name = data.name
    if data.description is not None:
        agent.description = data.description
    if data.system_prompt is not None:
        agent.system_prompt = data.system_prompt
    if data.enabled is not None:
        agent.is_active = data.enabled
        _agent_statuses[agent_id] = "disabled" if not data.enabled else "idle"
    if data.selected_skill_ids is not None:
        if data.selected_skill_ids:
            result = await db.execute(
                select(Skill.id).where(
                    Skill.id.in_(data.selected_skill_ids),
                    Skill.owner_id == current_user.id,
                )
            )
            allowed_skill_ids = set(result.scalars().all())
            invalid_skill_ids = [skill_id for skill_id in data.selected_skill_ids if skill_id not in allowed_skill_ids]
            if invalid_skill_ids:
                raise HTTPException(status_code=400, detail="存在无效的技能绑定")
        agent.selected_skill_ids = data.selected_skill_ids

    await db.commit()
    await db.refresh(agent)
    return AgentConfigOut(
        id=agent.role,
        name=agent.name,
        role=agent.role,
        description=agent.description,
        system_prompt=agent.system_prompt,
        enabled=agent.is_active,
        is_active=agent.is_active,
        selected_skill_ids=agent.selected_skill_ids or [],
    )


@router.get("/visibility/events", response_model=AgentVisibilityEventsResponse)
async def get_visibility_events(
    conversation_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
    agent_id: str | None = None,
    thread_id: str | None = None,
    event_type: str | None = None,
    started_after: str | None = None,
    ended_before: str | None = None,
    limit: int = Query(default=50, ge=1, le=200),
    offset: int = Query(default=0, ge=0),
):
    state = await _get_visibility_state(conversation_id, current_user=current_user, db=db)
    events = [dict(item) for item in state.get("event_trace") or []]
    filtered = _filter_events(
        events,
        agent_id=agent_id,
        thread_id=thread_id,
        event_type=event_type,
        started_after=_parse_visibility_datetime(started_after),
        ended_before=_parse_visibility_datetime(ended_before),
    )
    paged = filtered[offset:offset + limit]
    return AgentVisibilityEventsResponse(
        conversation_id=conversation_id,
        total=len(filtered),
        limit=limit,
        offset=offset,
        items=[_coerce_event_payload(item) for item in paged],
    )


@router.get("/visibility/topology", response_model=AgentVisibilityTopologyOut)
async def get_visibility_topology(
    conversation_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    state = await _get_visibility_state(conversation_id, current_user=current_user, db=db)
    tasks = [dict(item) for item in state.get("active_tasks") or []]
    task_results = [dict(item) for item in state.get("task_results") or []]
    nodes = _build_topology_nodes(state, tasks, task_results)
    root_agent_id = str(state.get("root_agent_id") or state.get("agent_id") or "") or None
    edges = [
        {"parent_agent_id": root_agent_id, "child_agent_id": node.agent_id}
        for node in nodes
        if node.parent_agent_id and root_agent_id and node.agent_id != root_agent_id
    ]
    return AgentVisibilityTopologyOut(
        conversation_id=conversation_id,
        root_agent_id=root_agent_id,
        current_agent=str(state.get("current_agent") or "") or None,
        nodes=nodes,
        edges=edges,
        tasks=_summarize_tasks(tasks, task_results),
        task_hierarchy=dict(state.get("task_hierarchy") or {}),
    )


@router.get("/visibility/tasks/{task_id}/evidence", response_model=AgentVisibilityEvidenceOut)
async def get_visibility_task_evidence(
    task_id: str,
    conversation_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    state = await _get_visibility_state(conversation_id, current_user=current_user, db=db)
    tasks = [dict(item) for item in state.get("active_tasks") or []]
    task = next((item for item in tasks if item.get("task_id") == task_id), None)
    task_results = [dict(item) for item in state.get("task_results") or []]
    result = next((item for item in task_results if item.get("task_id") == task_id), None)
    if task is None and result is None:
        raise HTTPException(status_code=404, detail="任务不存在")
    tool_outcomes = [
        dict(evidence)
        for evidence in (result or {}).get("evidence") or []
        if isinstance(evidence, dict) and evidence.get("tool_name")
    ]
    verification_entry = next(
        (
            dict(evidence)
            for evidence in (result or {}).get("evidence") or []
            if isinstance(evidence, dict) and evidence.get("type") == "verification"
        ),
        None,
    )
    verifier = {
        "status": (verification_entry or {}).get("status"),
        "summary": (verification_entry or {}).get("summary"),
        "evidence": [dict(item) for item in state.get("verification_evidence") or [] if item.get("task_id") == task_id],
    }
    return AgentVisibilityEvidenceOut(
        conversation_id=conversation_id,
        task_id=task_id,
        task=task,
        result=result,
        tool_outcomes=tool_outcomes,
        verifier=verifier,
    )


@router.get("/visibility/threads/{thread_id}/messages", response_model=AgentVisibilityThreadOut)
async def get_visibility_thread_messages(
    thread_id: str,
    conversation_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    state = await _get_visibility_state(conversation_id, current_user=current_user, db=db)
    items = [
        AgentVisibilityThreadMessageOut.model_validate(item)
        for item in state.get("message_trace") or []
        if isinstance(item, dict) and item.get("thread_id") == thread_id
    ]
    if not items:
        raise HTTPException(status_code=404, detail="线程不存在")
    return AgentVisibilityThreadOut(
        conversation_id=conversation_id,
        thread_id=thread_id,
        total=len(items),
        items=items,
    )


@router.get("/visibility/verifier", response_model=AgentVisibilityVerifierOut)
async def get_visibility_verifier(
    conversation_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    state = await _get_visibility_state(conversation_id, current_user=current_user, db=db)
    return AgentVisibilityVerifierOut(
        conversation_id=conversation_id,
        status=state.get("verification_status"),
        summary=state.get("verification_summary"),
        evidence=list(state.get("verification_evidence") or []),
    )


@router.post("", response_model=AgentOut, status_code=201)
async def create_agent(
    data: AgentCreate,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    if not current_user.is_superuser:
        raise HTTPException(status_code=403, detail="仅管理员可创建 Agent")
    if not data.spawn_permission:
        raise HTTPException(status_code=400, detail="缺少 spawn_permission，禁止直接创建 runtime agent")
    if data.role not in ALLOWED_AGENT_ROLES:
        raise HTTPException(status_code=400, detail="不支持的 Agent 角色")

    agent = Agent(
        name=data.name,
        role=data.role,
        description=data.description,
        system_prompt=data.system_prompt,
    )
    db.add(agent)
    await db.commit()
    await db.refresh(agent)
    return agent


@router.get("/{agent_id}", response_model=AgentOut)
async def get_agent(
    agent_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    result = await db.execute(select(Agent).where(Agent.id == agent_id))
    agent = result.scalar_one_or_none()
    if not agent:
        raise HTTPException(status_code=404, detail="Agent 不存在")
    return agent