#!/usr/bin/env python

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

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

"""
FastAPI application for Docker mode executor.

This module contains the FastAPI app and all HTTP endpoints for the executor.
It is only imported in Docker mode to avoid loading unnecessary dependencies
in local mode.
"""

import json
import os
import time
import uuid
from contextlib import asynccontextmanager

from fastapi import BackgroundTasks, FastAPI, HTTPException, Query, Request
from pydantic import BaseModel

from executor.services.agent_service import AgentService
from executor.services.heartbeat_service import start_heartbeat, stop_heartbeat
from executor.tasks import process, process_async, run_task, run_task_async

# Import the shared logger
from shared.logger import setup_logger
from shared.status import TaskStatus
from shared.telemetry.config import get_otel_config
from shared.telemetry.context import set_task_context, set_user_context
from shared.telemetry.context.large_data import log_json_body
from shared.telemetry.core import is_telemetry_enabled

# Use the shared logger setup function
logger = setup_logger("task_executor")

# Flag to track if skills have been initialized
_skills_initialized = False


# Define lifespan context manager for startup and shutdown events
@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Run task at startup if TASK_INFO is available
    """
    # Ensure /home/user directory exists (for Docker container compatibility)
    # Skip in Local mode as it's not needed and may fail on macOS
    from executor.config import config as executor_config

    if executor_config.EXECUTOR_MODE != "local":
        try:
            home_user_dir = "/home/user"
            if not os.path.exists(home_user_dir):
                os.makedirs(home_user_dir, mode=0o777, exist_ok=True)
                logger.info(f"Created {home_user_dir} directory")
        except Exception as e:
            logger.warning(f"Failed to create {home_user_dir} directory: {e}")

    # Initialize OpenTelemetry if enabled (configuration from shared/telemetry/config.py)
    otel_config = get_otel_config("wegent-executor")
    if otel_config.enabled:
        try:
            from shared.telemetry.context import restore_trace_context_from_env
            from shared.telemetry.core import init_telemetry

            init_telemetry(
                service_name=otel_config.service_name,
                enabled=otel_config.enabled,
                otlp_endpoint=otel_config.otlp_endpoint,
                sampler_ratio=otel_config.sampler_ratio,
                service_version="1.0.0",
                metrics_enabled=otel_config.metrics_enabled,
                capture_request_headers=otel_config.capture_request_headers,
                capture_request_body=otel_config.capture_request_body,
                capture_response_headers=otel_config.capture_response_headers,
                capture_response_body=otel_config.capture_response_body,
                max_body_size=otel_config.max_body_size,
            )
            logger.info("OpenTelemetry initialized successfully")

            # Apply instrumentation
            from shared.telemetry.instrumentation import (
                setup_opentelemetry_instrumentation,
            )

            setup_opentelemetry_instrumentation(app, logger)

            # Restore parent trace context from environment variables
            # This continues the trace started by executor_manager
            restore_trace_context_from_env()
            logger.debug("Restored trace context from environment")
        except Exception as e:
            logger.warning(f"Failed to initialize OpenTelemetry: {e}")

    try:
        if os.getenv("TASK_INFO"):
            # Generate a request_id for startup task execution
            # This ensures logs have a request_id even without HTTP request
            startup_request_id = str(uuid.uuid4())[:8]
            from shared.telemetry.context import set_request_context

            set_request_context(startup_request_id)

            logger.info("TASK_INFO environment variable found, attempting to run task")
            # Use async version to avoid "event loop already running" error
            status = await run_task_async()
            logger.info(f"Task execution status: {status}")
        else:
            logger.info(
                "No TASK_INFO environment variable found, skipping task execution"
            )
    except Exception as e:
        logger.exception(f"Error running task at startup: {str(e)}")

    # Start heartbeat service for sandbox health monitoring
    try:
        if start_heartbeat():
            logger.info("Heartbeat service started successfully")
    except Exception as e:
        logger.warning(f"Failed to start heartbeat service: {e}")

    # Initialize Claude Code for sandbox mode if AUTH_TOKEN is provided
    # This allows sandbox containers to download skills at startup
    try:
        # Use env_reader to support both env vars and file-based config
        from executor.config.env_reader import (
            get_auth_token,
            get_task_id,
            get_task_info,
        )

        auth_token = get_auth_token()
        task_id = get_task_id()

        # If AUTH_TOKEN env var is not set, try to extract from TASK_INFO
        # This handles non-sandbox tasks where auth_token is in the task data
        if not auth_token:
            task_info = get_task_info()
            if task_info:
                # Support both OpenAI format (metadata.auth_token) and legacy format (auth_token)
                metadata = task_info.get("metadata", {})
                auth_token = metadata.get("auth_token") or task_info.get("auth_token")
                if not task_id:
                    task_id = str(
                        metadata.get("task_id") or task_info.get("task_id", "")
                    )

        if auth_token and task_id:
            logger.info(
                f"AUTH_TOKEN found, initializing Claude Code for task {task_id}..."
            )
            _initialize_sandbox_claude(auth_token, task_id)

            # Pre-download all attachments for the task
            logger.info(f"Pre-downloading attachments for task {task_id}...")
            _predownload_task_attachments(auth_token, task_id)
    except Exception as e:
        logger.warning(f"Failed to initialize Claude Code for task: {e}")

    yield  # Application runs here

    # Stop heartbeat service
    try:
        stop_heartbeat()
        logger.info("Heartbeat service stopped")
    except Exception as e:
        logger.warning(f"Error stopping heartbeat service: {e}")

    # Shutdown OpenTelemetry
    if otel_config.enabled:
        from shared.telemetry.core import shutdown_telemetry

        shutdown_telemetry()
        logger.info("OpenTelemetry shutdown completed")


def _initialize_sandbox_claude(auth_token: str, task_id: str) -> None:
    """Initialize Claude Code for sandbox mode.

    Downloads and deploys skills based on task_id via Backend API.

    Args:
        auth_token: JWT token for API authentication
        task_id: Task ID for fetching associated skills
    """
    from executor.services.api_client import SkillDownloader, fetch_task_skills

    # Fetch skills via API
    logger.info(f"[SandboxInit] Fetching skills for task {task_id}...")
    skills_info = fetch_task_skills(task_id, auth_token)
    logger.info(
        f"[SandboxInit] Skills info: skills={skills_info.skills}, "
        f"preload_skills={skills_info.preload_skills}, "
        f"team_namespace={skills_info.team_namespace}"
    )

    # Combine skills and preload_skills
    all_skills = list(set(skills_info.skills + skills_info.preload_skills))
    if not all_skills:
        logger.info("[SandboxInit] No skills to deploy, skipping")
        return

    # Download and deploy skills
    logger.info(f"[SandboxInit] Deploying {len(all_skills)} skills: {all_skills}")
    downloader = SkillDownloader(
        auth_token=auth_token,
        team_namespace=skills_info.team_namespace,
        task_id=(
            int(task_id) if task_id else None
        ),  # Enable task-based authorization for shared teams
    )
    resolved_skill_map = {
        **(skills_info.skill_refs or {}),
        **(skills_info.preload_skill_refs or {}),
    }
    logger.info(
        "[SandboxInit] Resolved skill map prepared: keys=%s",
        sorted(resolved_skill_map.keys()),
    )
    result = downloader.download_and_deploy(
        all_skills,
        resolved_skill_map=resolved_skill_map or None,
    )
    logger.info(
        f"[SandboxInit] Skills deployment complete: "
        f"{result.success_count}/{result.total_count} deployed to {result.skills_dir}"
    )


def _predownload_task_attachments(auth_token: str, task_id: str) -> None:
    """Pre-download all attachments for a task at sandbox startup.

    Downloads all attachments for the task to /home/user directory,
    organized by subtask to match the sandbox path format:
    /home/user/{task_id}:executor:attachments/{subtask_id}/{filename}

    Args:
        auth_token: JWT token for API authentication
        task_id: Task ID for fetching attachments
    """
    from executor.services.api_client import fetch_task_attachments
    from executor.services.attachment_downloader import AttachmentDownloader

    # Define base attachments directory
    attachments_base_dir = "/home/user"

    # Fetch all attachments for the task
    logger.info(f"[SandboxInit] Fetching attachments for task {task_id}...")
    result = fetch_task_attachments(task_id, auth_token)

    if result.error:
        logger.warning(f"[SandboxInit] Failed to fetch attachments: {result.error}")
        return

    if not result.attachments:
        logger.info("[SandboxInit] No attachments to download")
        return

    # Group attachments by subtask_id for batch downloading
    from collections import defaultdict

    attachments_by_subtask = defaultdict(list)
    for att in result.attachments:
        attachments_by_subtask[att.subtask_id].append(
            {
                "id": att.id,
                "original_filename": att.original_filename,
                "subtask_id": att.subtask_id,
            }
        )

    # Download attachments per subtask using AttachmentDownloader
    total_success = 0
    total_failed = 0

    for subtask_id, attachments in attachments_by_subtask.items():
        logger.info(
            f"[SandboxInit] Downloading {len(attachments)} attachments for subtask {subtask_id}"
        )

        downloader = AttachmentDownloader(
            workspace=attachments_base_dir,
            task_id=task_id,
            subtask_id=str(subtask_id),
            auth_token=auth_token,
        )

        download_result = downloader.download_all(attachments)
        total_success += len(download_result.success)
        total_failed += len(download_result.failed)

    logger.info(
        f"[SandboxInit] Attachment pre-download complete: "
        f"{total_success}/{total_success + total_failed} downloaded to {attachments_base_dir}"
    )


async def _maybe_initialize_skills_from_env() -> None:
    """Initialize skills on first HTTP request (for preload sandbox mode).

    This function is called by HTTP middleware when the first request arrives.
    It reads auth_token and task_id from environment variables, Downward API files,
    or TASK_INFO, then downloads and deploys skills.

    Only applies to sandbox mode. For non-sandbox tasks, skills are downloaded
    automatically during agent initialization.
    """
    global _skills_initialized

    if _skills_initialized:
        return

    # Check if sandbox preload mode is enabled via environment variable
    import os

    preload_mode = os.getenv("PRELOAD_MODE", "").lower() == "true"
    if not preload_mode:
        _skills_initialized = True  # Preload mode not enabled, skip
        return

    # Read auth_token and task_id from env vars, Downward API files, or TASK_INFO
    from executor.config.env_reader import get_auth_token, get_task_id, get_task_info

    auth_token = get_auth_token()
    task_id = get_task_id()

    # If AUTH_TOKEN env var is not set, try to extract from TASK_INFO
    if not auth_token:
        task_info = get_task_info()
        if task_info:
            # Support both OpenAI format (metadata.auth_token) and legacy format (auth_token)
            metadata = task_info.get("metadata", {})
            auth_token = metadata.get("auth_token") or task_info.get("auth_token")
            if not task_id:
                task_id = str(metadata.get("task_id") or task_info.get("task_id", ""))

    if not auth_token or not task_id:
        logger.debug(
            "[PreloadMode] No auth_token or task_id available yet, skipping skills init"
        )
        return

    try:
        logger.info(
            f"[PreloadMode] First request received, initializing skills for task {task_id}..."
        )
        _initialize_sandbox_claude(auth_token, str(task_id))
        logger.info("[PreloadMode] Skills initialization completed")

        # Also pre-download attachments for the task
        logger.info(f"[PreloadMode] Pre-downloading attachments for task {task_id}...")
        _predownload_task_attachments(auth_token, str(task_id))
        logger.info("[PreloadMode] Attachments pre-download completed")

        _skills_initialized = True
    except Exception as e:
        logger.warning(f"[PreloadMode] Failed to initialize skills or attachments: {e}")
        # Don't fail the request, just log the warning
        _skills_initialized = True  # Mark as initialized to avoid retrying


def create_app() -> FastAPI:
    """Create and configure the FastAPI application."""
    return FastAPI(
        title="Task Executor API",
        description="API for executing tasks with agents",
        lifespan=lifespan,
    )


# Create FastAPI app
app = create_app()

# Register envd Connect RPC routes (enabled by default)
envd_enabled = os.getenv("ENVD_ENABLED", "true").lower() == "true"
if envd_enabled:
    try:
        from executor.envd.server import register_envd_routes

        register_envd_routes(app)
        logger.info("envd Connect RPC routes registered to main app")
    except Exception as e:
        logger.warning(f"Failed to register envd routes: {e}")


# Add middleware for request/response logging and OTEL capture
@app.middleware("http")
async def log_requests(request: Request, call_next):
    from opentelemetry import context
    from starlette.responses import StreamingResponse

    await _maybe_initialize_skills_from_env()

    # Skip logging for health check requests
    if request.url.path == "/":
        return await call_next(request)

    # Get request_id from header (propagated from upstream service) or generate new one
    request_id = request.headers.get("X-Request-ID") or str(uuid.uuid4())[:8]
    request.state.request_id = request_id

    start_time = time.time()

    # Always set request context for logging (works even without OTEL)
    from shared.telemetry.context import (
        extract_trace_context_from_headers,
        set_request_context,
    )

    set_request_context(request_id)

    # Extract and attach trace context from incoming HTTP headers for distributed tracing
    # This allows executor to continue the trace started by executor_manager
    otel_cfg = get_otel_config()
    if otel_cfg.enabled:
        try:
            # Convert headers to dict for extraction
            headers_dict = dict(request.headers)
            extracted_ctx = extract_trace_context_from_headers(headers_dict)
            if extracted_ctx:
                # Attach the extracted context so subsequent spans become children
                context.attach(extracted_ctx)
                logger.debug(
                    f"Attached trace context from headers for request {request_id}"
                )
        except Exception as e:
            logger.debug(f"Failed to extract trace context from headers: {e}")

    # Capture request body if OTEL is enabled and body capture is configured
    request_body = None
    if (
        otel_cfg.enabled
        and otel_cfg.capture_request_body
        and request.method in ("POST", "PUT", "PATCH")
    ):
        try:
            body_bytes = await request.body()
            if body_bytes:
                max_body_size = otel_cfg.max_body_size
                if len(body_bytes) <= max_body_size:
                    request_body = body_bytes.decode("utf-8", errors="replace")
                else:
                    request_body = (
                        body_bytes[:max_body_size].decode("utf-8", errors="replace")
                        + f"... [truncated, total size: {len(body_bytes)} bytes]"
                    )
        except Exception as e:
            logger.debug(f"Failed to capture request body: {e}")

    # Add OpenTelemetry span attributes if enabled
    if otel_cfg.enabled:
        try:
            from opentelemetry import trace

            from shared.telemetry.core import is_telemetry_enabled

            if is_telemetry_enabled():
                if request_body:
                    current_span = trace.get_current_span()
                    if current_span and current_span.is_recording():
                        log_json_body("http.request.body", request_body)
        except Exception as e:
            logger.debug(f"Failed to set OTEL context: {e}")

    # Pre-request logging
    logger.info(
        f"request : {request.method} {request.url.path} {request.query_params} {request_id}"
    )

    # Process request
    response = await call_next(request)
    process_time = (time.time() - start_time) * 1000

    # Capture response headers and body if OTEL is enabled
    if otel_cfg.enabled:
        try:
            from opentelemetry import trace

            from shared.telemetry.core import is_telemetry_enabled

            if is_telemetry_enabled():
                current_span = trace.get_current_span()
                if current_span and current_span.is_recording():
                    # Capture response headers
                    if otel_cfg.capture_response_headers:
                        for header_name, header_value in response.headers.items():
                            if header_name.lower() in (
                                "authorization",
                                "cookie",
                                "set-cookie",
                            ):
                                header_value = "[REDACTED]"
                            current_span.set_attribute(
                                f"http.response.header.{header_name}", header_value
                            )

                    # Capture response body (only for non-streaming responses)
                    if otel_cfg.capture_response_body:
                        if not isinstance(response, StreamingResponse):
                            try:
                                response_body_chunks = []
                                async for chunk in response.body_iterator:
                                    response_body_chunks.append(chunk)

                                response_body = b"".join(response_body_chunks)

                                max_body_size = otel_cfg.max_body_size
                                if response_body:
                                    if len(response_body) <= max_body_size:
                                        body_str = response_body.decode(
                                            "utf-8", errors="replace"
                                        )
                                    else:
                                        body_str = (
                                            response_body[:max_body_size].decode(
                                                "utf-8", errors="replace"
                                            )
                                            + f"... [truncated, total size: {len(response_body)} bytes]"
                                        )
                                    current_span.set_attribute(
                                        "http.response.body", body_str
                                    )

                                from starlette.responses import Response

                                response = Response(
                                    content=response_body,
                                    status_code=response.status_code,
                                    headers=dict(response.headers),
                                    media_type=response.media_type,
                                )
                            except Exception as e:
                                logger.debug(f"Failed to capture response body: {e}")
        except Exception as e:
            logger.debug(f"Failed to capture OTEL response: {e}")

    # Post-request logging
    logger.info(
        f"response: {request.method} {request.url.path} {request_id} {response.status_code} {process_time:.2f}ms"
    )

    response.headers["X-Request-ID"] = request_id
    return response


agent_service = AgentService()


@app.get("/")
async def health_check():
    """Health check endpoint for container readiness probes."""
    return {"status": "healthy", "service": "task_executor"}


class OpenAIBackgroundResponse(BaseModel):
    """Response model for OpenAI background mode"""

    id: str
    status: str
    message: str = ""


@app.post("/v1/responses", response_model=OpenAIBackgroundResponse)
async def openai_responses(request: Request):
    """
    OpenAI Responses API endpoint for background mode execution.

    This endpoint accepts OpenAI Responses API format requests with background=true.
    It converts the request to ExecutionRequest format and processes the task.

    Request format:
    {
        "model": "...",
        "input": "..." or [...],
        "instructions": "...",
        "stream": false,
        "background": true,
        "metadata": {
            "task_id": 123,
            "subtask_id": 456,
            ...
        },
        "model_config": {...}
    }

    Returns:
        OpenAI background response with queued status
    """
    # Read raw JSON data from request body
    body_bytes = await request.body()
    openai_request = json.loads(body_bytes)

    # Extract task identification from metadata
    metadata = openai_request.get("metadata", {})
    task_id = metadata.get("task_id", -1)
    subtask_id = metadata.get("subtask_id", -1)
    background = openai_request.get("background", False)

    logger.info(
        f"[v1/responses] Received OpenAI request: task_id={task_id}, "
        f"subtask_id={subtask_id}, background={background}, request={openai_request}"
    )

    # Set task and user context for tracing
    otel_config = get_otel_config()
    if otel_config.enabled and is_telemetry_enabled():
        set_task_context(task_id=task_id, subtask_id=subtask_id)
        user_id = metadata.get("user_id")
        user_name = metadata.get("user_name")
        if user_id or user_name:
            set_user_context(
                user_id=str(user_id) if user_id else None,
                user_name=user_name,
            )

    try:
        # Convert OpenAI request to ExecutionRequest format
        from shared.models import OpenAIRequestConverter

        execution_request = OpenAIRequestConverter.to_execution_request(openai_request)
        task_data = execution_request.to_dict()

        # Use async process function to handle task uniformly
        # IMPORTANT: Must use await process_async() instead of sync process()
        # because this endpoint is async and runs in FastAPI's event loop.
        # Using sync process() would cause _run_async() to create a new thread
        # with a temporary event loop, and any background tasks created in
        # execute() would be abandoned when that temporary loop closes.
        status = await process_async(task_data)

        logger.info(f"[v1/responses] Task {task_id} submitted, status={status.value}")

        # Return OpenAI background response format
        return OpenAIBackgroundResponse(
            id=f"resp_{subtask_id}",
            status="queued" if background else status.value,
            message=f"Task execution status: {status.value}",
        )

    except Exception as e:
        logger.exception(f"[v1/responses] Error executing task {task_id}: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error executing task: {str(e)}")


@app.delete("/api/tasks/session")
async def delete_session(
    task_id: int = Query(..., description="Task ID to delete session for")
):
    """
    Delete an agent session for a specific task_id
    """
    status, message = await agent_service.delete_session_async(task_id)

    if status == TaskStatus.SUCCESS:
        return {"message": message}
    else:
        raise HTTPException(status_code=404, detail=message)


@app.post("/api/tasks/cancel")
async def cancel_task(
    task_id: int = Query(..., description="Task ID to cancel"),
    background_tasks: BackgroundTasks = None,
):
    """
    Cancel the currently running task for a specific task_id
    Returns immediately, callback is sent asynchronously in background to avoid blocking executor_manager's cancel request
    """
    # Set task context for tracing
    otel_config = get_otel_config()
    if otel_config.enabled and is_telemetry_enabled():
        set_task_context(task_id=task_id)

    status, message = agent_service.cancel_task(task_id)

    if status == TaskStatus.SUCCESS:
        # Send cancel callback in background without blocking response
        if background_tasks:
            background_tasks.add_task(agent_service.send_cancel_callback_async, task_id)
        return {"message": message}
    else:
        raise HTTPException(status_code=400, detail=message)


@app.get("/api/tasks/sessions")
async def list_sessions():
    """
    List all active agent sessions
    """
    sessions = agent_service.list_sessions()
    return {"total": len(sessions), "sessions": sessions}


@app.delete("/api/tasks/claude/sessions")
async def close_all_claude_sessions():
    """
    Close all Claude client connections
    """
    try:
        status, message = await agent_service.close_all_claude_sessions()
        if status == TaskStatus.SUCCESS:
            return {"message": message}
        else:
            raise HTTPException(status_code=500, detail=message)
    except Exception as e:
        logger.exception(f"Error closing all Claude client connections: {str(e)}")
        raise HTTPException(
            status_code=500, detail=f"Error closing connections: {str(e)}"
        )


@app.delete("/api/tasks/sessions/close")
async def close_all_agent_sessions():
    """
    Close all agent connections regardless of type
    If an agent type doesn't support connection closing, it will be skipped
    """
    try:
        status, message, error_detail = await agent_service.close_all_agent_sessions()
        if status == TaskStatus.SUCCESS:
            return {"message": message}
        else:
            # Return 200 status code even with errors, as some agents may have closed successfully
            return {
                "message": message,
                "partial_success": True,
                "error_detail": error_detail,
            }
    except Exception as e:
        logger.exception(f"Error closing agent connections: {str(e)}")
        raise HTTPException(
            status_code=500, detail=f"Error closing connections: {str(e)}"
        )
