# TODO(christian): put annotations correctly.

import abc
import json
from typing import Any, TypeVar

from pydantic import BaseModel

from benchmark.tau_bench.model_utils.api.datapoint import (
    BinaryClassifyDatapoint,
    ClassifyDatapoint,
    Datapoint,
    GenerateDatapoint,
    ParseDatapoint,
    ParseForceDatapoint,
    ScoreDatapoint,
)
from benchmark.tau_bench.model_utils.api.types import PartialObj
from benchmark.tau_bench.model_utils.model.exception import ModelError
from benchmark.tau_bench.model_utils.model.general_model import GeneralModel
from benchmark.tau_bench.model_utils.model.utils import (
    add_md_close_tag,
    approx_num_tokens,
    display_choices,
    json_response_to_obj_or_partial_obj,
    optionalize_type,
    parse_json_or_json_markdown,
    try_classify_recover,
    type_to_json_schema_string,
)

T = TypeVar("T", bound=BaseModel)


class Score(BaseModel):
    score: int


class Classification(BaseModel):
    classification: str


def task_prompt(task: str, text: str) -> str:
    return f"# Task\n{task}\n\n{text}"


def force_json_prompt(text: str, with_prefix: bool = False) -> str:
    suffix = (
        'For example:\nassistant:```json\n{"key": "value"}\n```'
        if not with_prefix
        else "\n\n```json\n"
    )
    return f"{text}\n\nThe result should be a valid JSON object in a markdown block only. {suffix}"


def build_score_state(
    instruction: str,
    text: str,
    min: int,
    max: int,
    examples: list[ScoreDatapoint] | None = None,
) -> str:
    def display_sample(
        instr: str, t: str, min: int, max: int, response: int | None = None
    ) -> str:
        p = task_prompt(
            task='Score the following text with the provided instruction and range as an integer value in valid JSON:\n{"score": number}',
            text=force_json_prompt(
                f"Instruction:\n{instr}\n\nText:\n{t}\n\nRange:\n[{min}, {max}]",
                with_prefix=True,
            ),
        )
        if response is not None:
            # the json markdown block is opened in the prompt
            return f'{p}\n{{"score": {response}}}\n```'
        return p

    p = (
        "\n\n".join(
            [
                display_sample(ex.instruction, ex.text, min, max, ex.response)
                for ex in examples
            ]
        )
        if examples is not None
        else ""
    )
    return f"{p}\n\n{display_sample(instr=instruction, t=text, min=min, max=max)}"


def build_parse_force_state(
    instruction: str,
    typ: type[T] | dict[str, Any],
    text: str | None = None,
    examples: list[ParseForceDatapoint] | None = None,
) -> str:
    def display_sample(
        instr: str,
        t: str,
        ty: type[T] | dict[str, Any],
        response: T | dict[str, Any] | None = None,
    ) -> str:
        if isinstance(ty, dict):
            json_schema_string = json.dumps(ty)
        else:
            json_schema_string = type_to_json_schema_string(ty)
        text_insert = "" if t is None else f"\n\nText:\n{t}"
        input_text = force_json_prompt(
            text=f"Instruction:\n{instr}{text_insert}\n\nSchema:\n{json_schema_string}",
            with_prefix=True,
        )
        if response is not None:
            if isinstance(response, dict):
                response_display = json.dumps(response)
            else:
                response_display = response.model_dump_json()
            # the json markdown block is opened in the prompt
            return f"{input_text}\n{response_display}\n```"
        return input_text

    p = (
        "".join(
            [
                display_sample(
                    instr=ex.instruction,
                    t=ex.text,
                    ty=ex.typ,
                    response=ex.response,
                )
                for ex in examples
            ]
        )
        + "\n\n"
        if examples is not None and len(examples) > 0
        else ""
    )
    p += display_sample(instr=instruction, t=text, ty=typ)
    return task_prompt(
        task="Generate an object with the provided instruction, text, and schema.",
        text=p,
    )


