"""Task submission routes"""

import asyncio
import json
import uuid
from dataclasses import asdict
from datetime import datetime, timezone

from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel
from typing import Any

from agentkit.core.protocol import TaskMessage, TaskStatus
from agentkit.core.spec_manager import SpecManager

router = APIRouter(tags=["tasks"])


class SubmitTaskRequest(BaseModel):
    input_data: dict[str, Any]
    skill_name: str | None = None
    agent_name: str | None = None
    mode: str = "sync"  # "sync" or "async"

    # 输入数据大小限制（防止 OOM）
    model_config = {"json_schema_extra": {"max_input_size_bytes": 1024 * 1024}}  # 1MB


# 允许的 custom_handler 模块前缀白名单
_ALLOWED_HANDLER_PREFIXES = (
    "agentkit.",
    "app.agent_framework.",
)


def _validate_input_size(input_data: dict) -> None:
    """验证输入数据大小，防止超大 payload"""
    import json
    size = len(json.dumps(input_data, default=str).encode("utf-8"))
    if size > 1024 * 1024:  # 1MB
        raise HTTPException(
            status_code=413,
            detail=f"Input data too large: {size} bytes (max 1MB)",
        )


@router.get("/tasks")
async def list_tasks(status: str | None = None, limit: int = 100, req: Request = None):
    """List tasks"""
    store = req.app.state.task_store
    task_status = TaskStatus(status) if status else None
    records = store.list_tasks(status=task_status, limit=limit)
    return [r.to_dict() for r in records]


@router.post("/tasks")
async def submit_task(request: SubmitTaskRequest, req: Request):
    """Submit a task (Intent Router auto-routes to skill)"""
    # 输入大小验证
    _validate_input_size(request.input_data)

    pool = req.app.state.agent_pool
    skill_registry = req.app.state.skill_registry
    quality_gate = req.app.state.quality_gate
    output_standardizer = req.app.state.output_standardizer

    agent = None
    skill = None

    # 1. If agent_name specified, use that agent directly
    if request.agent_name:
        agent = pool.get_agent(request.agent_name)
        if agent is None:
            raise HTTPException(
                status_code=404,
                detail=f"Agent '{request.agent_name}' not found",
            )
        # Find the skill for this agent if available
        if agent._skill:
            skill = agent._skill

    # 2. If skill_name specified, use that skill
    elif request.skill_name:
        try:
            skill = skill_registry.get(request.skill_name)
        except Exception:
            raise HTTPException(
                status_code=404,
                detail=f"Skill '{request.skill_name}' not found",
            )
        # Get or create agent for this skill
        agent = pool.get_agent(request.skill_name)
        if agent is None:
            agent = await pool.create_agent_from_skill(request.skill_name)

    # 3. Otherwise, use default agent (REACT mode — LLM decides autonomously)
    else:
        all_skills = skill_registry.list_skills()
        if not all_skills:
            raise HTTPException(
                status_code=400,
                detail="No skills registered and no skill_name or agent_name specified",
            )
        agent = pool.get_agent("default")
        if agent is None:
            # Fallback: create from first available skill
            agent = await pool.create_agent_from_skill(all_skills[0].name)

    # 4. Async mode: submit to background runner
    if request.mode == "async":
        runner = req.app.state.runner
        task_id = await runner.submit(
            agent=agent,
            input_data=request.input_data,
            skill_name=request.skill_name,
            quality_gate=quality_gate,
            output_standardizer=output_standardizer,
            skill=skill,
        )
        return {"task_id": task_id, "status": "pending", "mode": "async"}

    # 5. Sync mode: existing blocking execution
    task = TaskMessage(
        task_id=str(uuid.uuid4()),
        agent_name=agent.name,
        task_type=agent.agent_type,
        priority=0,
        input_data=request.input_data,
        callback_url=None,
        created_at=datetime.now(timezone.utc),
    )

    task_result = await agent.execute(task)

    # 6. Run quality gate if skill available
    quality_result = None
    if skill:
        try:
            intent = getattr(skill.config, "intent", None)
            skill_context = None
            if intent is not None:
                keywords = list(intent.keywords) + list(intent.disambiguation_keywords)
                if keywords:
                    skill_context = {"intent_keywords": keywords}
            quality_result = await quality_gate.validate(
                task_result.output_data or {}, skill, skill_context=skill_context
            )
        except Exception:
            pass  # Quality gate failure shouldn't block the response

    # 7. Standardize output if skill available
    if skill:
        try:
            standard_output = await output_standardizer.standardize(
                raw_output=task_result.output_data or {},
                skill=skill,
                quality_result=quality_result,
            )
            return {
                "skill_name": standard_output.skill_name,
                "data": standard_output.data,
                "metadata": {
                    "version": standard_output.metadata.version,
                    "produced_at": standard_output.metadata.produced_at.isoformat(),
                    "quality_score": standard_output.metadata.quality_score,
                },
                "task_id": task.task_id,
                "status": task_result.status,
            }
        except Exception:
            pass  # Fall through to raw output

    # 8. Return raw result if no skill or standardization failed
    return {
        "task_id": task.task_id,
        "status": task_result.status,
        "output": task_result.output_data,
        "error_message": task_result.error_message,
    }


