"""Agent route handlers - /run, /stream, /batch endpoints."""

from __future__ import annotations

import asyncio
import json
from typing import TYPE_CHECKING, Any

from fastapi import HTTPException
from fastapi.responses import StreamingResponse

from ..schemas import (
    AgentRequest,
    AgentResponse,
    BatchAgentRequest,
    BatchAgentResponse,
    ContentPartPayload,
    StreamUIRequest,
    ToolCallResponse,
    UsageResponse,
)

if TYPE_CHECKING:
    pass


def _request_prompt_payload(
    prompt: str | None, prompt_parts: list[ContentPartPayload] | None
) -> str | list:
    """Convert request payload to runtime format."""
    if prompt_parts:
        return [item.to_runtime_part() for item in prompt_parts]
    return prompt or ""


def _serialize_content_parts(parts: list | None) -> list[ContentPartPayload]:
    """Serialize runtime content parts to response format."""
    from logos_adk.types import normalize_content_parts

    return [ContentPartPayload.from_runtime_part(part) for part in normalize_content_parts(parts)]


async def run_handler(
    agent: Any,
    request: AgentRequest,
    timeout: float = 120,
    artifact_sandbox: Any | None = None,
    observability: Any | None = None,
) -> AgentResponse:
    """Handle /run endpoint - synchronous agent execution."""
    from logos_adk.errors import GuardrailError

    prompt_payload = _request_prompt_payload(request.prompt, request.prompt_parts)
    session_id = request.session_id
    task_run = request.task_run

    run_id = None
    trace_id = None
    request_id = None
    workspace_id = None
    prompt_name = None
    prompt_version = None
    prompt_variant = None
    experiment_id = None
    model = None
    config_fingerprint = None

    content = ""
    content_parts: list = []
    tool_calls: list[ToolCallResponse] = []
    input_tokens = 0
    output_tokens = 0

    try:
        if task_run and getattr(agent, "_task_run_session_store", None):
            session_store = agent._task_run_session_store
            existing = (
                session_store.get(task_run.sessionId) if hasattr(session_store, "get") else None
            )
            if existing:
                session_id = task_run.sessionId

        response = await asyncio.wait_for(
            agent.run(
                prompt_payload,
                session_id=session_id,
                task_run_id=task_run.id if task_run else None,
            ),
            timeout=timeout,
        )

        content = str(getattr(response, "content", ""))
        content_parts = _serialize_content_parts(getattr(response, "content_parts", None))

        for tc in getattr(response, "tool_calls", []) or []:
            tool_calls.append(
                ToolCallResponse(
                    id=getattr(tc, "id", ""),
                    name=getattr(tc, "name", ""),
                    arguments=getattr(tc, "arguments", {}),
                )
            )

        usage = getattr(response, "usage", None)
        if usage:
            input_tokens = int(getattr(usage, "input_tokens", 0))
            output_tokens = int(getattr(usage, "output_tokens", 0))

        run_id = getattr(response, "run_id", None)
        trace_id = getattr(response, "trace_id", None)
        request_id = getattr(response, "request_id", None)
        session_id = session_id or getattr(response, "session_id", None)
        workspace_id = getattr(response, "workspace_id", None)
        prompt_name = getattr(response, "prompt_name", None)
        prompt_version = getattr(response, "prompt_version", None)
        prompt_variant = getattr(response, "prompt_variant", None)
        experiment_id = getattr(response, "experiment_id", None)
        model = getattr(response, "model", None)
        config_fingerprint = getattr(response, "config_fingerprint", None)

    except GuardrailError as e:
        raise HTTPException(
            status_code=422, detail="guardrail", headers={"X-Guardrail-Error": str(e)}
        )
    except asyncio.TimeoutError:
        raise HTTPException(status_code=504, detail="timeout")
    except Exception as e:
        import traceback

        raise HTTPException(status_code=500, detail=f"internal: {str(e)}\n{traceback.format_exc()}")

    return AgentResponse(
        content=content,
        content_parts=content_parts,
        tool_calls=tool_calls,
        usage=UsageResponse(input_tokens=input_tokens, output_tokens=output_tokens),
        run_id=run_id,
        trace_id=trace_id,
        request_id=request_id,
        session_id=session_id,
        workspace_id=workspace_id,
        prompt_name=prompt_name,
        prompt_version=prompt_version,
        prompt_variant=prompt_variant,
        experiment_id=experiment_id,
        model=model,
        config_fingerprint=config_fingerprint,
    )


