import asyncio
import functools
import logging
from collections.abc import Awaitable, Callable
from io import StringIO
from typing import Literal, ParamSpec, TypeVar

from pydantic import BaseModel, Field

from hatchet_sdk.clients.events import EventClient
from hatchet_sdk.logger import logger
from hatchet_sdk.runnables.contextvars import (
    ctx_action_key,
    ctx_additional_metadata,
    ctx_step_run_id,
    ctx_task_retry_count,
    ctx_worker_id,
    ctx_workflow_run_id,
)
from hatchet_sdk.utils.typing import (
    STOP_LOOP,
    STOP_LOOP_TYPE,
    JSONSerializableMapping,
    LogLevel,
)

T = TypeVar("T")
P = ParamSpec("P")


class ContextVarToCopyStr(BaseModel):
    name: Literal[
        "ctx_workflow_run_id",
        "ctx_step_run_id",
        "ctx_action_key",
        "ctx_worker_id",
    ]
    value: str | None


class ContextVarToCopyInt(BaseModel):
    name: Literal["ctx_task_retry_count"]
    value: int | None


class ContextVarToCopyDict(BaseModel):
    name: Literal["ctx_additional_metadata"]
    value: JSONSerializableMapping | None


class ContextVarToCopy(BaseModel):
    var: ContextVarToCopyStr | ContextVarToCopyDict | ContextVarToCopyInt = Field(
        discriminator="name"
    )


def copy_context_vars(
    ctx_vars: list[ContextVarToCopy],
    func: Callable[P, T],
    *args: P.args,
    **kwargs: P.kwargs,
) -> T:
    for var in ctx_vars:
        if var.var.name == "ctx_workflow_run_id":
            ctx_workflow_run_id.set(var.var.value)
        elif var.var.name == "ctx_step_run_id":
            ctx_step_run_id.set(var.var.value)
        elif var.var.name == "ctx_task_retry_count":
            ctx_task_retry_count.set(var.var.value)
        elif var.var.name == "ctx_action_key":
            ctx_action_key.set(var.var.value)
        elif var.var.name == "ctx_worker_id":
            ctx_worker_id.set(var.var.value)
        elif var.var.name == "ctx_additional_metadata":
            ctx_additional_metadata.set(var.var.value or {})
        else:
            raise ValueError(f"Unknown context variable name: {var.var.name}")

    return func(*args, **kwargs)


class LogRecord(BaseModel):
    message: str
    step_run_id: str
    level: LogLevel
    task_retry_count: int


class AsyncLogSender:
    def __init__(self, event_client: EventClient):
        self.event_client = event_client
        self.q = asyncio.Queue[LogRecord | STOP_LOOP_TYPE](
            maxsize=event_client.client_config.log_queue_size
        )

    async def consume(self) -> None:
        while True:
            record = await self.q.get()

            if record == STOP_LOOP:
                break

            try:
                await asyncio.to_thread(
                    self.event_client.log,
                    message=record.message,
                    step_run_id=record.step_run_id,
                    level=record.level,
                    task_retry_count=record.task_retry_count,
                )
            except Exception:
                logger.exception("failed to send log to Hatchet")

    def publish(self, record: LogRecord | STOP_LOOP_TYPE) -> None:
        try:
            self.q.put_nowait(record)
        except asyncio.QueueFull:
            logger.warning("log queue is full, dropping log message")


class LogForwardingHandler(logging.StreamHandler):  # type: ignore[type-arg]
    def __init__(self, log_sender: AsyncLogSender, stream: StringIO):
        super().__init__(stream)

        self.log_sender = log_sender

    def emit(self, record: logging.LogRecord) -> None:
        super().emit(record)

        log_entry = self.format(record)
        step_run_id = ctx_step_run_id.get()
        task_retry_count = ctx_task_retry_count.get()

        if not step_run_id:
            return

        self.log_sender.publish(
            LogRecord(
                message=log_entry,
                step_run_id=step_run_id,
                level=LogLevel.from_levelname(record.levelname),
                task_retry_count=task_retry_count or 0,
            )
        )


def capture_logs(
    logger: logging.Logger, log_sender: AsyncLogSender, func: Callable[P, Awaitable[T]]
) -> Callable[P, Awaitable[T]]:
    @functools.wraps(func)
    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        log_stream = StringIO()
        log_forwarder = LogForwardingHandler(log_sender, log_stream)
        log_forwarder.setLevel(logger.level)

        if logger.handlers:
            for handler in logger.handlers:
                if handler.formatter:
                    log_forwarder.setFormatter(handler.formatter)
                    break

            for handler in logger.handlers:
                for filter_obj in handler.filters:
                    log_forwarder.addFilter(filter_obj)

        if not any(h for h in logger.handlers if isinstance(h, LogForwardingHandler)):
            logger.addHandler(log_forwarder)

        try:
            result = await func(*args, **kwargs)
        finally:
            log_forwarder.flush()
            logger.removeHandler(log_forwarder)
            log_stream.close()

        return result

    return wrapper