@router.get("/tasks/{task_id}")
async def get_task_status(task_id: str, req: Request):
    """Get task status and result"""
    store = req.app.state.task_store
    record = store.get(task_id)
    if record is None:
        raise HTTPException(status_code=404, detail=f"Task '{task_id}' not found")
    return record.to_dict()


@router.post("/tasks/{task_id}/cancel")
async def cancel_task(task_id: str, req: Request):
    """Cancel a running task"""
    runner = req.app.state.runner

    # First, try cooperative cancellation via agent's CancellationToken
    pool = req.app.state.agent_pool
    agent_cancelled = False
    for agent in pool._agents.values() if hasattr(pool, '_agents') else []:
        if agent.cancel_task(task_id):
            agent_cancelled = True
            break

    # Also cancel the asyncio task via runner
    runner_cancelled = await runner.cancel(task_id)

    if not agent_cancelled and not runner_cancelled:
        raise HTTPException(status_code=400, detail="Task cannot be cancelled (not running or not found)")
    return {"task_id": task_id, "status": "cancelled"}


@router.post("/tasks/{task_id}/resume")
async def resume_task(task_id: str, req: Request, plan_id: str | None = None):
    """Resume a crashed pipeline from the last completed phase checkpoint.

    Reconstructs the team from the saved plan's expert names, creates a new
    TeamOrchestrator with the checkpoint manager, and calls resume().

    Args:
        task_id: Task ID from the URL path.
        plan_id: Optional plan ID. If not provided, falls back to task_id.
                 Needed because TeamPlan.id is auto-generated and may differ
                 from the task_id.
    """
    from agentkit.experts.orchestrator import TeamOrchestrator
    from agentkit.experts.router import ExpertTeamRouter
    from agentkit.experts.team import ExpertTeam
    from agentkit.orchestrator.checkpoint import PipelineCheckpoint

    app_state = req.app.state

    # Resolve plan_id: explicit param > task_id fallback
    resolved_plan_id = plan_id or task_id

    # P2 #10: 并发 resume 防护 — 同一 plan_id 的并发 resume 请求只允许一个执行
    lock_attr = f"_resume_lock_{resolved_plan_id}"
    lock = getattr(app_state, lock_attr, None)
    if lock is None:
        lock = asyncio.Lock()
        setattr(app_state, lock_attr, lock)
    if lock.locked():
        raise HTTPException(
            status_code=409,
            detail=f"Resume already in progress for plan '{resolved_plan_id}'",
        )
    async with lock:
        # 1. Create checkpoint manager
        checkpoint = PipelineCheckpoint(
            redis_client=getattr(app_state, "working_redis_client", None)
        )

        # 2. Load plan to get expert names
        plan_dict = await checkpoint.load_plan(resolved_plan_id)
        if plan_dict is None:
            raise HTTPException(
                status_code=404,
                detail=f"No checkpoint found for plan '{resolved_plan_id}'"
                f" (task_id='{task_id}')",
            )

        # 3. Extract unique expert names from plan
        expert_names: list[str] = []
        lead_name = plan_dict.get("lead_expert", "")
        if lead_name:
            expert_names.append(lead_name)
        for ph in plan_dict.get("phases", []):
            name = ph.get("assigned_expert", "")
            if name and name not in expert_names:
                expert_names.append(name)

        if not expert_names:
            raise HTTPException(
                status_code=400,
                detail="Cannot resume: no experts found in saved plan",
            )

        # 4. Resolve expert configs via ExpertTeamRouter
        template_registry = getattr(app_state, "expert_template_registry", None)
        if template_registry is None:
            from agentkit.experts.registry import ExpertTemplateRegistry

            template_registry = ExpertTemplateRegistry()

        team_router = ExpertTeamRouter(template_registry=template_registry)
        expert_configs = team_router.resolve_expert_configs(expert_names)
        if not expert_configs:
            raise HTTPException(
                status_code=400,
                detail="Cannot resume: failed to resolve expert configs",
            )

        lead_config = expert_configs[0]
        member_configs = expert_configs[1:] if len(expert_configs) > 1 else []

        # 5. Create team + orchestrator
        team = ExpertTeam(
            pool=app_state.agent_pool,
            template_registry=template_registry,
            redis_client=getattr(app_state, "working_redis_client", None),
        )
        await team.create_team(lead_config=lead_config, member_configs=member_configs)

        try:
            orchestrator = TeamOrchestrator(team=team, checkpoint=checkpoint)
            result = await orchestrator.resume(resolved_plan_id)
        finally:
            try:
                await team.dissolve()
            except Exception:
                pass

    return {
        "task_id": task_id,
        "plan_id": resolved_plan_id,
        "status": result.get("status", "unknown"),
        "result": result.get("result"),
        "phase_results": {
            pid: pr if isinstance(pr, dict) else {"content": str(pr)}
            for pid, pr in (result.get("phase_results") or {}).items()
        },
    }


