# SPDX-License-Identifier: Apache-2.0
"""Composite block combining LLM chat and text parsing with retry logic.

This module provides the LLMChatWithParsingRetryBlock that encapsulates the complete
LLM generation and parsing workflow with automatic retry on parsing failures.
"""

# Standard
from typing import Any, Optional

# Third Party
from datasets import Dataset
from pydantic import ConfigDict, Field, field_validator

# Local
from ...utils.error_handling import BlockValidationError
from ...utils.logger_config import setup_logger
from ..base import BaseBlock
from ..registry import BlockRegistry
from .llm_chat_block import LLMChatBlock
from .text_parser_block import TextParserBlock

logger = setup_logger(__name__)


class MaxRetriesExceededError(Exception):
    """Raised when maximum retry attempts are exceeded without achieving target count."""

    def __init__(self, target_count: int, actual_count: int, max_retries: int):
        self.target_count = target_count
        self.actual_count = actual_count
        self.max_retries = max_retries
        super().__init__(
            f"Failed to achieve target count {target_count} after {max_retries} retries. "
            f"Only got {actual_count} successful parses."
        )


@BlockRegistry.register(
    "LLMChatWithParsingRetryBlock",
    "llm",
    "Composite block combining LLM chat and text parsing with automatic retry on parsing failures",
)
class LLMChatWithParsingRetryBlock(BaseBlock):
    """Composite block for LLM generation with parsing retry logic.

    This block combines LLMChatBlock and TextParserBlock into a single cohesive block
    that automatically retries LLM generation when parsing fails, accumulating successful
    results until the target count is reached or max retries exceeded.

    Parameters
    ----------
    block_name : str
        Name of the block.
    input_cols : Union[str, List[str]]
        Input column name(s). Should contain the messages list.
    output_cols : Union[str, List[str]]
        Output column name(s) for parsed results.
    model : str
        Model identifier in LiteLLM format.
    api_base : Optional[str]
        Base URL for the API. Required for local models.
    api_key : Optional[str]
        API key for the provider. Falls back to environment variables.
    parsing_max_retries : int, optional
        Maximum number of retry attempts for parsing failures (default: 3).
        This is different from max_retries, which handles LLM network/API failures.

    ### LLM Generation Parameters ###
    async_mode : bool, optional
        Whether to use async processing (default: False).
    timeout : float, optional
        Request timeout in seconds (default: 120.0).
    max_retries : int, optional
        Maximum number of LLM retry attempts for network failures (default: 6).
    temperature : Optional[float], optional
        Sampling temperature (0.0 to 2.0).
    max_tokens : Optional[int], optional
        Maximum tokens to generate.
    top_p : Optional[float], optional
        Nucleus sampling parameter (0.0 to 1.0).
    frequency_penalty : Optional[float], optional
        Frequency penalty (-2.0 to 2.0).
    presence_penalty : Optional[float], optional
        Presence penalty (-2.0 to 2.0).
    stop : Optional[Union[str, List[str]]], optional
        Stop sequences.
    seed : Optional[int], optional
        Random seed for reproducible outputs.
    response_format : Optional[Dict[str, Any]], optional
        Response format specification (e.g., JSON mode).
    stream : Optional[bool], optional
        Whether to stream responses.
    n : Optional[int], optional
        Number of completions to generate per retry attempt.
    logprobs : Optional[bool], optional
        Whether to return log probabilities.
    top_logprobs : Optional[int], optional
        Number of top log probabilities to return.
    user : Optional[str], optional
        End-user identifier.
    extra_headers : Optional[Dict[str, str]], optional
        Additional headers to send with requests.
    extra_body : Optional[Dict[str, Any]], optional
        Additional parameters for the request body.
    provider_specific : Optional[Dict[str, Any]], optional
        Provider-specific parameters.

    ### Text Parser Parameters ###
    start_tags : List[str], optional
        List of start tags for tag-based parsing.
    end_tags : List[str], optional
        List of end tags for tag-based parsing.
    parsing_pattern : Optional[str], optional
        Regex pattern for custom parsing.
    parser_cleanup_tags : Optional[List[str]], optional
        List of tags to clean from parsed output.

    Examples
    --------
    >>> # Basic JSON parsing with retry
    >>> block = LLMChatWithParsingRetryBlock(
    ...     block_name="json_retry_block",
    ...     input_cols="messages",
    ...     output_cols="parsed_json",
    ...     model="openai/gpt-4",
    ...     parsing_max_retries=3,
    ...     parsing_pattern=r'"result":\s*"([^"]*)"',
    ...     n=3
    ... )

    >>> # Tag-based parsing with retry
    >>> block = LLMChatWithParsingRetryBlock(
    ...     block_name="tag_retry_block",
    ...     input_cols="messages",
    ...     output_cols=["explanation", "answer"],
    ...     model="anthropic/claude-3-sonnet-20240229",
    ...     parsing_max_retries=5,
    ...     start_tags=["<explanation>", "<answer>"],
    ...     end_tags=["</explanation>", "</answer>"],
    ...     n=2
    ... )
    """

    model_config = ConfigDict(
        extra="allow"
    )  # Allow extra fields for dynamic forwarding

    # --- Composite-specific configuration ---
    parsing_max_retries: int = Field(
        3, description="Maximum number of retry attempts for parsing failures"
    )

    # --- Parser configuration (required for internal TextParserBlock) ---
    start_tags: Optional[list[str]] = Field(
        None, description="Start tags for tag-based parsing"
    )
    end_tags: Optional[list[str]] = Field(
        None, description="End tags for tag-based parsing"
    )
    parsing_pattern: Optional[str] = Field(
        None, description="Regex pattern for custom parsing"
    )
    parser_cleanup_tags: Optional[list[str]] = Field(
        None, description="List of tags to clean from parsed output"
    )

    # Internal blocks - excluded from serialization
    llm_chat: Optional[LLMChatBlock] = Field(None, exclude=True)
    text_parser: Optional[TextParserBlock] = Field(None, exclude=True)

    @field_validator("input_cols")
    @classmethod
    def validate_single_input_col(cls, v):
        """Ensure exactly one input column."""
        if isinstance(v, str):
            return [v]
        if isinstance(v, list) and len(v) == 1:
            return v
        if isinstance(v, list) and len(v) != 1:
            raise ValueError(
                f"LLMChatWithParsingRetryBlock expects exactly one input column, got {len(v)}: {v}"
            )
        raise ValueError(f"Invalid input_cols format: {v}")

    @field_validator("parsing_max_retries")
    @classmethod
    def validate_parsing_max_retries(cls, v):
        """Ensure parsing_max_retries is positive."""
        if v < 1:
            raise ValueError("parsing_max_retries must be at least 1")
        return v

    def __init__(self, **kwargs):
        """Initialize with dynamic parameter routing."""
        super().__init__(**kwargs)
        self._create_internal_blocks(**kwargs)

        # Log initialization if model is configured
        if hasattr(self, "model") and self.model:
            logger.info(
                f"Initialized LLMChatWithParsingRetryBlock '{self.block_name}' with model '{self.model}'",
                extra={
                    "block_name": self.block_name,
                    "model": self.model,
                    "parsing_max_retries": self.parsing_max_retries,
                },
            )

    def _extract_params(self, kwargs: dict, block_class) -> dict:
        """Extract parameters for specific block class based on its model_fields."""
        # Exclude parameters that are handled by this wrapper
        wrapper_params = {
            "block_name",
            "input_cols",
            "output_cols",
            "parsing_max_retries",
        }

        # Extract parameters that the target block accepts
        params = {
            k: v
            for k, v in kwargs.items()
            if k in block_class.model_fields and k not in wrapper_params
        }

        # Also include declared fields from this composite block that the target block accepts
        for field_name in self.__class__.model_fields:
            if (
                field_name in block_class.model_fields
                and field_name not in wrapper_params
            ):
                field_value = getattr(self, field_name)
                if field_value is not None:  # Only forward non-None values
                    params[field_name] = field_value

        return params

    def _create_internal_blocks(self, **kwargs):
        """Create internal blocks with parameter routing."""
        # Route parameters to appropriate blocks
        llm_params = self._extract_params(kwargs, LLMChatBlock)
        parser_params = self._extract_params(kwargs, TextParserBlock)

        # 1. LLMChatBlock
        self.llm_chat = LLMChatBlock(
            block_name=f"{self.block_name}_llm_chat",
            input_cols=self.input_cols,
            output_cols=[f"{self.block_name}_raw_response"],
            **llm_params,
        )

        # 2. TextParserBlock
        self.text_parser = TextParserBlock(
            block_name=f"{self.block_name}_text_parser",
            input_cols=[f"{self.block_name}_raw_response"],
            output_cols=self.output_cols,
            **parser_params,
        )

    def __getattr__(self, name: str) -> Any:
        """Forward attribute access to appropriate internal block."""
        # Check each internal block to see which one has this parameter
        for block_attr, block_class in [
            ("llm_chat", LLMChatBlock),
            ("text_parser", TextParserBlock),
        ]:
            if hasattr(self, block_attr) and name in block_class.model_fields:
                internal_block = getattr(self, block_attr)
                if internal_block is not None:
                    return getattr(internal_block, name)
        raise AttributeError(
            f"'{self.__class__.__name__}' object has no attribute '{name}'"
        )

    def __setattr__(self, name: str, value: Any) -> None:
        """Handle dynamic parameter updates from flow.set_model_config()."""
        super().__setattr__(name, value)

        # Forward to appropriate internal blocks
        for block_attr, block_class in [
            ("llm_chat", LLMChatBlock),
            ("text_parser", TextParserBlock),
        ]:
            if hasattr(self, block_attr) and name in block_class.model_fields:
                setattr(getattr(self, block_attr), name, value)

    def _reinitialize_client_manager(self) -> None:
        """Reinitialize the internal LLM chat block's client manager.

        This should be called after model configuration changes to ensure
        the internal LLM chat block uses the updated model configuration.
        """
        if self.llm_chat and hasattr(self.llm_chat, "_reinitialize_client_manager"):
            # The parameters should already be forwarded via __setattr__ magic method
            # Just reinitialize the client manager with the current configuration
            self.llm_chat._reinitialize_client_manager()

    def generate(self, samples: Dataset, **kwargs: Any) -> Dataset:
        """Generate responses with parsing retry logic.

        For each input sample, this method:
        1. Generates LLM responses using the configured n parameter
        2. Attempts to parse the responses using TextParserBlock
        3. Counts successful parses and retries if below target
        4. Accumulates results across retry attempts
        5. Returns final dataset with all successful parses

        Parameters
        ----------
        samples : Dataset
            Input dataset containing the messages column.
        **kwargs : Any
            Additional keyword arguments passed to internal blocks.

        Returns
        -------
        Dataset
            Dataset with parsed results from successful generations.

        Raises
        ------
        BlockValidationError
            If model is not configured before calling generate().
        MaxRetriesExceededError
            If target count not reached after max retries for any sample.
        """
        # Validate that model is configured
        if not hasattr(self, "model") or not self.model:
            raise BlockValidationError(
                f"Model not configured for block '{self.block_name}'. "
                f"Call flow.set_model_config() before generating."
            )

        logger.info(
            f"Starting LLM generation with parsing retry for {len(samples)} samples",
            extra={
                "block_name": self.block_name,
                "model": self.model,
                "batch_size": len(samples),
                "parsing_max_retries": self.parsing_max_retries,
            },
        )

        all_results = []

        # Process each sample independently with retry logic
        for sample_idx, sample in enumerate(samples):
            # Determine target count for this sample (number of completions requested)
            target = kwargs.get("n", getattr(self, "n", None)) or 1

            logger.debug(
                f"Processing sample {sample_idx} with target count {target}",
                extra={
                    "block_name": self.block_name,
                    "sample_idx": sample_idx,
                    "target_count": target,
                },
            )

            if self.text_parser.expand_lists:
                # Current behavior for expand_lists=True: count rows directly
                sample_results = []
                total_parsed_count = 0

                # Retry loop for this sample
                for attempt in range(self.parsing_max_retries):
                    if total_parsed_count >= target:
                        break  # Already reached target

                    try:
                        # Generate LLM responses for this sample
                        temp_dataset = Dataset.from_list([sample])
                        llm_result = self.llm_chat.generate(temp_dataset, **kwargs)

                        # Parse the responses
                        parsed_result = self.text_parser.generate(llm_result, **kwargs)

                        # Count successful parses and accumulate results
                        new_parsed_count = len(parsed_result)
                        total_parsed_count += new_parsed_count
                        sample_results.extend(parsed_result)

                        logger.debug(
                            f"Attempt {attempt + 1} for sample {sample_idx}: {new_parsed_count} successful parses "
                            f"(total: {total_parsed_count}/{target})",
                            extra={
                                "block_name": self.block_name,
                                "sample_idx": sample_idx,
                                "attempt": attempt + 1,
                                "new_parses": new_parsed_count,
                                "total_parses": total_parsed_count,
                                "target_count": target,
                            },
                        )

                        if total_parsed_count >= target:
                            logger.debug(
                                f"Target reached for sample {sample_idx} after {attempt + 1} attempts",
                                extra={
                                    "block_name": self.block_name,
                                    "sample_idx": sample_idx,
                                    "attempts": attempt + 1,
                                    "final_count": total_parsed_count,
                                },
                            )
                            break

                    except Exception as e:
                        logger.warning(
                            f"Error during attempt {attempt + 1} for sample {sample_idx}: {e}",
                            extra={
                                "block_name": self.block_name,
                                "sample_idx": sample_idx,
                                "attempt": attempt + 1,
                                "error": str(e),
                            },
                        )
                        # Continue to next attempt
                        continue

            else:
                # New behavior for expand_lists=False: parse individual responses and accumulate
                accumulated_parsed_items = {col: [] for col in self.output_cols}
                total_parsed_count = 0

                # Retry loop for this sample
                for attempt in range(self.parsing_max_retries):
                    if total_parsed_count >= target:
                        break  # Already reached target

                    try:
                        # Generate LLM responses for this sample
                        temp_dataset = Dataset.from_list([sample])
                        llm_result = self.llm_chat.generate(temp_dataset, **kwargs)

                        # Get the raw responses (should be a list when n > 1)
                        raw_response_col = f"{self.block_name}_raw_response"
                        raw_responses = llm_result[0][raw_response_col]
                        if not isinstance(raw_responses, list):
                            raw_responses = [raw_responses]

                        # Parse each response individually and accumulate successful ones
                        new_parsed_count = 0
                        for response in raw_responses:
                            if total_parsed_count >= target:
                                break  # Stop if we've reached target

                            # Create temporary dataset with single response for parsing
                            temp_parse_data = [{**sample, raw_response_col: response}]
                            temp_parse_dataset = Dataset.from_list(temp_parse_data)

                            # Force expand_lists=True temporarily to get individual parsed items
                            original_expand_lists = self.text_parser.expand_lists
                            try:
                                self.text_parser.expand_lists = True
                                parsed_result = self.text_parser.generate(
                                    temp_parse_dataset, **kwargs
                                )
                            except Exception as parse_e:
                                logger.debug(
                                    f"Failed to parse individual response: {parse_e}"
                                )
                                continue
                            finally:
                                self.text_parser.expand_lists = original_expand_lists

                            # If parsing was successful, accumulate the results
                            if len(parsed_result) > 0:
                                for parsed_row in parsed_result:
                                    if total_parsed_count >= target:
                                        break

                                    # Only count as successful if ALL output columns are present
                                    if all(
                                        col in parsed_row for col in self.output_cols
                                    ):
                                        for col in self.output_cols:
                                            accumulated_parsed_items[col].append(
                                                parsed_row[col]
                                            )
                                        total_parsed_count += 1
                                        new_parsed_count += 1
                                    # If any column is missing, skip this parsed response entirely

                        logger.debug(
                            f"Attempt {attempt + 1} for sample {sample_idx}: {new_parsed_count} successful parses "
                            f"(total: {total_parsed_count}/{target})",
                            extra={
                                "block_name": self.block_name,
                                "sample_idx": sample_idx,
                                "attempt": attempt + 1,
                                "new_parses": new_parsed_count,
                                "total_parses": total_parsed_count,
                                "target_count": target,
                            },
                        )

                        if total_parsed_count >= target:
                            logger.debug(
                                f"Target reached for sample {sample_idx} after {attempt + 1} attempts",
                                extra={
                                    "block_name": self.block_name,
                                    "sample_idx": sample_idx,
                                    "attempts": attempt + 1,
                                    "final_count": total_parsed_count,
                                },
                            )
                            break

                    except Exception as e:
                        logger.warning(
                            f"Error during attempt {attempt + 1} for sample {sample_idx}: {e}",
                            extra={
                                "block_name": self.block_name,
                                "sample_idx": sample_idx,
                                "attempt": attempt + 1,
                                "error": str(e),
                            },
                        )
                        # Continue to next attempt
                        continue

                # Create final result row with accumulated lists
                if total_parsed_count > 0:
                    # Trim to exact target count if needed
                    for col in self.output_cols:
                        if len(accumulated_parsed_items[col]) > target:
                            accumulated_parsed_items[col] = accumulated_parsed_items[
                                col
                            ][:target]

                    # Only add the parsed output columns as lists, preserve other columns as-is
                    final_row = {**sample, **accumulated_parsed_items}
                    sample_results = [final_row]
                else:
                    sample_results = []

            # Check if we reached the target count
            if total_parsed_count < target:
                raise MaxRetriesExceededError(
                    target_count=target,
                    actual_count=total_parsed_count,
                    max_retries=self.parsing_max_retries,
                )

            # For expand_lists=True, trim results to exact target count if we exceeded it
            if self.text_parser.expand_lists and total_parsed_count > target:
                sample_results = sample_results[:target]
                logger.debug(
                    f"Trimmed sample {sample_idx} results from {total_parsed_count} to {target}",
                    extra={
                        "block_name": self.block_name,
                        "sample_idx": sample_idx,
                        "trimmed_from": total_parsed_count,
                        "trimmed_to": target,
                    },
                )

            # Add this sample's results to final dataset
            all_results.extend(sample_results)

        logger.info(
            f"LLM generation with parsing retry completed: {len(samples)} input samples → {len(all_results)} output rows",
            extra={
                "block_name": self.block_name,
                "input_samples": len(samples),
                "output_rows": len(all_results),
                "model": self.model,
            },
        )

        return Dataset.from_list(all_results)

    def _validate_custom(self, dataset: Dataset) -> None:
        """Custom validation for LLMChatWithParsingRetryBlock.

        This method validates the entire chain of internal blocks by simulating
        the data flow through each block to ensure they can all process the data correctly.
        """
        # Validate that required input column exists
        if len(self.input_cols) != 1:
            raise ValueError(
                f"LLMChatWithParsingRetryBlock expects exactly one input column, got {len(self.input_cols)}"
            )

        input_col = self.input_cols[0]
        if input_col not in dataset.column_names:
            raise ValueError(
                f"Required input column '{input_col}' not found in dataset. "
                f"Available columns: {dataset.column_names}"
            )

        # Validate parsing configuration
        has_regex = getattr(self, "parsing_pattern", None) is not None
        has_tags = bool(getattr(self, "start_tags", [])) or bool(
            getattr(self, "end_tags", [])
        )

        if not has_regex and not has_tags:
            raise ValueError(
                "LLMChatWithParsingRetryBlock requires at least one parsing method: "
                "either 'parsing_pattern' (regex) or 'start_tags'/'end_tags' (tag-based parsing)"
            )

        # Validate that internal blocks are initialized
        if not all([self.llm_chat, self.text_parser]):
            raise ValueError(
                "All internal blocks must be initialized before validation"
            )

        # Validate internal blocks
        try:
            logger.debug("Validating internal LLM chat block")
            self.llm_chat._validate_custom(dataset)

            # Create temporary dataset with expected LLM output for parser validation
            temp_data = []
            for sample in dataset:
                temp_sample = dict(sample)
                temp_sample[f"{self.block_name}_raw_response"] = "test output"
                temp_data.append(temp_sample)
            temp_dataset = Dataset.from_list(temp_data)

            logger.debug("Validating internal text parser block")
            self.text_parser._validate_custom(temp_dataset)

            logger.debug("All internal blocks validated successfully")

        except Exception as e:
            logger.error(f"Validation failed in internal blocks: {e}")
            raise ValueError(f"Internal block validation failed: {e}") from e

    def get_internal_blocks_info(self) -> dict[str, Any]:
        """Get information about the internal blocks.

        Returns
        -------
        Dict[str, Any]
            Information about each internal block.
        """
        return {
            "llm_chat": self.llm_chat.get_info() if self.llm_chat else None,
            "text_parser": self.text_parser.get_info() if self.text_parser else None,
        }

    def __repr__(self) -> str:
        """String representation of the block."""
        model = getattr(self, "model", "not_configured")
        return (
            f"LLMChatWithParsingRetryBlock(name='{self.block_name}', "
            f"model='{model}', parsing_max_retries={self.parsing_max_retries})"
        )
