"""
Job related helpers.
"""

import inspect
import json
import os
import traceback
from typing import Any, AsyncGenerator, Callable, Dict, Optional, Union, List

import aiohttp

from runpod.http_client import ClientSession, TooManyRequests
from runpod.serverless.modules.rp_logger import RunPodLogger

from ...version import __version__ as runpod_version
from ..utils import rp_debugger
from .rp_handler import is_generator
from .rp_http import send_result, stream_result
from .rp_tips import check_return_size
from .worker_state import WORKER_ID, REF_COUNT_ZERO, JobsProgress

JOB_GET_URL = str(os.environ.get("RUNPOD_WEBHOOK_GET_JOB")).replace("$ID", WORKER_ID)

log = RunPodLogger()
job_progress = JobsProgress()


def _job_get_url(batch_size: int = 1):
    """
    Prepare the URL for making a 'get' request to the serverless API (sls).

    This function constructs the appropriate URL for sending a 'get' request to the serverless API,
    ensuring that the request will be correctly routed and processed by the API.

    Returns:
        str: The prepared URL for the 'get' request to the serverless API.
    """

    if batch_size > 1:
        job_take_url = JOB_GET_URL.replace("/job-take/", "/job-take-batch/")
        job_take_url += f"&batch_size={batch_size}"
    else:
        job_take_url = JOB_GET_URL

    job_in_progress = "1" if job_progress.get_job_list() else "0"
    job_take_url += f"&job_in_progress={job_in_progress}"

    log.debug(f"rp_job | get_job: {job_take_url}")
    return job_take_url


async def get_job(
    session: ClientSession, num_jobs: int = 1
) -> Optional[List[Dict[str, Any]]]:
    """
    Get a job from the job-take API.

    `num_jobs = 1` will query the legacy singular job-take API.

    `num_jobs > 1` will query the batch job-take API.

    Args:
        session (ClientSession): The aiohttp ClientSession to use for the request.
        num_jobs (int): The number of jobs to get.
    """
    async with session.get(_job_get_url(num_jobs)) as response:
        log.debug(f"rp_job | Response: {type(response).__name__} {response.status}")

        if response.status == 204:
            log.debug("rp_job | Received 204 status, no jobs.")
            return

        if response.status == 400:
            log.debug("rp_job | Received 400 status, expected when FlashBoot is enabled.")
            return

        if response.status == 429:
            raise TooManyRequests(
                response.request_info,
                response.history,
                status=response.status,
                message=response.reason
            )

        # All other errors should raise an exception
        response.raise_for_status()

        # Verify if the content type is JSON
        if response.content_type != "application/json":
            log.debug(f"rp_job | Unexpected content type: {response.content_type}")
            return

        # Check if there is a non-empty content to parse
        if response.content_length == 0:
            log.debug("rp_job | No content to parse.")
            return

        try:
            jobs = await response.json()
            log.debug("rp_job | Received Job(s)")
        except aiohttp.ContentTypeError:
            log.debug(f"rp_job | Response content is not valid JSON. {response.content}")
            return
        except ValueError as json_error:
            log.debug(f"rp_job | Failed to parse JSON response: {json_error}")
            return

        # legacy job-take API
        if isinstance(jobs, dict):
            if "id" not in jobs or "input" not in jobs:
                raise Exception("Job has missing field(s): id or input.")
            return [jobs]

        # batch job-take API
        if isinstance(jobs, list):
            return jobs


