import inspect
import random
import string
import uuid
from abc import ABC, abstractmethod
from enum import Enum
from types import SimpleNamespace
from typing import Dict, List, Callable, Optional, Union, Any, get_type_hints

from dria_workflows import Workflow
from pydantic import BaseModel, create_model, ValidationError

from dria.client import Dria
from dria.models import TaskInput, CallbackType, Model
from dria.pipelines.callbacks import (
    aggregation_callback,
    broadcast_callback,
    default_callback,
    scatter_callback,
)
from dria.pipelines.config import StepConfig, PipelineConfig
from dria.pipelines.pipeline import Pipeline
from dria.pipelines.step import Step
from dria.utils.logging import logger


class StepType(str, Enum):
    FIRST = "first"
    MIDDLE = "middle"
    LAST = "last"


class StepTemplate(BaseModel, ABC):
    step_type: StepType = StepType.FIRST
    workflow: Optional[Workflow] = None
    params: SimpleNamespace = SimpleNamespace()
    config: StepConfig = StepConfig()

    class Config:
        arbitrary_types_allowed = True

    def __init__(self, config: Optional[StepConfig] = None, **kwargs):
        super().__init__(**kwargs)
        if config:
            self.config = config
        self._map_io(**kwargs)

    def get_params(self) -> dict:
        return vars(self.params)

    def _generate_dummy_inputs(self, func: Callable) -> Dict[str, Any]:
        dummy_inputs = {}
        type_hints = get_type_hints(func)
        signature = inspect.signature(func)

        for param_name, _ in signature.parameters.items():
            if param_name in ["self", "kwargs"]:
                continue
            param_type = type_hints.get(param_name, Any)
            dummy_inputs[param_name] = self._create_dummy_data(param_type)

        return dummy_inputs

    def _create_dummy_data(self, param_type):
        if param_type == str:
            return "".join(random.choices(string.ascii_letters, k=10))
        elif param_type == int:
            return random.randint(0, 100)
        elif param_type == float:
            return random.uniform(0, 100)
        elif param_type == bool:
            return random.choice([True, False])
        elif param_type == List[str]:
            return [
                "".join(random.choices(string.ascii_letters, k=5)) for _ in range(3)
            ]
        elif param_type == List[int]:
            return [random.randint(0, 100) for _ in range(3)]
        elif param_type == Dict[str, Any]:
            return {f"key_{i}": self._create_dummy_data(str) for i in range(3)}
        # Handling TypedDict types like 'Article'
        elif isinstance(param_type, type) and hasattr(param_type, "__annotations__"):
            dummy_data = {}
            # Generate dummy data for each field in the TypedDict
            for field, field_type in param_type.__annotations__.items():
                dummy_data[field] = self._create_dummy_data(field_type)
            return dummy_data
        else:
            return None

    def _map_io(self, **kwargs):

        for k, v in kwargs.items():
            setattr(self.params, k, v)

        self.params.inputs = []
        workflow_inputs = self._generate_dummy_inputs(self.create_workflow)
        task_input = TaskInput(**workflow_inputs)
        self.params.task_input = task_input
        for k, v in workflow_inputs.items():
            self.params.inputs.append(k)

        self.workflow = self.create_workflow(**workflow_inputs)

        if not self.workflow.return_value:
            raise ValueError("Workflow must have a return value.")

        self.params.output = self.workflow.return_value.input.key
        self.params.output_type = self.workflow.return_value.input.type
        self.params.output_json = self.workflow.return_value.to_json
        self.params.callback = self.callback
        self.params.callback_type = CallbackType.DEFAULT

    def broadcast(self, n: int) -> "StepTemplate":
        """

        Args:
            n: broadcast to N nodes

        Returns:

        """
        self.params.callback = broadcast_callback
        self.params.callback_params = {"n": n}
        self.params.callback_type = CallbackType.BROADCAST
        return self

    def scatter(self) -> "StepTemplate":
        """
        Scatter N outputs to N nodes
        Returns:

        """
        self.params.callback = scatter_callback
        self.params.callback_type = CallbackType.SCATTER
        return self

    def aggregate(self) -> "StepTemplate":
        """
        Aggregate N outputs to 1 node
        Returns:

        """
        self.params.callback = aggregation_callback
        self.params.callback_type = CallbackType.AGGREGATE
        return self

    def custom(self) -> "StepTemplate":
        """
        Custom callback
        Returns:

        """
        self.params.callback = self.callback
        self.params.callback_type = CallbackType.CUSTOM
        return self

    @abstractmethod
    def create_workflow(self, **kwargs) -> "Workflow":
        pass

    def set_models(self, models: Optional[List[Model]] = None) -> "StepTemplate":
        if models:
            self.config.models = models
        return self

    def set_workflow_params(
            self,
            min_compute: Optional[int] = None,
            max_time: Optional[int] = None,
            max_steps: Optional[int] = None,
            max_tokens: Optional[int] = None,
    ) -> "StepTemplate":
        if min_compute:
            self.config.min_compute = min_compute
        if max_time:
            self.config.max_time = max_time
        if max_steps:
            self.config.max_steps = max_steps
        if max_tokens:
            self.config.max_tokens = max_tokens
        return self

    def callback(self, step: "Step") -> Union[List[TaskInput], TaskInput]:
        if self.__class__.callback is not StepTemplate.callback:
            return self.__class__.callback(self, step)
        return default_callback(step)

    @staticmethod
    def parse(result) -> Any:
        pass


