""" Used to launch the FastAPI web server when worker is running in API mode. """

import os
import threading
import uuid
from dataclasses import dataclass
from typing import Any, Dict, Optional, Union

import requests
import uvicorn
from fastapi import APIRouter, FastAPI
from fastapi.encoders import jsonable_encoder
from fastapi.responses import RedirectResponse

from ...http_client import SyncClientSession
from ...version import __version__ as runpod_version
from .rp_handler import is_generator
from .rp_job import run_job, run_job_generator
from .rp_ping import Heartbeat
from .worker_state import Job, JobsProgress

RUNPOD_ENDPOINT_ID = os.environ.get("RUNPOD_ENDPOINT_ID", None)

TITLE = "RunPod | Development Worker API"

DESCRIPTION = """
The Development Worker API facilitates testing and debugging of your RunPod workers.
It offers a sandbox environment for executing code and simulating interactions with your worker, ensuring your applications can seamlessly transition to production on RunPod serverless platform.
Use this API for comprehensive testing of request submissions and result retrieval, mimicking the behavior of RunPod's operational environment.
---
*Note: This API serves as a local testing tool and will not be utilized once your worker is operational on the RunPod platform.*
"""

# Add CLI tool suggestion if RUNPOD_PROJECT_ID is not set.
if os.environ.get("RUNPOD_PROJECT_ID", None) is None:
    DESCRIPTION += """

    ℹ️ | Consider developing with our CLI tool to streamline your worker development process.

    >_  wget -qO- cli.runpod.net | sudo bash
    >_  runpodctl project create
    """

RUN_DESCRIPTION = """
Initiates processing jobs, returning a unique job ID.

**Parameters:**
- **input** (string): The data to be processed by the worker. This could be a string, JSON object, etc., depending on the worker's requirements.
- **webhook** (string, optional): A callback URL for result notification upon completion. If specified, the server will send a POST request to this URL with the job's result once it's available.

**Returns:**
- **job_id** (string): A unique identifier for the job, used with the `/stream` and `/status` endpoints for monitoring progress and checking job status.
"""

RUNSYNC_DESCRIPTION = """
Executes processing jobs synchronously, returning the job's output directly.

This endpoint is ideal for tasks where immediate result retrieval is necessary,
streamlining the execution process by eliminating the need for subsequent
status or result checks.

**Parameters:**
- **input** (string): The data to be processed by the worker. This should be in a format that the worker can understand (e.g., JSON, text, etc.).
- **webhook** (string, optional): A callback URL to which the result will be posted. While direct result retrieval is the primary operation mode for this endpoint, specifying a webhook allows for asynchronous result notification if needed.

**Returns:**
- **output** (Any): The direct output from the processing job, formatted according to the job's nature and the expected response structure. This could be a JSON object, plain text, or any data structure depending on the processing logic.
"""

STREAM_DESCRIPTION = """
Continuously aggregates the output of a processing job, returning the full output once the job is complete.

This endpoint is especially useful for jobs where the complete output needs to be accessed at once. It provides a consolidated view of the results post-completion, ensuring that users can retrieve the entire output without the need to poll multiple times or manage partial results.

**Parameters:**
- **job_id** (string): The unique identifier of the job for which output is being requested. This ID is used to track the job's progress and aggregate its output.

**Returns:**
- **output** (Any): The aggregated output from the job, returned as a single entity once the job has concluded. The format of the output will depend on the nature of the job and how its results are structured.
"""

STATUS_DESCRIPTION = """
Checks the completion status of a processing job and returns its output if the job is complete.

This endpoint is invaluable for monitoring the progress of a job and obtaining the output only after the job has fully completed. It simplifies the process of querying job completion and retrieving results, eliminating the need for continuous polling or result aggregation.

**Parameters:**
- **job_id** (string): The unique identifier for the job being queried. This ID is used to track and assess the status of the job.

**Returns:**
- **status** (string): The completion status of the job, typically 'complete' or 'in progress'. This status indicates whether the job has finished processing and if the output is ready for retrieval.
- **output** (Any, optional): The final output of the job, provided if the job is complete. The format and structure of the output depend on the job's nature and the data processing involved.

**Note:** The availability of the `output` field is contingent on the job's completion status. If the job is still in progress, this field may be omitted or contain partial results, depending on the implementation.
"""


# ------------------------------ Initializations ----------------------------- #
job_list = JobsProgress()
heartbeat = Heartbeat()


# ------------------------------- Input Objects ------------------------------ #
@dataclass
class Job:
    """Represents a job."""

    id: str
    input: Union[dict, list, str, int, float, bool]


