#!/usr/bin/env python

# SPDX-FileCopyrightText: 2025 Weibo, Inc.
#
# SPDX-License-Identifier: Apache-2.0

# -*- coding: utf-8 -*-

import asyncio
import os
from typing import Any, Dict, Optional, Tuple, Union

from executor.agents import Agent
from executor.services.agent_service import AgentService
from shared.logger import setup_logger
from shared.models import EmitterBuilder, ResponsesAPIEmitter, TransportFactory
from shared.models.execution import ExecutionRequest
from shared.models.openai_converter import get_metadata_field
from shared.status import TaskStatus
from shared.telemetry.decorators import (
    add_span_event,
    set_span_attribute,
    trace_async,
    trace_sync,
)

logger = setup_logger("task_processor")


def read_task_data() -> ExecutionRequest:
    """
    Read task data from environment variable or file.

    Returns:
        ExecutionRequest: Task data as ExecutionRequest object

    Raises:
        SystemExit: If TASK_INFO is not found
    """
    from executor.config.env_reader import get_task_info

    task_data = get_task_info()
    if task_data is None:
        logger.error("TASK_INFO not found")
        os._exit(1)
    return ExecutionRequest.from_dict(task_data)


def execute_task(agent: Agent) -> Tuple[TaskStatus, Optional[str]]:
    """
    Execute task
    This function is kept for backward compatibility and is now a wrapper around AgentService.execute_agent_task

    Args:
        agent (Agent): Agent instance

    Returns:
        tuple: (status: TaskStatus, error_message: str or None)
    """
    agent_service = AgentService()
    return agent_service.execute_agent_task(agent)


def _create_emitter(
    request: Union[ExecutionRequest, Dict[str, Any]],
) -> ResponsesAPIEmitter:
    """
    Create an emitter for the task from execution request or task data.

    Supports ExecutionRequest, OpenAI format dict, and legacy dict.

    Args:
        request: ExecutionRequest object or task data dict

    Returns:
        ResponsesAPIEmitter: Configured emitter for the task
    """
    from executor.config import config

    if isinstance(request, ExecutionRequest):
        task_id = request.task_id
        subtask_id = request.subtask_id
    else:
        task_id = get_metadata_field(request, "task_id", -1)
        subtask_id = get_metadata_field(request, "subtask_id", -1)

    return (
        EmitterBuilder()
        .with_task(task_id, subtask_id)
        .with_transport(
            TransportFactory.create_callback(callback_url=config.CALLBACK_URL)
        )
        .with_executor_info(
            name=os.getenv("EXECUTOR_NAME"),
            namespace=os.getenv("EXECUTOR_NAMESPACE"),
        )
        .build()
    )


def _extract_task_attributes(
    request: Union[ExecutionRequest, Dict[str, Any]],
) -> Dict[str, Any]:
    """Extract trace attributes from execution request or task data."""
    if isinstance(request, ExecutionRequest):
        attrs = {
            "task.id": str(request.task_id),
            "task.subtask_id": str(request.subtask_id),
            "task.title": request.task_title or "",
            "task.type": request.type or "online",
            "executor.name": os.getenv("EXECUTOR_NAME", ""),
        }
        # Extract user info if available
        user_data = request.user
        if user_data:
            if user_data.get("id"):
                attrs["user.id"] = str(user_data.get("id"))
            if user_data.get("name"):
                attrs["user.name"] = user_data.get("name")
    else:
        attrs = {
            "task.id": str(get_metadata_field(request, "task_id", -1)),
            "task.subtask_id": str(get_metadata_field(request, "subtask_id", -1)),
            "task.title": get_metadata_field(request, "task_title", ""),
            "task.type": get_metadata_field(request, "type", "online"),
            "executor.name": os.getenv("EXECUTOR_NAME", ""),
        }
        # Extract user info if available
        user_data = get_metadata_field(request, "user", {})
        if user_data:
            if user_data.get("id"):
                attrs["user.id"] = str(user_data.get("id"))
            if user_data.get("name"):
                attrs["user.name"] = user_data.get("name")
    return attrs