def build_parse_state(
    text: str,
    typ: type[T] | dict[str, Any],
    examples: list[ParseDatapoint] | None = None,
) -> str:
    instruction = "Parse the following text with the provided JSON schema."

    def display_sample(
        t: str,
        ty: type[T] | dict[str, Any],
        response: T | PartialObj | dict[str, Any] | None = None,
    ) -> str:
        if isinstance(ty, dict):
            json_schema_string = json.dumps(ty)
        else:
            optionalized_typ = optionalize_type(ty)
            json_schema_string = type_to_json_schema_string(optionalized_typ)
        # instruction is repeated to emphasize the task
        prompt = task_prompt(
            task=instruction,
            text=force_json_prompt(
                f"Text:\n{t}\n\nSchema:\n{json_schema_string}", with_prefix=True
            ),
        )
        if response is None:
            return prompt
        if isinstance(response, dict):
            response_display = json.dumps(response)
        else:
            response_display = response.model_dump_json()
        # the json markdown block is opened in the prompt
        json_response = f"{response_display}\n```"
        return f"{prompt}\n{json_response}"

    p = ""
    if examples is not None and len(examples) > 0:
        p = "\n\n".join(
            [
                display_sample(t=ex.text, ty=ex.typ, response=ex.response)
                for ex in examples
            ]
        )
    return f"{p}\n\n{display_sample(t=text, ty=typ)}"


def build_classify_state(
    instruction: str,
    text: str,
    options: list[str],
    examples: list[ClassifyDatapoint] | None = None,
) -> tuple[str, dict[str, int]]:
    def display_sample(
        instr: str, t: str, opts: list[str], response: int | None = None
    ) -> str | tuple[str, dict[str, int]]:
        choices_display, decode_map = display_choices(opts)
        input_text = force_json_prompt(
            f"Instruction:\n{instr}\n\nText:\n{t}\n\nChoices:\n{choices_display}",
            with_prefix=True,
        )
        prompt = task_prompt(task=instr, text=input_text)
        if response is not None:
            label = None
            for k, v in decode_map.items():
                if v == response:
                    label = k
                    break
            assert label is not None
            # the json markdown block is opened in the prompt
            json_display = f'{{"classification": "{label}"}}\n```'
            return f"{prompt}\n{json_display}"
        return prompt, decode_map

    p = 'Classify the following text with the provided instruction and choices. To classify, provide the key of the choice:\n{"classification": string}\n\nFor example, if the correct choice is \'Z. description of choice Z\', then provide \'Z\' as the classification as valid JSON:\n```json\n{"classification": "Z"}\n```'
    if examples is not None and len(examples) > 0:
        example_displays = "\n\n".join(
            [
                display_sample(
                    instr=ex.instruction,
                    t=ex.text,
                    opts=ex.options,
                    response=ex.response,
                )
                for ex in examples
            ]
        )
        p += f"\n\n{example_displays}"
    prompt, decode_map = display_sample(instr=instruction, t=text, opts=options)
    return f"{p}\n\n{prompt}", decode_map


def build_generate_state(
    instruction: str,
    text: str,
    examples: list[GenerateDatapoint] | None = None,
) -> str:
    def display_sample(instr: str, t: str, response: str | None = None) -> str:
        prompt = task_prompt(task=instr, text=t)
        if response is not None:
            return f"{prompt}\n\nText: {response}"
        return prompt

    prompt = (
        "\n\n".join([display_sample(ex.instruction, ex.text) for ex in examples])
        + "\n\n"
        if examples is not None and len(examples) > 0
        else ""
    )
    return f"{prompt}\n\n{display_sample(instruction, text)}\n\nText:"


