"""Sampling Strategies for budget forcing generation."""

from copy import deepcopy

import tqdm

from ...backends.ollama import OllamaModelBackend
from ...core import (
    Backend,
    BaseModelSubclass,
    Component,
    Context,
    FancyLogger,
    ModelOutputThunk,
    Requirement,
    S,
    SamplingResult,
    ValidationResult,
)
from ...stdlib import functional as mfuncs
from .base import RejectionSamplingStrategy
from .sampling_algos import think_budget_forcing


class BudgetForcingSamplingStrategy(RejectionSamplingStrategy):
    """Sampling strategy that enforces a token budget for chain-of-thought reasoning.

    Extends ``RejectionSamplingStrategy`` with explicit control over the ``<think>``
    block size and the answer block size. On each loop iteration,
    ``think_budget_forcing`` interleaves forced-thinking and final-answer generation,
    after which the standard rejection-sampling validation pass determines whether to
    accept or retry.

    Currently only supports the ``OllamaModelBackend``.

    Args:
        think_max_tokens (int | None): Tokens allocated for the thinking block.
            Defaults to ``4096``.
        answer_max_tokens (int | None): Tokens allocated for the answer block.
            ``None`` means unbounded.
        start_think_token (str | None): Token opening the thinking block.
            Defaults to ``"<think>"``.
        end_think_token (str | None): Token closing the thinking block.
            Defaults to ``"</think>"``.
        begin_response_token (str | None): Optional token opening the response
            block. Defaults to ``""``.
        end_response_token (str): Token closing the response block.
        think_more_suffix (str | None): Suffix to force continued thinking.
            Empty string disables forcing.
        answer_suffix (str | None): Suffix to elicit a final answer.
        loop_budget (int): Rejection-sampling loop count. Must be > 0.
        requirements (list[Requirement] | None): Requirements to validate against.

    """

    think_max_tokens: int | None
    answer_max_tokens: int | None
    start_think_token: str | None
    end_think_token: str | None
    begin_response_token: str | None
    end_response_token: str
    think_more_suffix: str | None
    answer_suffix: str | None

    def __init__(
        self,
        *,
        think_max_tokens: int | None = 4096,
        answer_max_tokens: int | None = None,
        start_think_token: str | None = "<think>",
        end_think_token: str | None = "</think>",
        begin_response_token: str | None = "",
        end_response_token: str = "",
        think_more_suffix: str | None = "",
        answer_suffix: str | None = "",
        loop_budget: int = 1,
        requirements: list[Requirement] | None,
    ):
        r"""Initialize BudgetForcingSamplingStrategy with token budget parameters and rejection-sampling settings.

        Raises:
            AssertionError: If loop_budget is not greater than 0.
        """
        super().__init__(loop_budget=loop_budget, requirements=requirements)
        self.think_max_tokens = think_max_tokens
        self.answer_max_tokens = answer_max_tokens
        self.start_think_token = start_think_token
        self.end_think_token = end_think_token
        self.begin_response_token = begin_response_token
        self.end_response_token = end_response_token
        self.think_more_suffix = think_more_suffix
        self.answer_suffix = answer_suffix

    async def sample(
        self,
        action: Component[S],
        context: Context,
        backend: Backend,
        requirements: list[Requirement] | None,
        *,
        validation_ctx: Context | None = None,
        format: type[BaseModelSubclass] | None = None,
        model_options: dict | None = None,
        tool_calls: bool = False,
        show_progress: bool = True,
    ) -> SamplingResult[S]:
        """This method performs a sampling operation based on the given instruction.

        Args:
            action : The action object to be sampled.
            context: The context to be passed to the sampling strategy.
            backend: The backend used for generating samples.
            requirements: List of requirements to test against (merged with global requirements).
            validation_ctx: Optional context to use for validation. If None, validation_ctx = ctx.
            format: output format for structured outputs.
            model_options: model options to pass to the backend during generation / validation.
            tool_calls: True if tool calls should be used during this sampling strategy.
            show_progress: if true, a tqdm progress bar is used. Otherwise, messages will still be sent to flog.

        Returns:
            SamplingResult: A result object indicating the success or failure of the sampling process.

        Raises:
            AssertionError: Asserts that all required components (repair, select_from_failure, validate, and generate) are provided before proceeding with the sampling.
        """
        validation_ctx = validation_ctx if validation_ctx is not None else context

        flog = FancyLogger.get_logger()

        sampled_results: list[ModelOutputThunk] = []
        sampled_scores: list[list[tuple[Requirement, ValidationResult]]] = []
        sampled_actions: list[Component] = []
        sample_contexts: list[Context] = []

        # The `logging_redirect_tqdm` approach did not work, so instead we will use the show_progress
        # flag to determine whether we should show the pbar.
        show_progress = show_progress and flog.getEffectiveLevel() <= FancyLogger.INFO

        reqs = []
        # global requirements supersede local requirements (global requirements can be defined by user)
        # Todo: re-evaluate if this makes sense
        if self.requirements is not None:
            reqs += self.requirements
        elif requirements is not None:
            reqs += requirements
        reqs = list(set(reqs))

        loop_count = 0
        loop_budget_range_iterator = (
            tqdm.tqdm(range(self.loop_budget))  # type: ignore
            if show_progress
            else range(self.loop_budget)  # type: ignore
        )

        next_action = deepcopy(action)
        next_context = context
        for _ in loop_budget_range_iterator:  # type: ignore
            loop_count += 1
            if not show_progress:
                flog.info(f"Running loop {loop_count} of {self.loop_budget}")

            # TODO
            # tool_calls is not supported for budget forcing
            assert tool_calls is False, (
                "tool_calls is not supported with budget forcing"
            )
            # TODO
            assert isinstance(backend, OllamaModelBackend), (
                "Only ollama backend supported with budget forcing"
            )
            # run a generation pass with budget forcing
            result = await think_budget_forcing(
                backend,
                next_action,
                ctx=context,
                format=format,
                tool_calls=tool_calls,
                think_max_tokens=self.think_max_tokens,
                answer_max_tokens=self.answer_max_tokens,
                start_think_token=self.start_think_token,
                end_think_token=self.end_think_token,
                think_more_suffix=self.think_more_suffix,
                answer_suffix=self.answer_suffix,
                model_options=model_options,
            )
            result_ctx = next_context

            # Sampling strategies may use different components from the original
            # action. This might cause discrepancies in the expected parsed_repr
            # type / value. Explicitly overwrite that here.
            result.parsed_repr = action.parse(result)

            # validation pass
            val_scores_co = mfuncs.avalidate(
                reqs=reqs,
                context=result_ctx,
                backend=backend,
                output=result,
                format=format,
                model_options=model_options,
                # tool_calls=tool_calls  # Don't support using tool calls in validation strategies.
            )
            val_scores = await val_scores_co

            # match up reqs with scores
            constraint_scores = list(zip(reqs, val_scores))

            # collect all data
            sampled_results.append(result)
            sampled_scores.append(constraint_scores)
            sampled_actions.append(next_action)
            sample_contexts.append(result_ctx)

            # if all vals are true -- break and return success
            if all(bool(s[1]) for s in constraint_scores):
                flog.info("SUCCESS")
                assert (
                    result._generate_log is not None
                )  # Cannot be None after generation.
                result._generate_log.is_final_result = True

                # SUCCESS !!!!
                return SamplingResult(
                    result_index=len(sampled_results) - 1,
                    success=True,
                    sample_generations=sampled_results,
                    sample_validations=sampled_scores,
                    sample_contexts=sample_contexts,
                    sample_actions=sampled_actions,
                )

            else:
                # log partial success and continue
                count_valid = len([s for s in constraint_scores if bool(s[1])])
                flog.info(f"FAILED. Valid: {count_valid}/{len(constraint_scores)}")

            # If we did not pass all constraints, update the instruction and try again.
            next_action, next_context = self.repair(
                next_context,
                result_ctx,
                sampled_actions,
                sampled_results,
                sampled_scores,
            )

        flog.info(
            f"Invoking select_from_failure after {len(sampled_results)} failed attempts."
        )

        # if no valid result could be determined, find a last resort.
        best_failed_index = self.select_from_failure(
            sampled_actions, sampled_results, sampled_scores
        )
        assert best_failed_index < len(sampled_results), (
            "The select_from_failure method did not return a valid result. It has to selected from failed_results."
        )

        assert (
            sampled_results[best_failed_index]._generate_log is not None
        )  # Cannot be None after generation.
        sampled_results[best_failed_index]._generate_log.is_final_result = True  # type: ignore

        return SamplingResult(
            result_index=best_failed_index,
            success=False,
            sample_generations=sampled_results,
            sample_validations=sampled_scores,
            sample_actions=sampled_actions,
            sample_contexts=sample_contexts,
        )