def _normalize_request(
    request: Union[ExecutionRequest, Dict[str, Any]],
) -> ExecutionRequest:
    """
    Normalize input to ExecutionRequest.

    Supports:
    - ExecutionRequest objects (returned as-is)
    - OpenAI Responses API format dicts (detected by "model" + "metadata" keys)
    - Legacy ExecutionRequest dict format

    Args:
        request: ExecutionRequest object or task data dict

    Returns:
        ExecutionRequest: Normalized request object
    """
    if isinstance(request, ExecutionRequest):
        return request
    # Detect OpenAI format by checking for "model" + "metadata" keys
    if isinstance(request, dict) and "model" in request and "metadata" in request:
        from shared.models import OpenAIRequestConverter

        return OpenAIRequestConverter.to_execution_request(request)
    return ExecutionRequest.from_dict(request)


def _run_async(coro):
    """Run an async coroutine in sync context."""
    import asyncio
    import concurrent.futures

    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = asyncio.get_event_loop()
        return loop.run_until_complete(coro)

    with concurrent.futures.ThreadPoolExecutor() as executor:
        future = executor.submit(asyncio.run, coro)
        return future.result()


@trace_sync(
    span_name="execute_task",
    tracer_name="executor.tasks",
    extract_attributes=_extract_task_attributes,
)
def process(request: Union[ExecutionRequest, Dict[str, Any]]) -> TaskStatus:
    """
    Process task and send callback with distributed tracing support.

    For subscription tasks, the container will exit after task completion
    (success or failure) since subscription tasks are one-time background
    executions that don't need to keep the container running.

    Args:
        request: ExecutionRequest object or task data dict (for backward compatibility)

    Returns:
        TaskStatus: Processing status
    """
    return _run_async(process_async(request))