@router.post("/tasks/stream")
async def stream_task(request: SubmitTaskRequest, req: Request):
    """Submit a task and stream ReAct events via SSE"""
    from sse_starlette.sse import EventSourceResponse

    pool = req.app.state.agent_pool
    skill_registry = req.app.state.skill_registry

    agent = None

    # Same agent resolution logic as submit_task
    if request.agent_name:
        agent = pool.get_agent(request.agent_name)
        if agent is None:
            raise HTTPException(
                status_code=404,
                detail=f"Agent '{request.agent_name}' not found",
            )
    elif request.skill_name:
        try:
            skill_registry.get(request.skill_name)
        except Exception:
            raise HTTPException(
                status_code=404,
                detail=f"Skill '{request.skill_name}' not found",
            )
        agent = pool.get_agent(request.skill_name)
        if agent is None:
            agent = await pool.create_agent_from_skill(request.skill_name)
    else:
        all_skills = skill_registry.list_skills()
        if not all_skills:
            raise HTTPException(
                status_code=400,
                detail="No skills registered and no skill_name or agent_name specified",
            )
        agent = pool.get_agent("default")
        if agent is None:
            # Fallback: create from first available skill
            agent = await pool.create_agent_from_skill(all_skills[0].name)

    async def event_generator():
        import logging
        from agentkit.core.exceptions import LLMProviderError
        from agentkit.core.react import ReActEngine

        stream_logger = logging.getLogger("agentkit.server.stream")

        # Use agent's ReAct config (max_steps, timeout)
        react_config = agent.get_react_config()
        react_engine = ReActEngine(
            llm_gateway=req.app.state.llm_gateway,
            max_steps=react_config["max_steps"],
        )

        # Build messages from input
        messages = [{"role": "user", "content": str(request.input_data)}]

        # Use public accessors instead of private attributes
        tools = agent.get_tools()
        model = agent.get_model()
        system_prompt = agent.get_system_prompt()
        timeout_seconds = react_config["timeout_seconds"]

        chunks_sent = 0
        try:
            async for event in react_engine.execute_stream(
                messages=messages,
                tools=tools,
                model=model,
                agent_name=agent.name,
                system_prompt=system_prompt,
                timeout_seconds=timeout_seconds,
            ):
                chunks_sent += 1
                yield {
                    "event": event.event_type,
                    "data": json.dumps({
                        "step": event.step,
                        "data": event.data,
                        "timestamp": event.timestamp,
                    }),
                }
        except LLMProviderError as e:
            if chunks_sent == 0:
                # No chunks sent yet — try fallback model from gateway
                fallback_model = req.app.state.llm_gateway._get_fallback_model(model)
                if fallback_model:
                    stream_logger.warning(
                        f"LLM provider failed for model '{model}', "
                        f"retrying with fallback '{fallback_model}'"
                    )
                    try:
                        async for event in react_engine.execute_stream(
                            messages=messages,
                            tools=tools,
                            model=fallback_model,
                            agent_name=agent.name,
                            system_prompt=system_prompt,
                            timeout_seconds=timeout_seconds,
                        ):
                            yield {
                                "event": event.event_type,
                                "data": json.dumps({
                                    "step": event.step,
                                    "data": event.data,
                                    "timestamp": event.timestamp,
                                }),
                            }
                    except LLMProviderError as fb_err:
                        stream_logger.error(
                            f"Fallback model '{fallback_model}' also failed: {fb_err}"
                        )
                        yield {
                            "event": "error",
                            "data": json.dumps({
                                "error": str(fb_err),
                                "fallback_attempted": True,
                            }),
                        }
                else:
                    stream_logger.error(f"LLM provider failed, no fallback available: {e}")
                    yield {
                        "event": "error",
                        "data": json.dumps({"error": str(e), "fallback_attempted": False}),
                    }
            else:
                # Chunks already sent — log and terminate gracefully
                stream_logger.error(
                    f"LLM provider failed during streaming (after {chunks_sent} events): {e}"
                )
                yield {
                    "event": "error",
                    "data": json.dumps({
                        "error": str(e),
                        "events_sent": chunks_sent,
                    }),
                }

    return EventSourceResponse(event_generator())