async def stream_handler(
    agent: Any,
    request: AgentRequest,
    timeout: float = 120,
    artifact_sandbox: Any | None = None,
    observability: Any | None = None,
) -> Any:
    """Handle /stream endpoint - streaming agent execution."""
    from fastapi.responses import StreamingResponse

    prompt_payload = _request_prompt_payload(request.prompt, request.prompt_parts)
    session_id = request.session_id
    task_run = request.task_run

    async def event_generator():
        session_id_local = session_id

        content = ""
        tool_calls: list[dict] = []

        task_run_id = task_run.id if task_run else None
        task_run_status = task_run.status if task_run else None

        try:
            if task_run and getattr(agent, "_task_run_session_store", None):
                session_store = agent._task_run_session_store
                existing = (
                    session_store.get(task_run.sessionId) if hasattr(session_store, "get") else None
                )
                if existing:
                    session_id_local = task_run.sessionId

            has_terminal = False
            async for event in agent.stream(
                prompt_payload,
                session_id=session_id_local,
                task_run_id=task_run_id,
            ):
                event_type = getattr(event, "type", None)

                is_terminal = getattr(event, "terminal", None)
                if is_terminal is True or event_type in {"done", "error"}:
                    has_terminal = True

                event_data = {
                    "type": event_type,
                    "task_run_id": task_run_id,
                    "task_run_status": task_run_status,
                }

                is_terminal = getattr(event, "terminal", None)
                if is_terminal is True or event_type in {"done", "error"}:
                    has_terminal = True

                if event_type == "text":
                    chunk = getattr(event, "text", "") or getattr(event, "delta", "")
                    if chunk:
                        content += chunk
                        yield f"data: {json.dumps({**event_data, 'type': 'text', 'delta': chunk, 'terminal': is_terminal if is_terminal is not None else False})}\n\n"

                elif event_type == "tool_start":
                    tool_call = {
                        "id": getattr(event, "tool_call_id", None) or getattr(event, "call_id", ""),
                        "name": getattr(event, "tool_name", None) or getattr(event, "name", ""),
                        "arguments": getattr(event, "tool_arguments", None)
                        or getattr(event, "arguments", {}),
                    }
                    tool_calls.append(tool_call)
                    yield f"data: {json.dumps({**event_data, 'type': 'tool_start', **tool_call, 'terminal': is_terminal if is_terminal is not None else False})}\n\n"

                elif event_type == "tool_end":
                    result = getattr(event, "result", None)
                    yield f"data: {json.dumps({**event_data, 'type': 'tool_end', 'result': result, 'terminal': is_terminal if is_terminal is not None else False})}\n\n"

                elif event_type == "error":
                    yield f"data: {json.dumps({**event_data, 'error': str(getattr(event, 'error', 'unknown')), 'terminal': True, 'retryable': getattr(event, 'retryable', False), 'task_run_status': 'failed'})}\n\n"

                elif event_type == "done":
                    response = getattr(event, "response", None)
                    if response:
                        session_id_local = session_id_local or getattr(response, "session_id", None)
                        content = str(getattr(response, "content", content))

                    done_event = {
                        **event_data,
                        "type": "done",
                        "terminal": True,
                        "resume_required": False,
                        "retryable": False,
                        "task_run_status": "completed",
                    }
                    yield f"data: {json.dumps(done_event)}\n\n"
                    break

        except asyncio.TimeoutError:
            yield f"data: {json.dumps({**event_data, 'type': 'error', 'detail': 'Agent execution timed out', 'terminal': True})}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({**event_data, 'type': 'error', 'detail': str(e), 'terminal': True})}\n\n"

        if not has_terminal:
            yield f"data: {json.dumps({**event_data, 'type': 'error', 'error': 'incomplete_execution', 'terminal': True, 'retryable': True, 'task_run_status': 'failed'})}\n\n"

    return StreamingResponse(event_generator(), media_type="text/event-stream")


