import json
import math
import os
from unittest.mock import MagicMock

import pytest
from lmwrapper.huggingface_wrapper import get_huggingface_lm
from lmwrapper.openai_wrapper import OpenAiModelNames, get_open_ai_lm
from lmwrapper.structs import ChatGptRoles, LmPrediction

from synthegrator.code_problems import CodeProblem, CodeSolution
from synthegrator.code_solver import LmCodeSolverAutoRegressive
from synthegrator.environments import ProjectWorkingDirectory
from synthegrator.few_shotting import FewShotConfig, FewShotLibrary
from synthegrator.lang_specs.lang_spec_python import PythonLangSpec
from synthegrator.lm_few_shotting_tools import FewShotRendererWrapper
from synthegrator.memory_fs import ProjectDir
from synthegrator.prompt_renderer_questions import (
    CallbackResponseParser,
    QuestionPromptRenderer,
)
from synthegrator.transformation_spec import (
    SolveStep,
    StsCodeClassificationQuestion,
    StsCodeQuestion,
    StsSelectorFullFile,
    StsSpecStatement,
    TransformationSpec,
)

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


def basic_question_renderer():
    code = "if __name__ == '__main__':\n    print(2 + 3)\n"
    wd = ProjectWorkingDirectory(ProjectDir.construct_with_files([("main.py", code)]))
    spec = TransformationSpec(
        statements=(
            StsSpecStatement(
                "main.py",
                StsSelectorFullFile(),
                StsCodeQuestion(
                    "What will the code print? Please give only the value that appears on standard out with no extra text.",
                ),
            ),
        ),
    )
    return CodeProblem(
        working_directory=wd,
        transformation_spec=spec,
    )


def basic_classification_question_problem():
    code = "if __name__ == '__main__':\n    print(2 + 3)\n"
    wd = ProjectWorkingDirectory(ProjectDir.construct_with_files([("main.py", code)]))
    spec = TransformationSpec(
        statements=(
            StsSpecStatement(
                "main.py",
                StsSelectorFullFile(),
                StsCodeClassificationQuestion(
                    "Which programming language is used?",
                    ("python", "java", "c++"),
                ),
            ),
        ),
    )
    problem = CodeProblem(
        working_directory=wd,
        transformation_spec=spec,
    )
    problem.known_solutions = [
        CodeSolution(
            solve_steps=[
                SolveStep(
                    "main.py",
                    spec.statements[0].static_mark_id,
                    "python",
                ),
            ],
            problem=problem,
            lang_spec=PythonLangSpec(),
        ),
    ]
    return problem


def basic_classification_question_problem_2():
    code = "if __name__ == '__main__':\n    print(2 + 3)\n"
    wd = ProjectWorkingDirectory(ProjectDir.construct_with_files([("main.py", code)]))
    spec = TransformationSpec(
        statements=(
            StsSpecStatement(
                "main.py",
                StsSelectorFullFile(),
                StsCodeClassificationQuestion(
                    "Which value will be printed?",
                    ("1", "5", "hello world", "2"),
                ),
            ),
        ),
    )
    problem = CodeProblem(
        working_directory=wd,
        transformation_spec=spec,
    )
    problem.known_solutions = [
        CodeSolution(
            solve_steps=[
                SolveStep(
                    "main.py",
                    spec.statements[0].static_mark_id,
                    "5",
                ),
            ],
            problem=problem,
            lang_spec=PythonLangSpec(),
        ),
    ]
    return problem


def test_simple():
    problem = basic_question_renderer()
    renderer = QuestionPromptRenderer()
    render = renderer.render(
        problem,
        MagicMock(),
    )
    assert (
        render.prompt.text == "```python\n"
        "if __name__ == '__main__':\n"
        "    print(2 + 3)\n"
        "```\n"
        "Question: What will the code print? Please give only the value that appears on standard out with no extra text.\n"
        "Answer:"
    )


@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="No api in github actions")
def test_simple_solve():
    problem = basic_question_renderer()
    renderer = QuestionPromptRenderer()
    parser = CallbackResponseParser()
    lm = get_open_ai_lm()
    solver = LmCodeSolverAutoRegressive(
        lm,
        prompt_renderer=renderer,
        response_parser=parser,
    )
    solution = solver.solve(problem)
    assert len(solution.solve_steps) == 1
    assert solution.solve_steps[0].path == "main.py"
    assert solution.solve_steps[0].value.strip() == "5"


