import dataclasses
from timeit import default_timer as timer

import numpy as np
from lmwrapper.abstract_predictor import LmPredictor
from lmwrapper.openai_wrapper import OpenAIPredictor
from lmwrapper.structs import (
    LmPrediction,
)

from synthegrator.code_problems import (
    CodeProblem,
    CodeSolution,
    LmCodeSolution,
)
from synthegrator.lang_specs.lang_spec_python import PythonLangSpec
from synthegrator.problem_rendering import (
    LmPromptRenderSingleEdit,
    PromptRenderer,
    PromptRendererSingleEditGeneric,
)
from synthegrator.prompting_test_case_selection import (
    PromptingTestCaseSelectionStrategy,
)
from synthegrator.response_parser import ResponseParser, ResponseParserSingleEdit
from synthegrator.transformation_spec import (
    MarkElement,
    StsPlaceTransforming,
)
from synthegrator.uncertainty_modeling import (
    DelayedUncertaintyEstimate,
    PrecomputedUncertaintyEstimate,
    ProbabilityIsCorrect,
)
from synthegrator.util import (
    estimate_best_overlap_span,
)


class BaseCodeSolver:
    def solve(
        self,
        problem: CodeProblem,
        randomizing_seed: int | None = None,
    ) -> CodeSolution:
        raise NotImplementedError

    @property
    def allows_multithreading(self) -> bool:
        return False


class LmBasedSolver(BaseCodeSolver):
    """A solver that internally is based on a language model"""

    def __init__(self, model: LmPredictor, include_lm_response: bool = False):
        self._uncertainty_hooks = []
        self.model = model
        self.include_lm_response = include_lm_response

    def add_uncertainty_hook(self, hook: "LmUncertaintyHook"):
        self._uncertainty_hooks.append(hook)

    def _augment_solution_with_uncertainty_hooks(
        self,
        solution: CodeSolution,
        lm_prediction: LmPrediction,
        modified_full_text: str,
    ):
        solution.uncertainty_estimates.extend(
            [
                hook.prepare(
                    lm_prediction,
                    self.model,
                    solution,
                    modified_full_text,
                )
                for hook in self._uncertainty_hooks
            ],
        )

    @property
    def allows_multithreading(self) -> bool:
        if isinstance(self.model, OpenAIPredictor):
            return True
        return False


class DummyCodeSolverAutoRegressive(LmBasedSolver):
    """
    Solves a code problem by peaking at the known solutions and trying to
    return it.
    """

    def __init__(self, model: LmPredictor, include_lm_response: bool = False) -> None:
        super().__init__(model, include_lm_response)

    def solve(
        self,
        problem: CodeProblem,
        randomizing_seed: int | None = None,
    ) -> CodeSolution:
        return problem.known_solutions[0]


def snap_line_indent(string: str, indent: int) -> str:
    """Makes every line in the string snap to the nearest multiple of indent"""
    lines = string.split("\n")
    for i in range(len(lines)):
        stripped = lines[i].lstrip()
        current_indent = len(lines[i]) - len(stripped)
        new_indent = round(current_indent / indent) * indent
        lines[i] = " " * new_indent + stripped
    return "\n".join(lines)