class CompletionModel(GeneralModel):
    @abc.abstractmethod
    def generate_from_prompt(
        self, prompt: str, temperature: float | None = None
    ) -> str:
        raise NotImplementedError

    @abc.abstractmethod
    def parse_force_from_prompt(
        self,
        prompt: str,
        typ: BaseModel | dict[str, Any],
        temperature: float | None = None,
    ) -> dict[str, Any]:
        raise NotImplementedError

    def handle_parse_force_response(self, prompt: str, content: str) -> dict[str, Any]:
        try:
            return parse_json_or_json_markdown(content)
        except (json.decoder.JSONDecodeError, ValueError) as e:
            raise ModelError(
                short_message=f"Failed to decode JSON: {content}",
                prompt=prompt,
                response=content,
            ) from e

    def _handle_classify_response(
        self, res: dict[str, int], decode_map: dict[str, int]
    ) -> int:
        if "classification" not in res:
            raise ModelError(f"Invalid response from model: {res}")
        choice = res["classification"]
        if choice not in decode_map:
            key = try_classify_recover(s=choice, decode_map=decode_map)
            if key is not None:
                return decode_map[key]
            raise ModelError(f"Invalid choice: {choice}")
        return decode_map[choice]

    def classify(
        self,
        instruction: str,
        text: str,
        options: list[str],
        examples: list[ClassifyDatapoint] | None = None,
        temperature: float | None = None,
    ) -> int:
        prompt, decode_map = build_classify_state(
            instruction, text, options, examples=examples
        )
        res = self.parse_force_from_prompt(
            prompt, typ=Classification, temperature=temperature
        )
        return self._handle_classify_response(res, decode_map)

    def parse(
        self,
        text: str,
        typ: type[T] | dict[str, Any],
        examples: list[ParseDatapoint] | None = None,
        temperature: float | None = None,
    ) -> T | PartialObj | dict[str, Any]:
        prompt = build_parse_state(text, typ, examples=examples)
        res = self.parse_force_from_prompt(
            prompt=prompt, typ=typ, temperature=temperature
        )
        return json_response_to_obj_or_partial_obj(response=res, typ=typ)

    def generate(
        self,
        instruction: str,
        text: str,
        examples: list[GenerateDatapoint] | None = None,
        temperature: float | None = None,
    ) -> str:
        prompt = build_generate_state(
            instruction=instruction, text=text, examples=examples
        )
        return self.generate_from_prompt(prompt=prompt, temperature=temperature)

    def _handle_parse_force_response(self, res: dict[str, Any], typ: type[T]) -> T:
        obj = json_response_to_obj_or_partial_obj(response=res, typ=typ)
        if isinstance(obj, dict):
            raise ModelError(f"Invalid response from model: {res}")
        return obj

    def parse_force(
        self,
        instruction: str,
        typ: type[T] | dict[str, Any],
        text: str | None = None,
        examples: list[ParseForceDatapoint] | None = None,
        temperature: float | None = None,
    ) -> T | dict[str, Any]:
        prompt = build_parse_force_state(
            instruction=instruction, text=text, typ=typ, examples=examples
        )
        res = self.parse_force_from_prompt(
            prompt=prompt, typ=typ, temperature=temperature
        )
        return self._handle_parse_force_response(res, typ)

    def _handle_score_response(
        self,
        res: dict[str, Any],
        min: int,
        max: int,
    ) -> int:
        if res is None or "score" not in res:
            raise ModelError(f"Invalid response from model: {res}")
        score = res["score"]
        if not isinstance(score, int):
            raise ModelError(f"Invalid score type: {type(score)}")
        if score < min or score > max:
            raise ModelError(f"Invalid score value: {score}")
        return score

    def score(
        self,
        instruction: str,
        text: str,
        min: int,
        max: int,
        examples: list[ScoreDatapoint] | None = None,
        temperature: float | None = None,
    ) -> int:
        prompt = build_score_state(instruction, text, min, max, examples=examples)
        res = self.parse_force_from_prompt(
            prompt=prompt, typ=Score, temperature=temperature
        )
        return self._handle_score_response(res, min, max)


def build_prompts(dps: list[Datapoint], include_response: bool = True) -> list[str]:
    if len(dps) == 0:
        return []
    typ = type(dps[0])
    for i, dp in enumerate(dps):
        if not isinstance(dp, typ):
            raise ValueError(
                f"All elements must be of type Datapoint, expected type {typ} at index {i}, got {type(dp)}"
            )
    if isinstance(dps[0], ParseDatapoint):
        build_func = build_parse_prompts
    elif isinstance(dps[0], BinaryClassifyDatapoint):
        build_func = build_binary_classify_prompts
    elif isinstance(dps[0], ClassifyDatapoint):
        build_func = build_classify_prompts
    elif isinstance(dps[0], ParseForceDatapoint):
        build_func = build_parse_force_prompts
    elif isinstance(dps[0], GenerateDatapoint):
        build_func = build_generate_prompts
    elif isinstance(dps[0], ScoreDatapoint):
        build_func = build_score_prompts
    else:
        raise ValueError(f"Unknown datapoint type: {type(dps[0])}")
    return build_func(dps, include_response)


def build_parse_prompts(
    dps: list[ParseDatapoint],
    include_response: bool = True,
) -> list[str]:
    datapoints = []
    for dp in dps:
        json_response_object = (
            dp.response.model_dump_json()
            if isinstance(dp.response, BaseModel)
            else json.dumps(dp.response)
        )
        prompt = build_parse_state(text=dp.text, typ=dp.typ)
        if include_response:
            json_response = add_md_close_tag(json_response_object)
            datapoints.append(prompt + json_response)
        else:
            datapoints.append(prompt)
    return datapoints


