from datetime import datetime
from typing import Any, Dict, List, Optional
from swarms.structs.serialization import SerializableMixin
from swarms.structs.agent import Agent


class SelfMoASeq(SerializableMixin):
    """
    Self-MoA-Seq: Sequential Self-Mixture of Agents

    An ensemble method that generates multiple outputs from a single
    high-performing model and aggregates them sequentially using a
    sliding window approach. This addresses context length constraints
    while maintaining the effectiveness of in-model diversity.

    Architecture:
    - Phase 1: Generate initial samples from the proposer model
    - Phase 2: Aggregate outputs using sliding window with synthesized bias
    - Phase 3: Iterate until all samples are processed
    """

    def __init__(
        self,
        name: str = "SelfMoASeq",
        description: str = "Self-MoA-Seq: Sequential Self-Mixture of Agents",
        model_name: str = "gpt-5.4",
        temperature: float = 0.7,
        window_size: int = 6,
        reserved_slots: int = 3,
        max_loops: int = 10,
        max_tokens: int = 2000,
        num_samples: int = 30,
        enable_logging: bool = True,
        log_level: str = "INFO",
        verbose: bool = True,
        proposer_model_name: Optional[str] = None,
        aggregator_model_name: Optional[str] = None,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        retry_backoff_multiplier: float = 2.0,
        retry_max_delay: float = 60.0,
        additional_kwargs: Dict[str, Any] = {},
        top_p: Optional[float] = None,
    ):
        # Store parameters
        self.model_name = model_name
        self.temperature = temperature
        self.window_size = window_size
        self.reserved_slots = reserved_slots
        self.max_loops = max_loops
        self.max_tokens = max_tokens
        self.num_samples = num_samples
        self.enable_logging = enable_logging
        self.log_level = log_level
        self.verbose = verbose

        # Retry configuration
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.retry_backoff_multiplier = retry_backoff_multiplier
        self.retry_max_delay = retry_max_delay

        self.setup()

        # Allow model overrides
        proposer_model = proposer_model_name or self.model_name
        aggregator_model = aggregator_model_name or self.model_name

        # Setup logging
        self._log(
            "info",
            f"Initializing Self-MoA-Seq with model: {self.model_name}",
        )

        # Initialize proposer agent (generates multiple samples)
        self.proposer = Agent(
            agent_name="SelfMoASeq-Proposer",
            system_prompt=(
                "You are a sample generator. Generate diverse, high-quality responses "
                "to the given task. Vary your approach while maintaining quality."
            ),
            model_name=proposer_model,
            temperature=self.temperature,
            max_loops=1,
            verbose=self.verbose,
            top_p=top_p,
        )

        # Initialize aggregator agent (synthesizes outputs)
        self.aggregator = Agent(
            agent_name="SelfMoASeq-Aggregator",
            system_prompt=(
                "You are an expert synthesizer. Analyze the provided responses and "
                "synthesize them into a single, high-quality output. Consider the "
                "strengths of each response and combine them effectively. Pay special "
                "attention to any highlighted best response, as it provides high-quality guidance."
            ),
            model_name=aggregator_model,
            temperature=0.0,  # Deterministic aggregation
            max_loops=1,
            verbose=self.verbose,
            top_p=top_p,
        )

        # Metrics tracking
        self.metrics: Dict[str, Any] = {
            "total_samples_generated": 0,
            "total_aggregations": 0,
            "total_tokens_used": 0,
            "execution_time_seconds": 0,
        }

        self._log("info", "Self-MoA-Seq initialization complete")

    def setup(self) -> None:
        """
        Validate constructor parameters stored on the instance.

        Raises:
            ValueError: If any configuration value falls outside its allowed
                range. See the constructor docstring for the full set of
                constraints.
        """
        if self.window_size < 2:
            raise ValueError("window_size must be at least 2")
        if self.reserved_slots >= self.window_size:
            raise ValueError(
                "reserved_slots must be less than window_size"
            )
        if not 0 <= self.temperature <= 2:
            raise ValueError("temperature must be between 0 and 2")
        if self.max_loops < 1:
            raise ValueError("max_loops must be at least 1")
        if self.num_samples < 2:
            raise ValueError("num_samples must be at least 2")
        if self.max_retries < 0:
            raise ValueError("max_retries must be non-negative")
        if self.retry_delay < 0:
            raise ValueError("retry_delay must be non-negative")
        if self.retry_backoff_multiplier < 1:
            raise ValueError("retry_backoff_multiplier must be >= 1")
        if self.retry_max_delay < self.retry_delay:
            raise ValueError("retry_max_delay must be >= retry_delay")

    def _generate_samples(
        self, task: str, num_samples: int
    ) -> List[str]:
        """
        Generate multiple samples from the proposer model.

        Args:
            task: The task description
            num_samples: Number of samples to generate

        Returns:
            List of generated samples
        """
        self._log(
            "info", f"Generating {num_samples} samples for task"
        )
        samples = []

        try:
            for i in range(num_samples):
                self._log(
                    "debug", f"Generating sample {i+1}/{num_samples}"
                )
                sample = self.proposer.run(task)
                samples.append(sample)
                self.metrics["total_samples_generated"] += 1

            self._log(
                "success",
                f"Successfully generated {len(samples)} samples",
            )
            return samples

        except Exception as e:
            self._log(
                "error", f"Error during sample generation: {str(e)}"
            )
            raise

    def _format_aggregation_prompt(
        self,
        task: str,
        samples: List[str],
        best_so_far: Optional[str] = None,
    ) -> str:
        """
        Format the aggregation prompt with sliding window.

        Args:
            task: Original task
            samples: List of samples to aggregate
            best_so_far: Previously synthesized best output

        Returns:
            Formatted aggregation prompt
        """
        prompt = f"Original Task:\n{task}\n\n"

        if best_so_far:
            prompt += f"Current Best Response (synthesized from previous iterations):\n{best_so_far}\n\n"

        prompt += "Candidate Responses to Synthesize:\n"
        for i, sample in enumerate(samples, 1):
            prompt += f"\n[Response {i}]:\n{sample}\n"

        prompt += (
            "\nProvide a comprehensive synthesis that combines the strengths of "
            "all responses while maintaining coherence and quality."
        )

        return prompt

    def _aggregate_window(
        self,
        task: str,
        window_samples: List[str],
        best_so_far: Optional[str] = None,
    ) -> str:
        """
        Aggregate a window of samples.

        Args:
            task: Original task
            window_samples: Samples in current window
            best_so_far: Best aggregation so far

        Returns:
            Synthesized output
        """
        self._log(
            "debug",
            f"Aggregating window of {len(window_samples)} samples",
        )

        try:
            prompt = self._format_aggregation_prompt(
                task,
                window_samples,
                best_so_far,
            )

            aggregated = self.aggregator.run(prompt)
            self.metrics["total_aggregations"] += 1

            self._log("debug", "Window aggregation complete")
            return aggregated

        except Exception as e:
            self._log(
                "error", f"Error during window aggregation: {str(e)}"
            )
            raise

    def run(
        self,
        task: str,
    ) -> Dict[str, Any]:
        """
        Execute Self-MoA-Seq on the given task.

        This method implements the sequential aggregation algorithm:
        1. Generate num_samples from the proposer model
        2. Use sliding window to aggregate in chunks
        3. Progressively synthesize outputs, biasing aggregator toward best
        4. Return final synthesized output

        Args:
            task: The task to process

        Returns:
            Dictionary containing:
                - final_output: The best synthesized response
                - all_samples: List of generated samples
                - aggregation_steps: Number of aggregation iterations
                - metrics: Performance metrics
        """
        self._log(
            "info",
            f"Starting Self-MoA-Seq run with {self.num_samples} samples",
        )
        start_time = datetime.now()

        try:
            # Validate input
            if not task or not isinstance(task, str):
                raise ValueError("task must be a non-empty string")

            # Phase 1: Generate samples
            self._log("info", "Phase 1: Generating initial samples")
            samples = self._generate_samples(task, self.num_samples)

            # Phase 2: Sequential aggregation with sliding window
            self._log("info", "Phase 2: Sequential aggregation")
            best_output = samples[0]
            aggregation_step = 0

            # Process samples in windows
            remaining_samples = samples[1:]

            while remaining_samples:
                aggregation_step += 1
                self._log(
                    "info",
                    f"Aggregation iteration {aggregation_step}, "
                    f"remaining samples: {len(remaining_samples)}",
                )

                # Create window: reserved slots + new samples
                window_size = min(
                    self.window_size - self.reserved_slots,
                    len(remaining_samples),
                )
                current_window = remaining_samples[:window_size]
                remaining_samples = remaining_samples[window_size:]

                # Aggregate with bias toward best output
                window_with_best = [best_output] + current_window
                best_output = self._aggregate_window(
                    task,
                    window_with_best,
                    best_output,
                )

                if aggregation_step >= self.max_loops:
                    self._log(
                        "warning",
                        f"Reached max aggregation loops ({self.max_loops})",
                    )
                    break

            # Calculate metrics
            elapsed = (datetime.now() - start_time).total_seconds()
            self.metrics["execution_time_seconds"] = elapsed

            result = {
                "final_output": best_output,
                "all_samples": samples,
                "aggregation_steps": aggregation_step,
                "metrics": self.metrics.copy(),
                "task": task,
                "timestamp": datetime.now().isoformat(),
            }

            self._log(
                "success",
                f"Self-MoA-Seq completed in {elapsed:.2f}s "
                f"with {aggregation_step} aggregation iterations",
            )

            if self.verbose:
                self._log(
                    "info",
                    f"Self MoA Seq result: {result}",
                )

            return result

        except Exception as e:
            self._log(
                "error", f"Fatal error in Self-MoA-Seq.run: {str(e)}"
            )
            raise

    def get_metrics(self) -> Dict[str, Any]:
        """Get current performance metrics."""
        return self.metrics.copy()
