"""
Dataclasses and helper functions for defining a
problem want to synthesize code for.
"""

import inspect
from abc import ABC, abstractmethod
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, NamedTuple

import xxhash
from lmwrapper.structs import LmPrediction, LmPrompt

from synthegrator.environments import (
    CodeExecutionEnvironment,
    ProjectWorkingDirectory,
)
from synthegrator.lang_specs.lang_spec import (
    LangSpec,
)
from synthegrator.lang_specs.lang_spec_java import JavaLangSpec
from synthegrator.lang_specs.lang_spec_python import (
    PythonLangSpec,
    assert_spec_is_python,
)
from synthegrator.memory_fs import ProjectDir
from synthegrator.sandboxing import (
    Cmd,
    CmdExecResult,
    TestsuiteExecResult,
    convert_pytest_outputs_to_testsuite_exec_result,
)
from synthegrator.transformation_spec import (
    MarkElement,
    SolveStep,
    StsValueQuery,
    TransformationSpec,
)
from synthegrator.uncertainty_modeling import DelayedUncertaintyEstimate

if TYPE_CHECKING:
    # A seeming hack to be able to import the SolutionEvaluator during
    #   the typing of CodeProblem without a circular import. Not quite
    #   how good of a solution this will be.
    from synthegrator.solution_eval import SolutionEvaluator


class TestCaseFormattedValue(NamedTuple):
    call_expression: str
    setup: str


@dataclass
class TestCase(ABC):
    __test__ = False
    is_hidden_test: bool
    """In a coding problem some test case might be hidden from the solver"""
    fail_message: str

    @abstractmethod
    def get_hash(self) -> str:
        hash_calc = xxhash.xxh64()
        if self.fail_message:
            hash_calc.update(self.fail_message)
        hash_calc.update(str(self.is_hidden_test))
        return hash_calc.hexdigest()

    def __post_init__(self):
        pass

    @abstractmethod
    def result_is_success(self, output: CmdExecResult) -> bool:
        # Deprecated with docker running
        raise NotImplementedError

    @abstractmethod
    def format_for_test_framework(
        self,
        lang_spec: LangSpec,
        override_test_id: str | None = None,
    ) -> str:
        """Formats for running on a test framework like pytest."""
        raise NotImplementedError


def _pytest_with_body(test_body: str, test_id, fixtures=None) -> str:
    if not fixtures:
        fixtures = []
    # Make a random id for the test
    test_body = "\n".join("    " + line for line in test_body.split("\n"))
    return f"def test_{test_id}({', '.join(fixtures)}):\n{test_body}"


@dataclass
class TestCaseExpressionIsEq(TestCase):
    __test__ = False
    """
    Checks whether calling expression equals a value. Currently implemented
    only for python
    """

    input_expression: str
    output: Any
    lang_spec: LangSpec

    def get_hash(self) -> str:
        hash_calc = xxhash.xxh64()
        hash_calc.update(super().get_hash())
        hash_calc.update(self.input_expression)
        hash_calc.update(str(self.output))
        hash_calc.update(self.lang_spec.get_lang_md_name())
        return hash_calc.hexdigest()

    def __post_init__(self):
        assert_spec_is_python(self.lang_spec)

    def result_is_success(self, output: CmdExecResult) -> bool:
        return output.output == self.output

    def format_for_test_framework(
        self,
        lang_spec: LangSpec,
        override_test_id: str | None = None,
    ) -> str:
        assert_spec_is_python(lang_spec)
        return _pytest_with_body(
            f"assert {self.input_expression} == {self.output}",
            override_test_id or self.get_hash(),
        )


@dataclass
class TestCaseMethodCallIsEq(TestCase):
    __test__ = False
    method_name: str
    test_input: tuple[Any, ...]
    output: Any

    def get_hash(self) -> str:
        hash_calc = xxhash.xxh64()
        hash_calc.update(super().get_hash())
        hash_calc.update(str(self.test_input))
        hash_calc.update(str(self.output))
        return hash_calc.hexdigest()

    def _format_method_call(
        self,
        lang_spec: LangSpec,
    ) -> str:
        assert_spec_is_python(lang_spec)
        return f"{self.method_name}({', '.join(map(str, self.test_input))})"

    def result_is_success(self, output: CmdExecResult) -> bool:
        return output.output == self.output

    def format_for_test_framework(
        self,
        lang_spec: LangSpec,
        override_test_id=None,
    ):
        assert_spec_is_python(lang_spec)
        return _pytest_with_body(
            f"assert {self._format_method_call(lang_spec)} ==" f" {self.output}",
            override_test_id or self.get_hash(),
        )