class LmCodeSolverAutoRegressive(LmBasedSolver):
    """
    A code solver that works by prompting a left-to-right language model
    and parsing the response.
    """

    def __init__(
        self,
        model: LmPredictor,
        prompt_renderer: PromptRenderer = PromptRendererSingleEditGeneric(),
        response_parser: ResponseParser = ResponseParserSingleEdit(),
        prompting_test_case_strategy: PromptingTestCaseSelectionStrategy = None,
        include_lm_response: bool = False,
        snap_indent: int | None = None,
        cache_lm: bool = True,
    ) -> None:
        super().__init__(model, include_lm_response)
        self.prompt_renderer = prompt_renderer
        self.response_parser = response_parser
        self.snap_indent = snap_indent
        self.cache_lm = cache_lm
        # if isinstance(model, OpenAIPredictor) and include_lm_response:
        #    msg = (
        #        "OpenAI models have changed to not support echoing/logprobs and this"
        #        " now needs to be handled. Currently not implemented."
        #    )
        #    raise NotImplementedError(
        #        msg,
        #    )
        if prompting_test_case_strategy is not None:
            msg = "Prompting test case strategy not implemented refactored"
            raise NotImplementedError(
                msg,
            )

    def solve(
        self,
        problem: CodeProblem,
        randomizing_seed: int | None = None,
    ) -> LmCodeSolution:
        start = timer()
        render = self.prompt_renderer.render(
            problem=problem,
            lm=self.model,
            prompt_seed=None,
        )
        if isinstance(render, LmPromptRenderSingleEdit):
            if render.changed_element is None:
                msg = "Insert element cannot be None"
                raise ValueError(msg)
            if not isinstance(render.changed_element, MarkElement):
                msg = (
                    "Render error. Expected changed element to be a MarkElement. Got: "
                    + str(type(render.changed_element))
                )
                raise ValueError(msg)
            if not isinstance(render.changed_element.verb, StsPlaceTransforming):
                msg = (
                    "Render error. Expected changed element verb to be "
                    "a verb StsPlaceTransforming. Got: "
                    + str(type(render.changed_element.verb))
                )
                raise ValueError(msg)
        logprobs = render.prompt.logprobs
        if self.include_lm_response and (logprobs is None or logprobs == 0):
            logprobs = 1
        render = dataclasses.replace(
            render,
            prompt=dataclasses.replace(
                render.prompt,
                temperature=0,
                max_tokens=1000,
                cache=self.cache_lm,
                logprobs=logprobs,
                echo=False,  # Broken now with logprobs
            ),
        )

        resp = self.model.predict(render.prompt)
        solve_steps = self.response_parser.parse(render, resp, problem)
        if self.snap_indent:
            # TODO: move this into response parser?
            solve_steps = (
                dataclasses.replace(
                    solve_steps[0],
                    value=snap_line_indent(solve_steps[0].value, self.snap_indent),
                ),
            )
        end = timer()

        # Uncertainty stuff (TODO: refector)
        # if render.prompt.is_text_a_chat():
        #    mod_text = LmChatDialog(
        #        [
        #            *render.prompt.text,
        #            LmChatTurn(
        #                ChatGptRoles.assistant,
        #                solution.solve_steps[0].value,
        #            ),
        #        ],
        #    )
        # else:
        #    mod_text = render.code_prefix + solution.solve_steps[0].value
        # self._augment_solution_with_uncertainty_hooks(solution, resp, mod_text)
        lang_spec = PythonLangSpec()
        for step in solve_steps:
            if step.path is not None:
                lang_spec = problem.get_lang_spec_for_path(step.path)
                break

        return LmCodeSolution(
            solve_steps,
            problem,
            lang_spec=lang_spec,
            time=end - start,
            prompt=render.prompt,
            lm_prediction=resp if self.include_lm_response else None,
        )


class LmUncertaintyHook:
    def prepare(
        self,
        lm_prediction: LmPrediction,
        model: LmPredictor,
        solution: CodeSolution,
        modified_full_text: str,  # The prompt + completion (but with
        # possible changes to completion like truncation)
    ) -> DelayedUncertaintyEstimate:
        pass

    def get_name(self) -> str:
        return type(self).__name__

    def __eq__(self, other):
        return type(self).__name__ == type(other).__name__

    def __hash__(self):
        return hash(type(self).__name__)


class LmUncertaintyProbTotal(LmUncertaintyHook):
    def prepare(
        self,
        lm_prediction: LmPrediction,
        model: LmPredictor,
        solution: CodeSolution,
        modified_full_text: str,
    ) -> DelayedUncertaintyEstimate:
        start, end = estimate_best_overlap_span(
            lm_prediction.completion_text,
            larger_tokens=lm_prediction.completion_tokens,
            mod_string=solution.solve_steps[0].value,
            tokenizer_func=model.tokenize,
        )
        return PrecomputedUncertaintyEstimate[ProbabilityIsCorrect](
            ProbabilityIsCorrect(
                name_key="token_prob_total",
                value=np.exp(
                    np.sum(lm_prediction.completion_logprobs[start:end]),
                ),
            ),
        )

    def get_name(self) -> str:
        return "token_prob_total"


class LmUncertaintyProbAverage(LmUncertaintyHook):
    def prepare(
        self,
        lm_prediction: LmPrediction,
        model: LmPredictor,
        solution: CodeSolution,
        modified_full_text: str,
    ) -> DelayedUncertaintyEstimate:
        start, end = estimate_best_overlap_span(
            lm_prediction.completion_text,
            larger_tokens=lm_prediction.completion_tokens,
            mod_string=solution.solve_steps[0].value,
            tokenizer_func=model.tokenize,
        )
        out_estimate = ProbabilityIsCorrect(
            name_key="token_prob_avg",
            value=float(
                np.mean(np.exp(lm_prediction.completion_logprobs[start:end])),
            ),
        )
        return PrecomputedUncertaintyEstimate[ProbabilityIsCorrect](
            out_estimate,
        )

    def get_name(self) -> str:
        return "token_prob_avg"
