import contextvars
import uuid
from contextlib import asynccontextmanager
from typing import Any, Optional

from r2r.base import User

from .kv_logger import KVLoggingSingleton

run_id_var = contextvars.ContextVar("run_id", default=None)


class RunManager:
    def __init__(self, logger: KVLoggingSingleton):
        self.logger = logger
        self.run_info = {}

    def generate_run_id(self) -> uuid.UUID:
        return uuid.uuid4()

    async def set_run_info(
        self, pipeline_type: str, run_id: Optional[uuid.UUID] = None
    ):
        run_id = run_id or run_id_var.get()
        if run_id is None:
            run_id = self.generate_run_id()
            token = run_id_var.set(run_id)
            self.run_info[run_id] = {"pipeline_type": pipeline_type}
        else:
            token = run_id_var.set(run_id)
        return run_id, token

    async def get_run_info(self):
        run_id = run_id_var.get()
        return self.run_info.get(run_id, None)

    async def log_run_info(
        self,
        key: str,
        value: Any,
        is_info_log: bool = False,
        user: Optional[User] = None,
    ):
        if run_id := run_id_var.get():
            await self.logger.log(
                log_id=run_id,
                key=key,
                value=value,
                user_id=user.id if user else None,
                is_info_log=is_info_log,
            )

    async def clear_run_info(self, token: contextvars.Token):
        run_id = run_id_var.get()
        run_id_var.reset(token)
        if run_id and run_id in self.run_info:
            del self.run_info[run_id]


@asynccontextmanager
async def manage_run(
    run_manager: RunManager,
    pipeline_type: str,
    run_id: Optional[uuid.UUID] = None,
):
    run_id, token = await run_manager.set_run_info(pipeline_type, run_id)
    try:
        yield run_id
    finally:
        # Note: Do not clear the run info to ensure the run ID remains the same
        run_id_var.reset(token)