@dataclass
class TestJob:
    """Represents a test job.
    input can be any type of data.
    """

    id: Optional[str] = None
    input: Optional[Union[dict, list, str, int, float, bool]] = None
    webhook: Optional[str] = None


@dataclass
class DefaultRequest:
    """Represents a test input."""

    input: Dict[str, Any]
    webhook: Optional[str] = None


# ------------------------------ Output Objects ------------------------------ #
@dataclass
class JobOutput:
    """Represents the output of a job."""

    id: str
    status: str
    output: Optional[Union[dict, list, str, int, float, bool]] = None
    error: Optional[str] = None


@dataclass
class StreamOutput:
    """Stream representation of a job."""

    id: str
    status: str = "IN_PROGRESS"
    stream: Optional[Union[dict, list, str, int, float, bool]] = None
    error: Optional[str] = None


# ------------------------------ Webhook Sender ------------------------------ #
def _send_webhook(url: str, payload: Dict[str, Any]) -> bool:
    """
    Sends a webhook to the provided URL.

    Args:
        url (str): The URL to send the webhook to.
        payload (Dict[str, Any]): The JSON payload to send.

    Returns:
        bool: True if the request was successful, False otherwise.
    """
    with SyncClientSession() as session:
        try:
            response = session.post(url, json=payload, timeout=10)
            response.raise_for_status()  # Raises exception for 4xx/5xx responses
            return True
        except requests.RequestException as err:
            print(f"WEBHOOK | Request to {url} failed: {err}")
            return False