async def handle_job(session: ClientSession, config: Dict[str, Any], job) -> dict:
    if is_generator(config["handler"]):
        is_stream = True
        generator_output = run_job_generator(config["handler"], job)
        log.debug("Handler is a generator, streaming results.", job["id"])

        job_result = {"output": []}
        async for stream_output in generator_output:
            log.debug(f"Stream output: {stream_output}", job["id"])

            if type(stream_output.get("output")) == dict:
                if stream_output["output"].get("error"):
                    stream_output = {"error": str(stream_output["output"]["error"])}

            if stream_output.get("error"):
                job_result = stream_output
                break

            if config.get("return_aggregate_stream", False):
                job_result["output"].append(stream_output["output"])

            await stream_result(session, stream_output, job)
    else:
        is_stream = False
        job_result = await run_job(config["handler"], job)

    # If refresh_worker is set, pod will be reset after job is complete.
    if config.get("refresh_worker", False):
        log.info("refresh_worker flag set, stopping pod after job.", job["id"])
        job_result["stopPod"] = True

    # If rp_debugger is set, debugger output will be returned.
    if config["rp_args"].get("rp_debugger", False) and isinstance(job_result, dict):
        job_result["output"]["rp_debugger"] = rp_debugger.get_debugger_output()
        log.debug("rp_debugger | Flag set, returning debugger output.", job["id"])

        # Calculate ready delay for the debugger output.
        ready_delay = (config["reference_counter_start"] - REF_COUNT_ZERO) * 1000
        job_result["output"]["rp_debugger"]["ready_delay_ms"] = ready_delay
    else:
        log.debug("rp_debugger | Flag not set, skipping debugger output.", job["id"])
        rp_debugger.clear_debugger_output()

    # Send the job result back to JOB_DONE_URL
    await send_result(session, job_result, job, is_stream=is_stream)


async def run_job(handler: Callable, job: Dict[str, Any]) -> Dict[str, Any]:
    """
    Run the job using the handler.

    Args:
        handler (Callable): The handler function to use.
        job (Dict[str, Any]): The job to run.

    Returns:
        Dict[str, Any]: The result of running the job.
    """
    log.info("Started.", job["id"])
    run_result = {}

    try:
        handler_return = handler(job)
        job_output = (
            await handler_return
            if inspect.isawaitable(handler_return)
            else handler_return
        )

        log.debug(f"Handler output: {job_output}", job["id"])

        if isinstance(job_output, dict):
            error_msg = job_output.pop("error", None)
            refresh_worker = job_output.pop("refresh_worker", None)
            run_result["output"] = job_output

            if error_msg:
                run_result["error"] = error_msg
            if refresh_worker:
                run_result["stopPod"] = True

        elif isinstance(job_output, bool):
            run_result = {"output": job_output}

        else:
            run_result = {"output": job_output}

        if run_result.get("output") == {}:
            run_result.pop("output")

        check_return_size(run_result)  # Checks the size of the return body.

    except Exception as err:
        error_info = {
            "error_type": str(type(err)),
            "error_message": str(err),
            "error_traceback": traceback.format_exc(),
            "hostname": os.environ.get("RUNPOD_POD_HOSTNAME", "unknown"),
            "worker_id": os.environ.get("RUNPOD_POD_ID", "unknown"),
            "runpod_version": runpod_version,
        }

        log.error("Captured Handler Exception", job["id"])
        log.error(json.dumps(error_info, indent=4))
        run_result = {"error": json.dumps(error_info)}

    finally:
        log.debug(f"run_job return: {run_result}", job["id"])

    return run_result


async def run_job_generator(
    handler: Callable, job: Dict[str, Any]
) -> AsyncGenerator[Dict[str, Union[str, Any]], None]:
    """
    Run generator job used to stream output.
    Yields output partials from the generator.
    """
    is_async_gen = inspect.isasyncgenfunction(handler)
    log.debug(
        "Using Async Generator" if is_async_gen else "Using Standard Generator",
        job["id"],
    )

    try:
        job_output = handler(job)

        if is_async_gen:
            async for output_partial in job_output:
                log.debug(f"Async Generator output: {output_partial}", job["id"])
                yield {"output": output_partial}
        else:
            for output_partial in job_output:
                log.debug(f"Generator output: {output_partial}", job["id"])
                yield {"output": output_partial}

    except Exception as err:
        log.error(err, job["id"])
        yield {"error": f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"}
    finally:
        log.info("Finished running generator.", job["id"])