@dataclass
class TestCaseArbitraryMethod(TestCase):
    method_body: str
    lang_spec: LangSpec
    fixtures: tuple[str, ...] = ()

    def get_hash(self) -> str:
        hash_calc = xxhash.xxh64()
        hash_calc.update(super().get_hash())
        hash_calc.update(str(self.method_body))
        hash_calc.update(
            str(self.lang_spec.get_lang_md_name()),
        )  # TODO: add other info from spec
        hash_calc.update(str(self.fixtures))
        return hash_calc.hexdigest()

    def result_is_success(self, output: CmdExecResult) -> bool:
        return output.completed

    def format_for_test_framework(
        self,
        lang_spec: LangSpec,
        override_test_id: str | None = None,
    ) -> str:
        assert_spec_is_python(lang_spec)
        return _pytest_with_body(
            self.method_body,
            override_test_id or self.get_hash(),
            self.fixtures,
        )


class DiscoveredTestsuite(TestCase):
    """
    A test case for when running a suite that already exists in the environment.
    It specifies the necessary commands to get the valid suite xml.
    """

    def __init__(
        self,
        cmds: Sequence[Cmd | str],
        cmds_results_to_suite: (
            Callable[[Sequence[CmdExecResult]], TestsuiteExecResult] | None
        ) = None,
    ):
        self.is_hidden_test = True
        self.fail_message = f"Fail discovered test {cmds}"
        self.cmds = Cmd.from_cmds_str_mix(cmds)
        self.cmds_results_to_suite = cmds_results_to_suite
        if self.cmds_results_to_suite is None:
            if len(self.cmds) < 2:
                msg = "Cmds length invalid. Failed run?"
                raise Exception(msg)

            if "cat " not in self.cmds[-1].cmd:
                msg = "Cat command not executed. Failed run?"
                raise Exception(msg)

            self.cmds_results_to_suite = self._default_cmds_results_suite

    def get_hash(self) -> str:
        hash_calc = xxhash.xxh64()
        hash_calc.update(super().get_hash())
        hash_calc.update(str(self.cmds))
        hash_calc.update(inspect.getsource(self.cmds_results_to_suite))
        return hash_calc.hexdigest()

    @staticmethod
    def _default_cmds_results_suite(
        cmds_results: Sequence[CmdExecResult],
    ) -> TestsuiteExecResult:
        return convert_pytest_outputs_to_testsuite_exec_result(
            cmds_results[-2],
            cmds_results[-1],
        )

    def result_is_success(self, output: CmdExecResult) -> bool:
        return output.completed

    def format_for_test_framework(
        self,
        lang_spec: LangSpec,
        override_test_id: str | None = None,
    ) -> str:
        raise NotImplementedError


@dataclass
class CodeSolution:
    solve_steps: Sequence[SolveStep]
    problem: "CodeProblem"
    lang_spec: LangSpec = None  # TODO remove this eventually
    time: float = None
    """Time in seconds of how long the solver took to produce the solution."""

    def __post_init__(self):
        if not isinstance(self.solve_steps, Sequence):
            msg = "Solve steps must be a Sequence"
            raise TypeError(msg)
        if not all(isinstance(step, SolveStep) for step in self.solve_steps):
            msg = "All solve steps must be a SolveStep"
            raise TypeError(msg)

    def apply(self) -> ProjectDir:
        from synthegrator.solution_eval import apply_solution

        return apply_solution(self)

    def pull_value_answers(self) -> list[tuple[MarkElement[StsValueQuery], str]]:
        from synthegrator.solution_eval import pull_value_answers_from_solution

        return pull_value_answers_from_solution(self)

    def dict_serialize(self) -> dict:
        out = {
            "solve_steps": sorted(
                [step.dict_serialize() for step in self.solve_steps],
                key=lambda x: (x["path"], x["mark_id"]),
            ),
            "problem": self.problem.dict_serialize(),
            "time": self.time,
            "lang_spec": self.lang_spec.dict_serialize() if self.lang_spec else None,
            "number_of_solve_steps": len(self.solve_steps),
        }
        if len(self.solve_steps) == 1:
            out["only_solve_value"] = self.solve_steps[0].value
            out["only_solve_path"] = self.solve_steps[0].path
        return out


@dataclass
class LmCodeSolution(CodeSolution):
    uncertainty_estimates: list[DelayedUncertaintyEstimate] = None
    prompt: LmPrompt | None = None
    lm_prediction: LmPrediction | None = None

    def __post_init__(self):
        super().__post_init__()
        if self.uncertainty_estimates is None:
            self.uncertainty_estimates = []

    def dict_serialize(
        self,
        remove_redundant_prompt: bool = True,
    ) -> dict:
        out = super().dict_serialize()
        # out["uncertainty_estimates"] = [
        #    est.dict_serialize()
        #    for est in self.uncertainty_estimates
        # ]
        if self.prompt:
            out["prompt"] = self.prompt.dict_serialize()
        if self.lm_prediction:
            out["lm_prediction"] = self.lm_prediction.dict_serialize()
            if remove_redundant_prompt and self.prompt:
                # There will also be the prompt in the response which
                # is redundant
                del out["lm_prediction"]["prompt"]
        return out


