from contextlib import asynccontextmanager
import asyncio
import logging
import uuid
from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, AsyncGenerator, Optional

from pydantic import BaseModel

from ..utils import generate_run_id

logger = logging.getLogger(__name__)


class PipeType(Enum):
    INGESTOR = "ingestor"
    EVAL = "eval"
    GENERATOR = "generator"
    SEARCH = "search"
    TRANSFORM = "transform"
    OTHER = "other"


class PipeRunInfo(BaseModel):
    """Information about a pipe run."""

    run_id: uuid.UUID


class AsyncState:
    """A state object for storing data between pipes."""

    def __init__(self):
        self.data = {}
        self.lock = asyncio.Lock()

    async def update(self, outer_key: str, values: dict):
        """Update the state with new values."""
        async with self.lock:
            if not isinstance(values, dict):
                raise ValueError("Values must be contained in a dictionary.")
            if outer_key not in self.data:
                self.data[outer_key] = {}
            for inner_key, inner_value in values.items():
                self.data[outer_key][inner_key] = inner_value

    async def get(self, outer_key: str, inner_key: str, default=None):
        """Get a value from the state."""
        async with self.lock:
            if outer_key not in self.data:
                raise ValueError(
                    f"Key {outer_key} does not exist in the state."
                )
            if inner_key not in self.data[outer_key]:
                return default or {}
            return self.data[outer_key][inner_key]

    async def delete(self, outer_key: str, inner_key: Optional[str] = None):
        """Delete a value from the state."""
        async with self.lock:
            if outer_key in self.data and not inner_key:
                del self.data[outer_key]
            else:
                if inner_key not in self.data[outer_key]:
                    raise ValueError(
                        f"Key {inner_key} does not exist in the state."
                    )
                del self.data[outer_key][inner_key]


@asynccontextmanager
async def manage_run_info(
    pipe: "AsyncPipe", run_id: Optional[uuid.UUID] = None
):
    try:
        run_id = run_id or generate_run_id()
        pipe._run_info = PipeRunInfo(run_id=run_id)
        yield
    finally:
        pipe.run_id = None


class AsyncPipe(ABC):
    """An asynchronous pipe for processing data."""

    class PipeConfig(BaseModel):
        """Configuration for a pipe."""

        name: str = "default_pipe"

        class Config:
            extra = "forbid"
            arbitrary_types_allowed = True

    class Input(BaseModel):
        """Input for a pipe."""

        message: AsyncGenerator[Any, None]

        class Config:
            extra = "forbid"
            arbitrary_types_allowed = True

    def __init__(
        self,
        type: PipeType = PipeType.OTHER,
        config: Optional[PipeConfig] = None,
    ):
        self._config = config or self.PipeConfig()
        self._run_info = None
        self._type = type

        logger.info(f"Initialized pipe {self.config.name} of type {self.type}")

    @property
    def config(self) -> PipeConfig:
        return self._config

    @property
    def type(self) -> PipeType:
        return self._type

    @property
    def run_info(self) -> PipeRunInfo:
        return self._run_info

    async def run(
        self,
        input: Input,
        state: AsyncState,
        run_id: Optional[uuid.UUID] = None,
        *args: Any,
        **kwargs: Any,
    ) -> AsyncGenerator[Any, None]:
        async with manage_run_info(self, run_id):
            result = self._run_logic(input, state)
            return result

    @abstractmethod
    async def _run_logic(
        self, input: Input, state: AsyncState, *args: Any, **kwargs: Any
    ) -> AsyncGenerator[Any, None]:
        pass