def build_binary_classify_prompts(
    dps: list[BinaryClassifyDatapoint],
    include_response: bool = True,
) -> list[str]:
    return build_classify_prompts(
        [
            ClassifyDatapoint(
                instruction=dp.instruction,
                text=dp.text,
                options=["true", "false"],
                response=0 if dp.response else 1,
            )
            for dp in dps
        ],
        include_response=include_response,
    )


def build_classify_prompts(
    dps: list[ClassifyDatapoint],
    include_response: bool = True,
) -> list[str]:
    def label_idx_to_label_json(idx: int, decode_map: dict[str, int]) -> str:
        label = None
        for k, v in decode_map.items():
            if v == idx:
                label = k
                break
        if label is None:
            raise ValueError(f"Label index {idx} not found in decode map")
        return f'{{"classification": "{label}"}}'

    datapoints = []
    for dp in dps:
        prompt, decode_map = build_classify_state(
            instruction=dp.instruction, text=dp.text, options=dp.options
        )
        if include_response:
            json_response_object = label_idx_to_label_json(
                idx=dp.response, decode_map=decode_map
            )
            json_response = add_md_close_tag(json_response_object)
            datapoints.append(prompt + json_response)
        else:
            datapoints.append(prompt)
    return datapoints


def build_parse_force_prompts(
    dps: list[ParseForceDatapoint],
    include_response: bool = True,
) -> list[str]:
    datapoints = []
    for dp in dps:
        json_response_obj = (
            dp.response.model_dump_json()
            if isinstance(dp.response, BaseModel)
            else json.dumps(dp.response)
        )
        prompt = build_parse_force_state(
            instruction=dp.instruction,
            text=dp.text,
            typ=dp.typ,
        )
        if include_response:
            json_response = add_md_close_tag(json_response_obj)
            datapoints.append(prompt + json_response)
        else:
            datapoints.append(prompt)
    return datapoints


def build_generate_prompts(
    dps: list[GenerateDatapoint], include_response: bool = True
) -> list[str]:
    datapoints = []
    for dp in dps:
        prompt = build_generate_state(instruction=dp.instruction, text=dp.text)
        if include_response:
            datapoints.append(prompt + dp.response)
        else:
            datapoints.append(prompt)
    return datapoints


def build_score_prompts(
    dps: list[ScoreDatapoint],
    include_response: bool = True,
) -> list[str]:
    datapoints = []
    for dp in dps:
        json_response_object = f'{{"score": {dp.response}}}'
        prompt = build_score_state(
            instruction=dp.instruction,
            text=dp.text,
            min=dp.min,
            max=dp.max,
        )
        if include_response:
            json_response = add_md_close_tag(json_response_object)
            datapoints.append(prompt + json_response)
        else:
            datapoints.append(prompt)
    return datapoints


# TODO: handle examples
def approx_prompt_str(dp: Datapoint, include_response: bool = False) -> str:
    return build_prompts(dps=[dp], include_response=include_response)[0]


# TODO: handle examples
def approx_cost_for_datapoint(
    dp: Datapoint,
    price_per_input_token: float,
) -> float:
    """For now, we approximate the cost of a datapoint as the cost of the input (output tokens are priced as input tokens as well)."""
    prompt = approx_prompt_str(dp, include_response=True)
    assert isinstance(prompt, str)
    return price_per_input_token * approx_num_tokens(prompt)


# TODO: handle examples
def approx_latency_for_datapoint(
    dp: Datapoint, latency_ms_per_output_token: float
) -> float:
    if isinstance(dp, BinaryClassifyDatapoint | ClassifyDatapoint):
        approx_response = '{"classification": 0}'
    elif isinstance(dp, ParseDatapoint):
        # this is extremely approximate
        approx_response = (
            '{"street": "main st", "city": "san francisco", "state": "CA"}'
        )
    elif isinstance(dp, GenerateDatapoint):
        # this is extremely approximate
        approx_response = "This is a generated text response."
    elif isinstance(dp, ParseForceDatapoint):
        # this is extremely approximate
        approx_response = (
            '{"street": "main st", "city": "san francisco", "state": "CA"}'
        )
    elif isinstance(dp, ScoreDatapoint):
        approx_response = '{"score": 0}'
    else:
        raise ValueError(f"Unsupported datapoint type: {type(dp)}")
    return latency_ms_per_output_token * approx_num_tokens(approx_response)