# ---------------------------------------------------------------------------
# Spec endpoints
# ---------------------------------------------------------------------------


def _get_spec_manager(req: Request) -> SpecManager:
    """Get or create SpecManager from app state."""
    if not hasattr(req.app.state, "spec_manager") or req.app.state.spec_manager is None:
        req.app.state.spec_manager = SpecManager()
    return req.app.state.spec_manager


class UpdateSpecRequest(BaseModel):
    """Request body for updating a spec."""

    goal: str | None = None
    steps: list[dict[str, Any]] | None = None
    status: str | None = None
    metadata: dict[str, Any] | None = None


@router.get("/specs")
async def list_specs(status: str | None = None, req: Request = None):
    """List all specs, optionally filtered by status."""
    mgr = _get_spec_manager(req)
    specs = mgr.list_specs(status=status)
    return [asdict(s) for s in specs]


@router.get("/specs/{spec_id}")
async def get_spec(spec_id: str, req: Request):
    """Get a specific spec by ID."""
    mgr = _get_spec_manager(req)
    spec = mgr.get(spec_id)
    if spec is None:
        raise HTTPException(status_code=404, detail=f"Spec '{spec_id}' not found")
    return asdict(spec)


@router.put("/specs/{spec_id}")
async def update_spec(spec_id: str, body: UpdateSpecRequest, req: Request):
    """Update a spec (e.g., edit steps, change status)."""
    mgr = _get_spec_manager(req)
    kwargs: dict[str, Any] = {}
    if body.goal is not None:
        kwargs["goal"] = body.goal
    if body.steps is not None:
        kwargs["steps"] = body.steps
    if body.status is not None:
        kwargs["status"] = body.status
    if body.metadata is not None:
        kwargs["metadata"] = body.metadata
    spec = mgr.update(spec_id, **kwargs)
    if spec is None:
        raise HTTPException(status_code=404, detail=f"Spec '{spec_id}' not found")
    return asdict(spec)


@router.post("/specs/{spec_id}/confirm")
async def confirm_spec(spec_id: str, req: Request):
    """Confirm a spec for execution."""
    mgr = _get_spec_manager(req)
    spec = mgr.confirm(spec_id)
    if spec is None:
        raise HTTPException(status_code=404, detail=f"Spec '{spec_id}' not found")
    return asdict(spec)