class StepBuilder:
    """
    A builder class for creating and configuring pipelines steps.

    This class provides a fluent interface for step creation and configuration,
    allowing for easy setup of step properties such as callbacks and configurations.
    """

    def __init__(
            self,
            name: Optional[str] = None,
            input: Optional[Union[TaskInput, List[TaskInput]]] = None,
            workflow: Optional[Union[Callable, Workflow]] = None,
            config: StepConfig = StepConfig(),
    ):
        """
        Initialize a new StepBuilder instance.

        Args:
            name (Optional[str]): The name of the step. If not provided, a UUID will be generated.
            input (Optional[Union[TaskInput, List[TaskInput]]]): The input(s) for the step.
            workflow (Optional[Callable]): The workflow function for the step.
            config (StepConfig): The configuration for the step. Defaults to an empty StepConfig.
        """
        self.step = Step(
            name=name or str(uuid.uuid4()),
            input=input,
            workflow=workflow,
            config=config,
        )
        logger.debug(f"Created StepBuilder with name: {self.step.name}")

    def add_callback(
            self,
            callback: Callable,
            callback_type: CallbackType,
            params: Optional[Dict] = None,
    ) -> "StepBuilder":
        """
        Set a custom callback for the step.

        Args:
            callback (Callable): Custom callback function.
            callback_type (CallbackType): Type of the custom callback.
            params (Optional[Dict]): Extra parameters for the custom callback.

        Returns:
            StepBuilder: Self instance for method chaining.
        """
        if not callable(callback):
            raise ValueError("Custom callback must be callable")

        if not any(
                param.annotation in [Step, "Step"]
                for param in inspect.signature(callback).parameters.values()
        ):
            raise ValueError("Custom callback must have a parameter of type Step")

        self.step.callback = callback
        self.step.callback_type = callback_type
        self.step.callback_params = params or {}
        logger.debug(f"Set custom callback for step: {self.step.name}")
        return self

    def set_config(self, config: StepConfig) -> "StepBuilder":
        """
        Set the configuration for the step.

        Args:
            config (StepConfig): The configuration for the step.

        Returns:
            StepBuilder: Self instance for method chaining.
        """
        self.step.config = config
        logger.debug(f"Set config for step: {self.step.name}")
        return self

    def build(self) -> Step:
        """
        Build and return the configured Step instance.

        Returns:
            Step: The fully configured Step instance.
        """
        if not self.step.config:
            self.step.config = StepConfig()

        if self.step.callback is None:
            self.step.callback = default_callback
            self.step.callback_type = CallbackType.DEFAULT

        logger.info(f"Built step: {self.step.name}")
        return self.step


