import os
from unittest.mock import MagicMock

import pytest
from lmwrapper.openai_wrapper import OpenAiModelNames, get_open_ai_lm
from lmwrapper.structs import LmPrediction

from synthegrator.code_problems import (
    CodeProblem,
    CodeSolution,
    DiscoveredTestsuite,
)
from synthegrator.code_solver import LmCodeSolverAutoRegressive
from synthegrator.environments import CodeExecutionEnvironment, ProjectWorkingDirectory
from synthegrator.lang_specs.lang_spec_java import JavaLangSpec
from synthegrator.memory_fs import ProjectDir
from synthegrator.problem_rendering import (
    render_single_edit_codeproblem_autoregressive,
)
from synthegrator.sandboxing import JAVA_DEFAULT_DOCKER_ENV
from synthegrator.solution_eval import evaluate_code_problem_execution
from synthegrator.transformation_spec import (
    SolveStep,
    StsInsert,
    StsSelectorCharSpan,
    StsSpecStatement,
    TransformationSpec,
)

IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"


def get_java_double_problem():
    content = (
        # fmt: off
        "import java.util.*;\n"
        "import java.lang.*;\n"
        "\n"
        "class Solution {\n"
        "    /**\n"
        "    Double the input number.\n"
        "    */\n"
        "    public int doubleNumber(int x) {\n"
        "}\n"
        # fmt: on
    )
    content_test = (
        # fmt: off
        "import org.junit.jupiter.api.Test;\n"
        "import static org.junit.jupiter.api.Assertions.assertEquals;\n"
        "\n"
        "class SolutionTest {\n"
        "\n"
        "    @Test\n"
        "    void testDouble() {\n"
        "        Solution solution = new Solution();\n"
        "\n"
        "        assertEquals(4, solution.doubleNumber(2));\n"
        "        assertEquals(0, solution.doubleNumber(0));\n"
        "        assertEquals(-6, solution.doubleNumber(-3));\n"
        "        assertEquals(20, solution.doubleNumber(10));\n"
        "    }\n"
        "}\n"
        # fmt: on
    )
    select_char = len(content) - len("}\n")
    problem = CodeProblem(
        instructions="Complete this method",
        working_directory=ProjectWorkingDirectory(
            files=ProjectDir.construct_with_files(
                [
                    ("src/main/java/Solution.java", content),
                    ("src/test/java/SolutionTest.java", content_test),
                ],
            ),
        ),
        transformation_spec=TransformationSpec(
            statements=(
                StsSpecStatement(
                    "src/main/java/Solution.java",
                    StsSelectorCharSpan(select_char, select_char),
                    StsInsert(
                        lang_spec_name="java",
                        set_base_space_indent=8,
                        stop_at_block_end=True,
                        set_ending_text="}",
                    ),
                ),
            ),
        ),
        environment=CodeExecutionEnvironment(JAVA_DEFAULT_DOCKER_ENV),
        instructions_are_essential=False,
        test_cases=[
            DiscoveredTestsuite(
                cmds=[
                    "mvn test",
                    "cat target/surefire-reports/TEST-SolutionTest.xml",
                ],
            ),
        ],
    )
    render = render_single_edit_codeproblem_autoregressive(problem)
    assert render.prompt.text == (
        # fmt: off
        "import java.util.*;\n"
        "import java.lang.*;\n"
        "\n"
        "class Solution {\n"
        "    /**\n"
        "    Double the input number.\n"
        "    */\n"
        "    public int doubleNumber(int x) {\n"
    )
    problem.known_solutions.append(
        CodeSolution(
            solve_steps=(
                SolveStep(
                    path=render.path,
                    mark_id=render.changed_element.mark_id,
                    value=" " * 8 + "return x * 2;\n    }",
                ),
            ),
            problem=problem,
            lang_spec=JavaLangSpec(),
        ),
    )
    return problem


def test_simple_method_problem_known_solution():
    problem = get_java_double_problem()
    eval = evaluate_code_problem_execution(problem, problem.known_solutions[0])
    assert eval.all_passed()