# ---------------------------------------------------------------------------- #
#                                  API Worker                                  #
# ---------------------------------------------------------------------------- #
class WorkerAPI:
    """Used to launch the FastAPI web server when the worker is running in API mode."""

    def __init__(self, config: Dict[str, Any]):
        """
        Initializes the WorkerAPI class.
        1. Starts the heartbeat thread.
        2. Initializes the FastAPI web server.
        3. Sets the handler for processing jobs.
        """
        # Start the heartbeat thread.
        heartbeat.start_ping()

        self.config = config

        tags_metadata = [
            {
                "name": "Synchronously Submit Request & Get Job Results",
                "description": "Endpoints for submitting job requests and getting the results.",
            },
            {
                "name": "Submit Job Requests",
                "description": "Endpoints for submitting job requests.",
            },
            {
                "name": "Check Job Results",
                "description": "Endpoints for checking the status of a job and getting the results.",
            },
        ]

        # Initialize the FastAPI web server.
        self.rp_app = FastAPI(
            title=TITLE,
            description=DESCRIPTION,
            version=runpod_version,
            docs_url="/",
            openapi_tags=tags_metadata,
        )

        # Create an APIRouter and add the route for processing jobs.
        api_router = APIRouter()

        # Docs Redirect /docs -> /
        api_router.add_api_route(
            "/docs", lambda: RedirectResponse(url="/"), include_in_schema=False
        )

        if RUNPOD_ENDPOINT_ID:
            api_router.add_api_route(
                f"/{RUNPOD_ENDPOINT_ID}/realtime", self._realtime, methods=["POST"]
            )

        # Simulation endpoints.
        api_router.add_api_route(
            "/run",
            self._sim_run,
            methods=["POST"],
            response_model_exclude_none=True,
            summary="Mimics the behavior of the run endpoint.",
            description=RUN_DESCRIPTION,
            tags=["Submit Job Requests"],
        )
        api_router.add_api_route(
            "/runsync",
            self._sim_runsync,
            methods=["POST"],
            response_model_exclude_none=True,
            summary="Mimics the behavior of the runsync endpoint.",
            description=RUNSYNC_DESCRIPTION,
            tags=["Synchronously Submit Request & Get Job Results"],
        )
        api_router.add_api_route(
            "/stream/{job_id}",
            self._sim_stream,
            methods=["POST"],
            response_model_exclude_none=True,
            summary="Mimics the behavior of the stream endpoint.",
            description=STREAM_DESCRIPTION,
            tags=["Check Job Results"],
        )
        api_router.add_api_route(
            "/status/{job_id}",
            self._sim_status,
            methods=["POST"],
            response_model_exclude_none=True,
            summary="Mimics the behavior of the status endpoint.",
            description=STATUS_DESCRIPTION,
            tags=["Check Job Results"],
        )

        # Include the APIRouter in the FastAPI application.
        self.rp_app.include_router(api_router)

    def start_uvicorn(self, api_host="localhost", api_port=8000, api_concurrency=1):
        """
        Starts the Uvicorn server.
        """
        uvicorn.run(
            self.rp_app,
            host=api_host,
            port=int(api_port),
            workers=int(api_concurrency),
            log_level=os.environ.get("UVICORN_LOG_LEVEL", "info"),
            access_log=False,
        )

    # ----------------------------- Realtime Endpoint ---------------------------- #
    async def _realtime(self, job: Job):
        """
        Performs model inference on the input data using the provided handler.
        If handler is not provided, returns an error message.
        """
        job_list.add(job.id)

        # Process the job using the provided handler, passing in the job input.
        job_results = await run_job(self.config["handler"], job.__dict__)

        job_list.remove(job.id)

        # Return the results of the job processing.
        return jsonable_encoder(job_results)

    # ---------------------------------------------------------------------------- #
    #                             Simulation Endpoints                             #
    # ---------------------------------------------------------------------------- #

    # ------------------------------------ run ----------------------------------- #
    async def _sim_run(self, job_request: DefaultRequest) -> JobOutput:
        """Development endpoint to simulate run behavior."""
        assigned_job_id = f"test-{uuid.uuid4()}"
        job_list.add({
            "id": assigned_job_id,
            "input": job_request.input,
            "webhook": job_request.webhook
        })
        return jsonable_encoder({"id": assigned_job_id, "status": "IN_PROGRESS"})

    # ---------------------------------- runsync --------------------------------- #
    async def _sim_runsync(self, job_request: DefaultRequest) -> JobOutput:
        """Development endpoint to simulate runsync behavior."""
        assigned_job_id = f"test-{uuid.uuid4()}"
        job = TestJob(id=assigned_job_id, input=job_request.input)

        if is_generator(self.config["handler"]):
            generator_output = run_job_generator(self.config["handler"], job.__dict__)
            job_output = {"output": []}
            async for stream_output in generator_output:
                job_output["output"].append(stream_output["output"])
        else:
            job_output = await run_job(self.config["handler"], job.__dict__)

        if job_output.get("error", None):
            return jsonable_encoder(
                {"id": job.id, "status": "FAILED", "error": job_output["error"]}
            )

        if job_request.webhook:
            thread = threading.Thread(
                target=_send_webhook,
                args=(job_request.webhook, job_output),
                daemon=True,
            )
            thread.start()

        return jsonable_encoder(
            {"id": job.id, "status": "COMPLETED", "output": job_output["output"]}
        )

    # ---------------------------------- stream ---------------------------------- #
    async def _sim_stream(self, job_id: str) -> StreamOutput:
        """Development endpoint to simulate stream behavior."""
        stashed_job = job_list.get(job_id)
        if stashed_job is None:
            return jsonable_encoder(
                {"id": job_id, "status": "FAILED", "error": "Job ID not found"}
            )

        job = TestJob(id=job_id, input=stashed_job.input)

        if is_generator(self.config["handler"]):
            generator_output = run_job_generator(self.config["handler"], job.__dict__)
            stream_accumulator = []
            async for stream_output in generator_output:
                stream_accumulator.append({"output": stream_output["output"]})
        else:
            return jsonable_encoder(
                {
                    "id": job_id,
                    "status": "FAILED",
                    "error": "Stream not supported, handler must be a generator.",
                }
            )

        job_list.remove(job.id)

        if stashed_job.webhook:
            thread = threading.Thread(
                target=_send_webhook,
                args=(stashed_job.webhook, stream_accumulator),
                daemon=True,
            )
            thread.start()

        return jsonable_encoder(
            {"id": job_id, "status": "COMPLETED", "stream": stream_accumulator}
        )

    # ---------------------------------- status ---------------------------------- #
    async def _sim_status(self, job_id: str) -> JobOutput:
        """Development endpoint to simulate status behavior."""
        stashed_job = job_list.get(job_id)
        if stashed_job is None:
            return jsonable_encoder(
                {"id": job_id, "status": "FAILED", "error": "Job ID not found"}
            )

        job = TestJob(id=stashed_job.id, input=stashed_job.input)

        if is_generator(self.config["handler"]):
            generator_output = run_job_generator(self.config["handler"], job.__dict__)
            job_output = {"output": []}
            async for stream_output in generator_output:
                job_output["output"].append(stream_output["output"])
        else:
            job_output = await run_job(self.config["handler"], job.__dict__)

        job_list.remove(job.id)

        if job_output.get("error", None):
            return jsonable_encoder(
                {"id": job_id, "status": "FAILED", "error": job_output["error"]}
            )

        if stashed_job.webhook:
            thread = threading.Thread(
                target=_send_webhook,
                args=(stashed_job.webhook, job_output),
                daemon=True,
            )
            thread.start()

        return jsonable_encoder(
            {"id": job_id, "status": "COMPLETED", "output": job_output["output"]}
        )