class PipelineBuilder:
    """
    A builder class for creating and configuring pipelines.

    This class provides methods to add steps to a pipelines and build the final Pipeline instance.
    """

    def __init__(self, config: PipelineConfig, dria_client: Dria):
        """
        Initialize a new PipelineBuilder instance.

        Args:
            config (PipelineConfig): The configuration for the pipelines.
            dria_client (Dria): The Dria client instance.
        """
        self.steps: List[Step] = []
        self.templates: List[StepTemplate] = []
        self.config = config
        self.dria_client = dria_client
        self.pipeline_input: Optional[Dict] = None
        logger.info("Initialized PipelineBuilder")

    def _check_unique_step_names(self):
        """
        Check if all step names in the pipelines are unique.

        Raises:
            ValueError: If duplicate step names are found.
        """
        step_names = set()
        for step in self.steps:
            if step.name in step_names:
                raise ValueError(
                    f"Duplicate step name found: {step.name}. Each step must have a unique name."
                )
            step_names.add(step.name)

    def input(self, **kwargs):
        self.pipeline_input = kwargs
        return self

    def _validate_input_types(self, step_template: StepTemplate):
        if not self.pipeline_input:
            raise ValueError(
                "Pipeline input not set. Use .input() method before adding steps."
            )

        expected_types = get_type_hints(step_template.create_workflow)
        input_model = create_model(
            "InputModel",
            **{k: (v, ...) for k, v in expected_types.items() if k != "return"},
        )

        try:
            input_model(**self.pipeline_input)
        except ValidationError as e:
            error_messages = []
            for error in e.errors():
                if error["type"] == "value_error.missing":
                    error_messages.append(f"Missing input field: {error['loc'][0]}")
                else:
                    error_messages.append(
                        f"Invalid input for field '{error['loc'][0]}': {error['msg']}"
                    )
            raise ValueError("\n".join(error_messages))

    def _add_step(self, step: Step) -> "PipelineBuilder":
        """
        Add a step to the pipelines.

        Args:
            step (Step): Step to add to the pipelines.

        Returns:
            PipelineBuilder: Self instance for method chaining.

        Raises:
            ValueError: If the step configuration is invalid.
        """
        try:
            # Check if step.workflow is callable
            if not callable(step.workflow):
                workflow = step.workflow
            else:
                workflow = step.workflow({})
            step_inputs = [p.name for task in workflow.tasks for p in task.inputs]
            step_outputs = [p.key for task in workflow.tasks for p in task.outputs]
            previous_step = self.steps[-1] if self.steps else None
            step.input_keys = list(
                set([key for key in step_inputs if key not in step_outputs])
            )
            if previous_step:
                previous_step.next_step_input = step.input_keys

            if step.callback_type in (
                    CallbackType.BROADCAST,
                    CallbackType.AGGREGATE,
                    CallbackType.SCATTER,
            ):
                if len(step.input_keys) > 1:
                    raise ValueError(
                        "Broadcast and Aggregate callbacks only support a single input for this step. "
                        "If you want to use multiple inputs, use a custom callback."
                    )

            self.steps.append(step)
        except Exception as e:
            logger.error(f"Error adding step {step.name}: {str(e)}")
            raise
        return self

    def __lshift__(self, other: StepTemplate):
        """
        Add a step to the pipelines using the given StepTemplate.
        Args:
            other:

        Returns:

        """
        task_input = None
        if len(self.steps) == 0:
            self._validate_input_types(other)
            task_input = TaskInput(**self.pipeline_input)

        _step = StepBuilder(
            name=other.__class__.__name__ + f".{len(self.steps)}",
            config=other.config,
            input=task_input,
            workflow=other.workflow,
        ).add_callback(other.params.callback, other.params.callback_type)

        self._add_step(_step.build())
        self.templates.append(other)
        return self

    def build(self) -> Pipeline:
        """
        Build and return the configured Pipeline instance.

        Returns:
            Pipeline: The fully configured Pipeline instance.

        Raises:
            Exception: If there's an error during pipelines building.
        """
        try:
            pipeline = Pipeline(client=self.dria_client, config=self.config)

            self._check_unique_step_names()
            # Note: Uncomment the following line if you want to enforce no callback on the last step
            # self._check_last_step_callback()

            for step in self.steps:
                pipeline.add_step(step)
            logger.info("Built pipelines successfully")
            return pipeline
        except Exception as e:
            logger.error(f"Error building pipelines: {str(e)}")
            raise
