"""Task submission routes"""

import json
import uuid
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

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
    intent_router = req.app.state.intent_router
    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 Intent Router to find matching skill
    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",
            )
        try:
            routing_result = await intent_router.route(request.input_data, all_skills)
            skill = skill_registry.get(routing_result.matched_skill)
            # Get or create agent for this skill
            agent = pool.get_agent(routing_result.matched_skill)
            if agent is None:
                agent = await pool.create_agent_from_skill(routing_result.matched_skill)
        except (ValueError, RuntimeError) as e:
            raise HTTPException(status_code=400, detail=str(e))

    # 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:
            quality_result = await quality_gate.validate(task_result.output_data or {}, skill)
        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/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
    intent_router = req.app.state.intent_router

    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",
            )
        try:
            routing_result = await intent_router.route(request.input_data, all_skills)
            skill_registry.get(routing_result.matched_skill)
            agent = pool.get_agent(routing_result.matched_skill)
            if agent is None:
                agent = await pool.create_agent_from_skill(routing_result.matched_skill)
        except (ValueError, RuntimeError) as e:
            raise HTTPException(status_code=400, detail=str(e))

    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())