async def stream_ui_handler(
    agent: Any,
    request: StreamUIRequest,
    timeout: float = 120,
    artifact_sandbox: Any | None = None,
    observability: Any | None = None,
) -> Any:
    """Handle /stream-ui endpoint - streaming UI events with pause/resume support."""
    prompt_payload = _request_prompt_payload(request.prompt, request.prompt_parts)
    return _stream_ui_response(
        agent,
        prompt_payload=prompt_payload,
        session_id=request.session_id,
        correlation_id=request.correlation_id,
        task_run_id=request.task_run.id if request.task_run else None,
        task_run_status=request.task_run.status if request.task_run else None,
    )


async def resume_handler(
    agent: Any,
    request: StreamUIRequest,
    timeout: float = 120,
    artifact_sandbox: Any | None = None,
    observability: Any | None = None,
) -> Any:
    """Handle /stream-ui/continue style resume requests."""
    resume_prompt = request.user_input or request.prompt or ""
    correlation_id = request.continue_correlation_id or request.correlation_id
    node_name = request.continue_from_node
    return _resume_ui_response(
        agent,
        prompt=resume_prompt,
        session_id=request.session_id,
        correlation_id=correlation_id,
        node_name=node_name,
        task_run_id=request.task_run.id if request.task_run else request.task_run_id,
        task_run_status=request.task_run.status if request.task_run else None,
    )


def _stream_ui_response(
    agent: Any,
    *,
    prompt_payload: str | list,
    session_id: str | None,
    correlation_id: str | None,
    task_run_id: str | None,
    task_run_status: str | None,
) -> StreamingResponse:
    async def event_source():
        has_terminal = False
        default_event_data = {
            "task_run_id": task_run_id,
            "task_run_status": task_run_status,
        }
        try:
            if hasattr(agent, "stream_ui_events"):
                iterator = agent.stream_ui_events(
                    prompt_payload,
                    session_id=session_id,
                    correlation_id=correlation_id,
                    task_run_id=task_run_id,
                )
            else:
                iterator = agent.stream(
                    prompt_payload,
                    session_id=session_id,
                    correlation_id=correlation_id,
                    task_run_id=task_run_id,
                )

            async for event in iterator:
                payload = _normalize_ui_event(event, default_event_data=default_event_data)
                if payload is None:
                    continue
                if payload.get("terminal") is True:
                    has_terminal = True
                yield f"data: {json.dumps(payload)}\n\n"

        except asyncio.TimeoutError:
            yield f"data: {json.dumps({**default_event_data, 'type': 'error', 'event': 'error', 'error': 'timeout', 'terminal': True, 'retryable': True, 'task_run_status': 'failed'})}\n\n"
            has_terminal = True
        except Exception as exc:
            yield f"data: {json.dumps({**default_event_data, 'type': 'error', 'event': 'error', 'detail': str(exc), 'terminal': True, 'retryable': False, 'task_run_status': 'failed'})}\n\n"
            has_terminal = True

        if not has_terminal:
            yield f"data: {json.dumps({**default_event_data, 'type': 'error', 'event': 'error', 'error': 'incomplete_execution', 'terminal': True, 'retryable': True, 'task_run_status': 'failed'})}\n\n"

    return StreamingResponse(event_source(), media_type="text/event-stream")