def test_simple_solve_mocked():
    problem = basic_question_renderer()
    renderer = QuestionPromptRenderer()
    parser = CallbackResponseParser()
    lm = get_open_ai_lm(OpenAiModelNames.gpt_3_5_turbo_instruct)
    lm.predict = MagicMock(return_value=LmPrediction("5\n", None, None))
    solver = LmCodeSolverAutoRegressive(
        lm,
        prompt_renderer=renderer,
        response_parser=parser,
    )
    solution = solver.solve(problem)
    assert len(solution.solve_steps) == 1
    assert solution.solve_steps[0].path == "main.py"
    assert solution.solve_steps[0].value.strip() == "5"


def test_classication_render():
    problem = basic_classification_question_problem()
    renderer = QuestionPromptRenderer()
    render = renderer.render(
        problem,
        MagicMock(),
    )
    # fmt: off
    assert render.prompt.text == (
        "```python\n"
        "if __name__ == '__main__':\n"
        "    print(2 + 3)\n"
        "```\n"
        "Question: Which programming language is used? Answer with only the letter corresponding to the best answer.\n"
        "A) python\n"
        "B) java\n"
        "C) c++\n"
        "Answer:"
    )
    # fmt: on


def test_classication_render_few_shot():
    problem = basic_classification_question_problem()
    few_shot_config = FewShotConfig(
        FewShotLibrary([basic_classification_question_problem_2()]),
        num_examples=1,
        change_examples_between_solves=False,
    )
    renderer = QuestionPromptRenderer()
    parser = CallbackResponseParser()
    renderer = FewShotRendererWrapper(
        renderer,
        parser,
        few_shot_config,
    )
    lm = MagicMock()
    lm.is_chat_model = False
    render = renderer.render(
        problem,
        lm,
    )
    print(render.prompt.text)
    # fmt: off
    assert render.prompt.text == (
        "```python\n"
        "if __name__ == '__main__':\n"
        "    print(2 + 3)\n"
        "```\n"
        "Question: Which value will be printed? Answer with only the letter corresponding to the best answer.\n"
        "A) 1\n"
        "B) 5\n"
        "C) hello world\n"
        "D) 2\n"
        "Answer: B"
        "\n###\n"
        "```python\n"
        "if __name__ == '__main__':\n"
        "    print(2 + 3)\n"
        "```\n"
        "Question: Which programming language is used? Answer with only the letter corresponding to the best answer.\n"
        "A) python\n"
        "B) java\n"
        "C) c++\n"
        "Answer:"
    )
    # fmt: on


def test_classication_render_few_shot_chat():
    problem = basic_classification_question_problem()
    few_shot_config = FewShotConfig(
        FewShotLibrary([basic_classification_question_problem_2()]),
        num_examples=1,
        change_examples_between_solves=False,
    )
    renderer = QuestionPromptRenderer()
    parser = CallbackResponseParser()
    renderer = FewShotRendererWrapper(
        renderer,
        parser,
        few_shot_config,
    )
    lm = MagicMock()
    lm.is_chat_model = True
    render = renderer.render(
        problem,
        lm,
    )
    print(render.prompt.text)
    # fmt: off
    assert render.prompt.is_text_a_chat()
    assert len(render.prompt.text) == 3
    assert render.prompt.text[0].content == (
        "```python\n"
        "if __name__ == '__main__':\n"
        "    print(2 + 3)\n"
        "```\n"
        "Question: Which value will be printed? Answer with only the letter corresponding to the best answer.\n"
        "A) 1\n"
        "B) 5\n"
        "C) hello world\n"
        "D) 2\n"
        "Answer:"
    )
    assert render.prompt.text[0].role == ChatGptRoles.user
    assert render.prompt.text[1].content == "B"
    assert render.prompt.text[1].role == ChatGptRoles.assistant
    assert render.prompt.text[2].content == (
        "```python\n"
        "if __name__ == '__main__':\n"
        "    print(2 + 3)\n"
        "```\n"
        "Question: Which programming language is used? Answer with only the letter corresponding to the best answer.\n"
        "A) python\n"
        "B) java\n"
        "C) c++\n"
        "Answer:"
    )
    assert render.prompt.text[2].role == ChatGptRoles.user
    # fmt: on


