import dataclasses

from lmwrapper.abstract_predictor import LmPredictor
from lmwrapper.structs import ChatGptRoles, LmChatDialog, LmChatTurn

from synthegrator.code_problems import CodeProblem
from synthegrator.few_shotting import FewShotConfig
from synthegrator.problem_rendering import LmPromptRender, PromptRenderer
from synthegrator.response_parser import (
    AnswerForSolutionInvalidForSolve,
    AnswerForSolutionNeverSupported,
    ResponseParser,
)


def get_answer_for_each(
    problems: list[CodeProblem],
    renderer: PromptRenderer,
    parser: ResponseParser,
    lm: LmPredictor,
) -> list[tuple[LmPromptRender, str]]:
    out = []
    for problem in problems:
        render = renderer.render(problem, lm)
        if len(problem.known_solutions) == 0:
            continue
        try:
            ans = parser.answer_for_solution(
                render, problem, problem.known_solutions[0].solve_steps
            )
        except AnswerForSolutionNeverSupported:
            return []
        except AnswerForSolutionInvalidForSolve:
            continue
        ans_starts_with_exactly_one_space = ans.startswith(" ") and not ans.startswith(
            "  "
        )
        if lm.is_chat_model and ans_starts_with_exactly_one_space:
            ans = ans[1:]
        out.append((render, ans))
    return out


def get_default_text_seperator() -> str:
    return "\n###\n"


def combine_for_few_shot(
    render_and_answers: list[tuple[LmPromptRender, str]],
    lm: LmPredictor,
    text_seperator: str = get_default_text_seperator(),
) -> LmChatDialog | str:
    if lm.is_chat_model:
        messages = []
        for render, ans in render_and_answers:
            if not render.prompt.is_text_a_chat():
                text = render.prompt.get_text_as_chat()
            else:
                text = render.prompt.text
            assert text[-1].role == ChatGptRoles.user
            messages.extend(text)
            messages.append(LmChatTurn(ChatGptRoles.assistant, ans))
        return LmChatDialog(messages)
    messages = []
    for render, ans in render_and_answers:
        assert not render.prompt.is_text_a_chat()
        text = render.prompt.text
        assert isinstance(text, str)
        messages.append(text + ans)
    if len(messages) == 0:
        return ""
    return text_seperator.join(messages) + text_seperator


def generic_prepare_few_shot(
    target_problem: CodeProblem,
    few_shot_config: FewShotConfig,
    renderer: PromptRenderer,
    parser: ResponseParser,
    lm: LmPredictor,
) -> LmPromptRender:
    problems = few_shot_config.library.sample_problems(
        n=few_shot_config.num_examples,
        target_problem=target_problem,
    )
    render_and_answers = get_answer_for_each(problems, renderer, parser, lm)
    few_shot_context = combine_for_few_shot(render_and_answers, lm)
    target_render = renderer.render(target_problem, lm)
    new_prompt = dataclasses.replace(
        target_render.prompt,
        text=few_shot_context
        + (
            target_render.prompt.get_text_as_chat()
            if lm.is_chat_model
            else target_render.prompt.get_text_as_string_default_form()
        ),
    )
    new_render = dataclasses.replace(
        target_render,
        prompt=new_prompt,
    )
    return new_render


class FewShotRendererWrapper(PromptRenderer):
    """
    Wraps a prompt render to be add support
    for few-shotting for a given renderer as long
    as the response parser supports answer_for_solution.
    """

    def __init__(
        self,
        renderer: PromptRenderer,
        response_parser: ResponseParser,
        few_shot_config: FewShotConfig,
    ):
        super().__init__()
        self.renderer = renderer
        self.parser = response_parser
        self.few_shot_config = few_shot_config

    def render(
        self, problem: CodeProblem, lm: LmPredictor, prompt_seed: int | None = None
    ) -> LmPromptRender:
        return generic_prepare_few_shot(
            problem, self.few_shot_config, self.renderer, self.parser, lm
        )