@trace_async(
    span_name="execute_task",
    tracer_name="executor.tasks",
    extract_attributes=_extract_task_attributes,
)
async def process_async(request: Union[ExecutionRequest, Dict[str, Any]]) -> TaskStatus:
    """
    Process task and send callback with distributed tracing support (async version).

    This is the async version of process() that should be used when called from
    an async context (e.g., FastAPI lifespan).

    For subscription tasks, the container will exit after task completion
    (success or failure) since subscription tasks are one-time background
    executions that don't need to keep the container running.

    Args:
        request: ExecutionRequest object or task data dict (for backward compatibility)

    Returns:
        TaskStatus: Processing status
    """
    # Normalize to ExecutionRequest for type safety
    task_data = _normalize_request(request)

    # Create emitter for the task lifecycle
    emitter = _create_emitter(task_data)

    # Check if this is a subscription task
    is_subscription = task_data.is_subscription

    # Extract validation_id for validation tasks
    validation_params = task_data.validation_params or {}
    validation_id = validation_params.get("validation_id")

    # Send task started event using emitter
    result = await emitter.start()
    if not result or result.get("status") != TaskStatus.SUCCESS.value:
        logger.error("Failed to send 'start' event")
        set_span_attribute("error", True)
        set_span_attribute("error.message", "Failed to send start event")
        # For subscription tasks, exit container on failure
        if is_subscription:
            logger.info(
                "Subscription task failed to start, exiting container with code 1"
            )
            os._exit(1)
        return TaskStatus.FAILED

    add_span_event("task_started_event_sent")

    # Both subscription and regular tasks are backgrounded so the HTTP response
    # returns RUNNING immediately — well within executor-manager's 5s dispatch timeout.
    # Subscription tasks call os._exit() at the end of the background coroutine
    # once the task completes (success or failure).
    agent_service = AgentService()

    async def _run_in_background():
        try:
            status, error_message = await agent_service.execute_task(task_data)
            message = (
                "Task executed successfully"
                if status in [TaskStatus.SUCCESS, TaskStatus.COMPLETED]
                else error_message
            )
            logger.info(
                f"Background task completed: task_id={task_data.task_id}, status={status}"
            )

            # Only emit a terminal callback here when the agent returned a
            # terminal status directly. ClaudeCode/Agno regular tasks often
            # return RUNNING here and emit their own terminal events later.
            if not is_subscription:
                if status in [TaskStatus.SUCCESS, TaskStatus.COMPLETED]:
                    await emitter.done(content=message or "Task completed")
                elif status == TaskStatus.FAILED:
                    await emitter.error(message or "Task execution failed")
        except Exception as e:
            error_msg = f"Unexpected error during background task execution: {str(e)}"
            logger.exception(error_msg)
            status = TaskStatus.FAILED
            message = error_msg
            await emitter.error(message)
            if is_subscription:
                os._exit(1)
            return

        # For subscription tasks, exit container after completion.
        # NOTE: Do NOT call emitter.done()/emitter.error() here for regular subscription tasks.
        # The agent's execute_async() already sent response.completed with the actual content.
        # Calling emitter.done(content="") again would overwrite the result with empty string.
        # Exception: validation tasks don't send their own done event, so we handle them here.
        if is_subscription:
            if validation_id and status in [TaskStatus.SUCCESS, TaskStatus.COMPLETED]:
                import json

                done_result = {"validation_id": validation_id, "stage": "completed"}
                try:
                    agent = agent_service.get_agent(task_data.task_id)
                    if agent and hasattr(agent, "validation_result"):
                        done_result["value"] = json.dumps(agent.validation_result)
                except Exception as e:
                    logger.warning(f"Failed to retrieve validation result: {e}")
                await emitter.done(content=json.dumps(done_result))

            exit_code = 0 if status in [TaskStatus.SUCCESS, TaskStatus.COMPLETED] else 1
            logger.info(
                f"Subscription task completed with status {status.value}, "
                f"exiting container with code {exit_code}"
            )
            os._exit(exit_code)

    def _handle_background_task_exception(task: asyncio.Task) -> None:
        """
        Global exception handler for background tasks.

        This is the last line of defense to ensure any uncaught exception
        in the background task is reported to the frontend via emitter.
        Even if the internal try-except fails (e.g., emitter.error() throws),
        this callback will catch it and attempt to notify the user.
        """
        try:
            # Check if task raised an exception
            exc = task.exception()
            if exc is not None:
                error_msg = f"[FATAL] Uncaught exception in background task: {type(exc).__name__}: {str(exc)}"
                logger.critical(error_msg)

                # Attempt to send error notification to frontend
                # Use a new event loop since we're in a callback context
                try:
                    # Create a fresh emitter in case the original one is corrupted
                    fallback_emitter = _create_emitter(task_data)

                    # Try to send error via the fallback emitter
                    async def _send_fallback_error():
                        try:
                            await fallback_emitter.error(
                                f"Task execution failed unexpectedly: {str(exc)}",
                                code="fatal_error",
                            )
                        except Exception as send_err:
                            logger.critical(
                                f"[FATAL] Failed to send fallback error notification: {send_err}"
                            )

                    # Run in a new event loop to avoid issues with the current context
                    try:
                        loop = asyncio.get_running_loop()
                        asyncio.create_task(_send_fallback_error())
                    except RuntimeError:
                        # No running loop, create a new one
                        asyncio.run(_send_fallback_error())

                except Exception as fallback_err:
                    logger.critical(
                        f"[FATAL] All error notification attempts failed: {fallback_err}"
                    )

                # For subscription tasks, exit with error code
                if is_subscription:
                    logger.critical(
                        f"Subscription task {task_data.task_id} failed fatally, "
                        f"exiting container with code 1"
                    )
                    os._exit(1)
        except asyncio.CancelledError:
            # Task was cancelled, not an error
            logger.info(
                f"Background task for task_id={task_data.task_id} was cancelled"
            )
        except asyncio.InvalidStateError:
            # Task hasn't completed yet (shouldn't happen in done callback)
            pass
        except Exception as callback_err:
            # Even the exception handler failed - log and exit for subscription tasks
            logger.critical(f"[FATAL] Exception handler itself failed: {callback_err}")
            if is_subscription:
                os._exit(1)

    # Create background task with exception handler callback
    background_task = asyncio.create_task(_run_in_background())
    background_task.add_done_callback(_handle_background_task_exception)
    logger.info(
        f"Task {task_data.task_id} dispatched to background, returning RUNNING immediately"
    )
    return TaskStatus.RUNNING


def run_task() -> TaskStatus:
    """
    Main function, used to read task data and process it

    Returns:
        TaskStatus: Processing status
    """
    execution_request = read_task_data()
    return process(execution_request)


async def run_task_async() -> TaskStatus:
    """
    Main function (async version), used to read task data and process it.

    This is the async version of run_task() that should be used when called from
    an async context (e.g., FastAPI lifespan).

    Returns:
        TaskStatus: Processing status
    """
    execution_request = read_task_data()
    return await process_async(execution_request)