def run_test_simple_classification_solve(lm, known_ans: dict[str, float] | None = None):
    problem = basic_classification_question_problem()
    renderer = QuestionPromptRenderer()
    parser = CallbackResponseParser()
    solver = LmCodeSolverAutoRegressive(
        lm,
        prompt_renderer=renderer,
        response_parser=parser,
    )
    solution = solver.solve(problem)
    assert len(solution.solve_steps) == 1
    assert solution.solve_steps[0].path == "main.py"
    prob_dict = json.loads(solution.solve_steps[0].value)
    assert len(prob_dict) == 3
    assert set(prob_dict.keys()) == {"python", "java", "c++"}
    assert prob_dict["python"] > 0.5
    if known_ans is not None:
        assert len(prob_dict) == len(known_ans)
        for k, v in known_ans.items():
            assert math.isclose(prob_dict[k], v, abs_tol=0.0001)


@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="No api in github actions")
def test_simple_classification_solve():
    lm = get_open_ai_lm()
    run_test_simple_classification_solve(lm)


@pytest.mark.skip(reason="need to merge hf top token")
def test_simple_classification_solve_hf():
    lm = get_huggingface_lm(
        "Salesforce/codegen-350M-multi",
        trust_remote_code=True,
    )
    run_test_simple_classification_solve(lm)


def test_simple_classification_solve_mocked():
    lm = get_open_ai_lm()
    mocked_pred = MagicMock()
    mocked_pred.completion_text = "A"
    mocked_pred.top_token_logprobs = [
        {
            "A": math.log(0.8),
        },
    ]
    lm.predict = MagicMock(return_value=mocked_pred)
    run_test_simple_classification_solve(
        lm,
        known_ans={"c++": 0.1, "python": 0.8, "java": 0.1},
    )


def test_simple_classification_solve_mocked2():
    lm = get_open_ai_lm()
    mocked_pred = MagicMock()
    mocked_pred.completion_text = "A"
    mocked_pred.top_token_logprobs = [
        {
            " a": math.log(0.8),
        },
    ]
    lm.predict = MagicMock(return_value=mocked_pred)
    run_test_simple_classification_solve(
        lm,
        known_ans={"c++": 0.1, "python": 0.8, "java": 0.1},
    )


def test_simple_classification_solve_mocked3():
    lm = get_open_ai_lm()
    mocked_pred = MagicMock()
    mocked_pred.completion_text = "A"
    mocked_pred.top_token_logprobs = [
        {
            "A": math.log(0.8),
            "B": math.log(0.15),
        },
    ]
    lm.predict = MagicMock(return_value=mocked_pred)
    run_test_simple_classification_solve(
        lm,
        known_ans={"c++": 0.05, "python": 0.8, "java": 0.15},
    )


def test_simple_classification_solve_mocked_norm_mass():
    lm = get_open_ai_lm()
    mocked_pred = MagicMock()
    mocked_pred.completion_text = "A"
    mocked_pred.top_token_logprobs = [
        {
            "A": math.log(0.4),
            "B": math.log(0.1),
            "C": math.log(0.00000000001),
        },
    ]
    lm.predict = MagicMock(return_value=mocked_pred)
    run_test_simple_classification_solve(
        lm,
        known_ans={"c++": 0.0, "python": 0.8, "java": 0.2},
    )


def test_simple_classification_agg():
    lm = get_open_ai_lm()
    mocked_pred = MagicMock()
    mocked_pred.completion_text = "A"
    mocked_pred.top_token_logprobs = [
        {
            "A": math.log(0.5),
            "B": math.log(0.1),
            "C": math.log(0.00000000001),
            " python": math.log(0.4),
            "sdf": math.log(0.4),
        },
    ]
    lm.predict = MagicMock(return_value=mocked_pred)
    run_test_simple_classification_solve(
        lm,
        known_ans={"c++": 0.0, "python": 0.9, "java": 0.1},
    )


def test_ans_to_solution():
    problem = basic_classification_question_problem()
    renderer = QuestionPromptRenderer(randomize_multiple_choice=False)
    parser = CallbackResponseParser()
    render = renderer.render(problem, lm=MagicMock())
    ans = parser.answer_for_solution(
        render,
        problem,
        solve_steps=[
            SolveStep(
                "main.py",
                problem.transformation_spec.statements[0].static_mark_id,
                "python",
            ),
        ],
    )
    assert ans == " A"
    ans = parser.answer_for_solution(
        render,
        problem,
        solve_steps=[
            SolveStep(
                "main.py",
                problem.transformation_spec.statements[0].static_mark_id,
                "java",
            ),
        ],
    )
    assert ans == " B"