DEFAULT_CODE_ENV = CodeExecutionEnvironment()
DEFAULT_LANG_SPECS = (
    (".py", PythonLangSpec()),
    (".java", JavaLangSpec()),
)


@dataclass
class CodeProblem:
    instructions: str = None
    """The main instructions describing what we want a solver to do in this problem"""
    working_directory: ProjectWorkingDirectory = None
    """The working directory that contains the relevant files for the problem.
    Optionally, the files can exist already in the execution environment, but then
    they are opaque to the solvers."""
    transformation_spec: TransformationSpec = None
    """Outlines what a solver is and isn't allowed or expected to edit"""
    test_cases: list[TestCase] = field(default_factory=list)
    known_solutions: list[CodeSolution] = field(default_factory=list)
    environment: CodeExecutionEnvironment = DEFAULT_CODE_ENV
    dataset_name: str = None
    """Labeling which dataset this problem came from"""
    problem_id: str = None
    """Some sort of identifier within the dataset. The pair of (dataset_name, problem_id)
    should be unique."""
    instructions_are_essential: bool = True
    """Whether the instructions include key information for solving the problem.
    For some method or line completion tasks or when the task is embedded into
    documents in the code, the instructions might not be essential. However
    if you are generating code with no documentation or other guide on the task,
    then the instructions might be essential. This can influence the behaviour
    of autoregressive models which might perform better when the instructions
    are naturally integrated into the code."""
    extension_to_lang_specs: tuple[tuple[str, LangSpec], ...] = DEFAULT_LANG_SPECS
    """A tuple that contains all the language specs that are used in the problem.
    Each element is a tuple of (extension, lang_spec). Paths that end with the
    extension will use the associate language spec"""
    preferred_solution_evaluator: "SolutionEvaluator" = None
    """Preferred solution evaluator for the problem. Can be used in case certain
     problems use metrics. If set to `None` then a SolutionEvaluator.
    will be used if there are test cases. Otherwise it will be None."""
    instructable_metadata: dict[str, Any] = None
    """NOT IMPLEMENTED. Intended for things that certain solvers might using
    while solving, but might do it in their own specially formatted way. Some
    examples might include expected commit message or something. Several of
    the standard metadata (eg, a commit message, issue text and discussion, etc)
    are intended to be standardized with certain keys (TODO)."""
    noninstructable_metadata: dict[str, Any] = None
    """Intended for extra problem metadata (for example which split a dataset example
    came from or other info) which is not intended for use by the solver but
    just for processing after the fact."""
    state_info = None
    """NOT IMPLEMENTED. intended to package extra external information available when
    solving the problem (like traces, error messages, etc.). This still needs
    to be further designed out and is not implemented."""
    past_solve_context: list[CodeSolution] = None
    """Context from some iterations trying to solve this"""

    def __post_init__(self):
        if self.state_info is not None:
            raise NotImplementedError
        if not isinstance(self.known_solutions, list):
            msg = "known_solutions must be a list"
            raise TypeError(msg)
        if not all(
            isinstance(s, CodeSolution) for s in self.known_solutions
        ) or isinstance(self.known_solutions, CodeSolution):
            msg = "known_solutions must be a list of CodeSolution"
            raise TypeError(msg)
        if self.preferred_solution_evaluator is None and (
            self.test_cases
            or len(list(self.transformation_spec.classification_question_statements()))
            > 0
        ):
            from synthegrator.solution_eval import SolutionEvaluator

            self.preferred_solution_evaluator = SolutionEvaluator.default_instance()

    def get_lang_spec_for_path(self, path: str) -> LangSpec | None:
        for ext, lang_spec in self.extension_to_lang_specs:
            if path.endswith(ext):
                return lang_spec
        return None

    def get_non_hidden_test_cases(self) -> list[TestCase]:
        return [t for t in self.test_cases if not t.is_hidden_test]

    def dict_serialize(self, just_serialize_identifiers: bool = True) -> dict:
        out = {
            "dataset_name": self.dataset_name,
            "problem_id": self.problem_id,
        }
        if not just_serialize_identifiers:
            raise NotImplementedError
        return out

    def __eq__(self, other):
        if not isinstance(other, CodeProblem):
            return False
        # We basically compare everything except:
        #   - known_solutions since they might reference back to this problem
        return (
            self.problem_id == other.problem_id
            and self.instructions == other.instructions
            and self.working_directory == other.working_directory
            and self.transformation_spec == other.transformation_spec
            and self.test_cases == other.test_cases
            and self.environment == other.environment
            and self.dataset_name == other.dataset_name
            and self.instructions_are_essential == other.instructions_are_essential
            and self.extension_to_lang_specs == other.extension_to_lang_specs
            and self.past_solve_context == other.past_solve_context
        )