@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="No API key")
@pytest.mark.skip(reason="Need to sort to through include lm things")
def test_simple_method_problem_gpt():
    problem = get_java_double_problem()
    lm = get_open_ai_lm(
        OpenAiModelNames.gpt_3_5_turbo_instruct,
    )
    solver = LmCodeSolverAutoRegressive(lm, cache_lm=False, include_lm_response=True)
    solution = solver.solve(problem)
    print("prompt", solution.lm_prediction.prompt)
    print("completion text", solution.lm_prediction.completion_text)
    print(solution.solve_steps[0])
    eval = evaluate_code_problem_execution(problem, solution)
    assert eval.all_passed()
    # Disable lm resp
    solver = LmCodeSolverAutoRegressive(lm, cache_lm=False, include_lm_response=False)
    solution = solver.solve(problem)
    assert solution.lm_prediction is None


@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="No API key")
def test_simple_method_problem_gpt_chat():
    problem = get_java_double_problem()
    lm = get_open_ai_lm(OpenAiModelNames.gpt_3_5_turbo)
    solver = LmCodeSolverAutoRegressive(lm)
    solution = solver.solve(problem)
    print(solution.solve_steps[0])
    eval = evaluate_code_problem_execution(problem, solution)
    assert eval.all_passed()


@pytest.mark.parametrize(
    "gen_return",
    [
        "    return x * 2;\n    }\n",
        "    return x * 2;\n    }\n} class ExtraCode() {}\n int x =",
        "    return x * 2;\n",
    ],
)
def test_simple_method_problem_mocked(gen_return: str):
    problem = get_java_double_problem()
    lm = get_open_ai_lm(OpenAiModelNames.gpt_3_5_turbo)
    lm.predict = MagicMock(
        return_value=LmPrediction(
            gen_return,
            None,
            None,
        ),
    )
    solver = LmCodeSolverAutoRegressive(lm)
    solution = solver.solve(problem)
    print(solution.solve_steps[0])
    assert "ExtraCode" not in solution.solve_steps[0].value
    eval = evaluate_code_problem_execution(problem, solution)
    assert eval.all_passed()


@pytest.mark.parametrize(
    "gen_return",
    [
        "    return x * 2;\n    }\n",
        "    return x * 2;\n    }",
        "    return x * 2;\n    }\n} class ExtraCode() {}",
        "    return x * 2;\n    }\n void ExtraCode() {}\n}",
    ],
)
def test_truncate_func(gen_return):
    lang_spec = JavaLangSpec()
    out = lang_spec.truncate_function_completion(gen_return)
    assert out == "    return x * 2;\n    }"


@pytest.mark.parametrize(
    "gen_return",
    [
        "    return x * 2;}",
        "    return x * 2;}\n}",
    ],
)
def test_truncate_func_2(gen_return):
    lang_spec = JavaLangSpec()
    out = lang_spec.truncate_function_completion(gen_return)
    assert out == "    return x * 2;}"


@pytest.mark.parametrize(
    "gen_return",
    [
        "        int x = 3;\n        return x * 2;\n    } class ExtraCode() {}",
    ],
)
def test_truncate_func_3(gen_return):
    lang_spec = JavaLangSpec()
    out = lang_spec.truncate_function_completion(gen_return)
    assert out == "        int x = 3;\n        return x * 2;\n    }"


def test_truncate_func_loop():
    lang_spec = JavaLangSpec()
    completion = (
        # fmt: off
        "        int x = 3;\n"
        "        for (int i = 0; i < 10; i++) {\n"
        "            x += i;\n"
        "        }\n"
        "        return x * 2;\n"
        "    }\n"
        "}\n"
        "class ExtraCode() {}"
        # fmt: on
    )
    out = lang_spec.truncate_function_completion(completion)
    assert out == "\n".join(completion.splitlines()[:-2])


def test_truncate_func_inner_func():
    lang_spec = JavaLangSpec()
    completion = (
        # fmt: off
        "        int x = 3;\n"
        "        void bar(int x) {\n"
        "            x += 1;\n"
        "        }\n"
        "        return x * 2;\n"
        "    }\n"
        "}\n"
        "class ExtraCode() {}"
        # fmt: on
    )
    out = lang_spec.truncate_function_completion(completion)
    assert out == "\n".join(completion.splitlines()[:-2])


# def test_simple_method_problem_codet5():
#    problem = _get_problem()
#    lm = get_huggingface_lm("Salesforce/codet5p-220m")
#    solver = LmCodeSolverAutoRegressive(lm)
#    solution = solver.solve(problem, lang_spec=JavaLangSpec())
#    print(solution.solve_steps[0])
#    eval = evaluate_code_problem(problem, solution)
#    assert eval.all_passed()