def _resume_ui_response(
    agent: Any,
    *,
    prompt: str,
    session_id: str | None,
    correlation_id: str | None,
    node_name: str | None,
    task_run_id: str | None,
    task_run_status: str | None,
) -> StreamingResponse:
    async def event_source():
        default_event_data = {
            "task_run_id": task_run_id,
            "task_run_status": task_run_status,
        }
        if not hasattr(agent, "resume"):
            yield f"data: {json.dumps({**default_event_data, 'type': 'error', 'event': 'error', 'error': 'resume_not_supported', 'terminal': True, 'retryable': False, 'task_run_status': 'failed'})}\n\n"
            return

        has_terminal = False
        try:
            async for event in agent.resume(
                prompt,
                session_id=session_id,
                correlation_id=correlation_id,
                node_name=node_name,
                task_run_id=task_run_id,
            ):
                payload = _normalize_ui_event(event, default_event_data=default_event_data)
                if payload is None:
                    continue
                if payload.get("terminal") is True:
                    has_terminal = True
                yield f"data: {json.dumps(payload)}\n\n"
        except asyncio.TimeoutError:
            yield f"data: {json.dumps({**default_event_data, 'type': 'error', 'event': 'error', 'error': 'timeout', 'terminal': True, 'retryable': True, 'task_run_status': 'failed'})}\n\n"
            has_terminal = True
        except Exception as exc:
            yield f"data: {json.dumps({**default_event_data, 'type': 'error', 'event': 'error', 'detail': str(exc), 'terminal': True, 'retryable': False, 'task_run_status': 'failed'})}\n\n"
            has_terminal = True

        if not has_terminal:
            yield f"data: {json.dumps({**default_event_data, 'type': 'error', 'event': 'error', 'error': 'incomplete_execution', 'terminal': True, 'retryable': True, 'task_run_status': 'failed'})}\n\n"

    return StreamingResponse(event_source(), media_type="text/event-stream")


def _normalize_ui_event(
    event: Any,
    *,
    default_event_data: dict[str, Any],
) -> dict[str, Any] | None:
    if isinstance(event, dict):
        payload = {**default_event_data, **event}
        payload.setdefault("event", payload.get("event_type") or payload.get("type") or "event")
        payload.setdefault("type", payload.get("event_type") or payload.get("event") or "event")

        status = payload.get("status")
        if payload["event"] == "graph_completed":
            payload["terminal"] = True
            if status == "paused":
                payload.setdefault("resume_required", True)
                payload["task_run_status"] = "paused"
            elif status == "completed":
                payload.setdefault("resume_required", False)
                payload.setdefault("retryable", False)
                payload["task_run_status"] = "completed"
            elif status == "failed":
                payload.setdefault("retryable", False)
                payload["task_run_status"] = "failed"
        elif payload["event"] == "pause":
            payload.setdefault("terminal", True)
            payload.setdefault("resume_required", True)
            payload["task_run_status"] = "paused"
        elif payload["event"] == "done":
            payload.setdefault("terminal", True)
            payload.setdefault("resume_required", False)
            payload.setdefault("retryable", False)
            payload["task_run_status"] = "completed"
        elif payload["event"] == "error" or payload["type"] == "error":
            payload.setdefault("terminal", True)
            payload.setdefault("retryable", False)
            payload["task_run_status"] = "failed"
        return payload

    event_type = getattr(event, "type", None)
    is_terminal = getattr(event, "terminal", None)
    event_data = dict(default_event_data)

    if event_type == "text":
        chunk = getattr(event, "text", "") or getattr(event, "delta", "")
        if not chunk:
            return None
        return {
            **event_data,
            "type": "text",
            "event": "text",
            "delta": chunk,
            "terminal": is_terminal if is_terminal is not None else False,
        }

    if event_type == "tool_start":
        return {
            **event_data,
            "type": "tool_start",
            "event": "tool_start",
            "id": getattr(event, "tool_call_id", ""),
            "name": getattr(event, "tool_name", ""),
            "arguments": getattr(event, "tool_arguments", {}),
            "terminal": is_terminal if is_terminal is not None else False,
        }

    if event_type == "tool_end":
        return {
            **event_data,
            "type": "tool_end",
            "event": "tool_end",
            "result": getattr(event, "result", None),
            "terminal": is_terminal if is_terminal is not None else False,
        }

    if event_type == "pause":
        return {
            **event_data,
            "type": "pause",
            "event": "pause",
            "reason": getattr(event, "reason", "approval_required"),
            "terminal": True,
            "resume_required": True,
            "task_run_status": "paused",
        }

    if event_type == "done":
        response = getattr(event, "response", None)
        content = str(getattr(response, "content", "")) if response is not None else ""
        return {
            **event_data,
            "type": "done",
            "event": "done",
            "content": content,
            "terminal": True,
            "resume_required": False,
            "retryable": False,
            "task_run_status": "completed",
        }

    if event_type == "error":
        return {
            **event_data,
            "type": "error",
            "event": "error",
            "error": str(getattr(event, "error", "unknown")),
            "terminal": True,
            "retryable": getattr(event, "retryable", False),
            "task_run_status": "failed",
        }

    return None


async def batch_handler(
    agent: Any,
    request: BatchAgentRequest,
    timeout: float = 120,
    artifact_sandbox: Any | None = None,
    observability: Any | None = None,
) -> BatchAgentResponse:
    """Handle /batch endpoint - parallel agent execution."""
    prompts = request.prompts
    session_ids = request.session_ids or [None] * len(prompts)
    max_concurrency = min(request.max_concurrency, 10)

    semaphore = asyncio.Semaphore(max_concurrency)

    async def run_single(prompt: str, session_id: str | None) -> AgentResponse:
        async with semaphore:
            try:
                response = await asyncio.wait_for(
                    agent.run(prompt, session_id=session_id),
                    timeout=timeout,
                )

                content = str(getattr(response, "content", ""))
                content_parts = _serialize_content_parts(getattr(response, "content_parts", None))

                tool_calls = []
                for tc in getattr(response, "tool_calls", []) or []:
                    tool_calls.append(
                        ToolCallResponse(
                            id=getattr(tc, "id", ""),
                            name=getattr(tc, "name", ""),
                            arguments=getattr(tc, "arguments", {}),
                        )
                    )

                usage = getattr(response, "usage", None)
                input_tokens = int(getattr(usage, "input_tokens", 0)) if usage else 0
                output_tokens = int(getattr(usage, "output_tokens", 0)) if usage else 0

                return AgentResponse(
                    content=content,
                    content_parts=content_parts,
                    tool_calls=tool_calls,
                    usage=UsageResponse(input_tokens=input_tokens, output_tokens=output_tokens),
                    run_id=getattr(response, "run_id", None),
                    trace_id=getattr(response, "trace_id", None),
                    request_id=getattr(response, "request_id", None),
                    session_id=session_id or getattr(response, "session_id", None),
                    workspace_id=getattr(response, "workspace_id", None),
                    prompt_name=getattr(response, "prompt_name", None),
                    prompt_version=getattr(response, "prompt_version", None),
                    prompt_variant=getattr(response, "prompt_variant", None),
                    experiment_id=getattr(response, "experiment_id", None),
                    model=getattr(response, "model", None),
                    config_fingerprint=getattr(response, "config_fingerprint", None),
                )
            except asyncio.TimeoutError:
                raise HTTPException(
                    status_code=504, detail=f"Batch execution timed out for prompt: {prompt[:50]}"
                )
            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Batch execution failed: {str(e)}")

    tasks = [run_single(prompt, session_id) for prompt, session_id in zip(prompts, session_ids)]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    items = []
    for i, result in enumerate(results):
        if isinstance(result, Exception):
            items.append(
                AgentResponse(
                    content=f"Error: {str(result)}",
                    content_parts=[],
                    tool_calls=[],
                    usage=UsageResponse(input_tokens=0, output_tokens=0),
                )
            )
        else:
            items.append(result)

    return BatchAgentResponse(items=items)


__all__ = ["run_handler", "stream_handler", "batch_handler"]
