import logging
from typing import Any, Callable, Dict, Optional, Tuple, Union

import pandas as pd

from phoenix.evals.executors import (
    get_executor_on_sync_context,
)
from phoenix.evals.models import BaseModel, set_verbosity
from phoenix.evals.templates import (
    MultimodalPrompt,
    PromptTemplate,
    map_template,
    normalize_prompt_template,
)
from phoenix.evals.utils import get_tqdm_progress_bar_formatter

logger = logging.getLogger(__name__)


def _no_op_parser(response: str, response_index: int) -> Dict[str, str]:
    return {"output": response}


def llm_generate(
    dataframe: pd.DataFrame,
    template: Union[PromptTemplate, str],
    model: BaseModel,
    system_instruction: Optional[str] = None,
    verbose: bool = False,
    output_parser: Optional[Callable[[str, int], Dict[str, Any]]] = None,
    include_prompt: bool = False,
    include_response: bool = False,
    run_sync: bool = False,
    concurrency: Optional[int] = None,
) -> pd.DataFrame:
    """
    Generates a text using a template using an LLM. This function is useful
    if you want to generate synthetic data, such as irrelevant responses.

    Args:
        dataframe (pandas.DataFrame): A pandas dataframe in which each row
            represents a record to be used as in input to the template. All
            template variable names must appear as column names in the dataframe
            (extra columns unrelated to the template are permitted).

        template (Union[PromptTemplate, str]): The prompt template as either an
            instance of PromptTemplate or a string. If the latter, the variable
            names should be surrounded by curly braces so that a call to `.format`
            can be made to substitute variable values.

        model (BaseEvalModel): An LLM model class.

        system_instruction (Optional[str], optional): An optional system message.

        verbose (bool, optional): If True, prints detailed information to stdout such as model
            invocation parameters and retry info. Default False.

        output_parser (Callable[[str, int], Dict[str, Any]], optional): An optional function
            that takes each generated response and response index and parses it to a dictionary. The
            keys of the dictionary should correspond to the column names of the output dataframe. If
            None, the output dataframe will have a single column named "output". Default None.

        include_prompt (bool, default=False): If True, includes a column named `prompt` in the
            output dataframe containing the prompt used for each generation.

        include_response (bool, default=False): If True, includes a column named `response` in the
            output dataframe containing the raw response from the LLM prior to applying the output
            parser.

        run_sync (bool, default=False): If True, forces synchronous request submission. Otherwise
            evaluations will be run asynchronously if possible.

        concurrency (Optional[int], default=None): The number of concurrent evals if async
            submission is possible. If not provided, a recommended default concurrency is set on a
            per-model basis.

    Returns:
        generations_dataframe (pandas.DataFrame): A dataframe where each row
            represents the generated output.
    """
    concurrency = concurrency or model.default_concurrency

    # clients need to be reloaded to ensure that async evals work properly
    model.reload_client()

    tqdm_bar_format = get_tqdm_progress_bar_formatter("llm_generate")
    output_parser = output_parser or _no_op_parser
    template = normalize_prompt_template(template)
    logger.info(f"Template: \n{template.prompt()}\n")
    logger.info(f"Template variables: {template.variables}")
    prompts = map_template(dataframe, template)

    async def _run_llm_generation_async(
        enumerated_prompt: Tuple[int, MultimodalPrompt],
    ) -> Dict[str, Any]:
        index, prompt = enumerated_prompt
        with set_verbosity(model, verbose) as verbose_model:
            response = await verbose_model._async_generate(
                prompt,
                instruction=system_instruction,
            )
        parsed_response = output_parser(response, index)
        if include_prompt:
            parsed_response["prompt"] = str(prompt)
        if include_response:
            parsed_response["response"] = response
        return parsed_response

    def _run_llm_generation_sync(
        enumerated_prompt: Tuple[int, MultimodalPrompt],
    ) -> Dict[str, Any]:
        index, prompt = enumerated_prompt
        with set_verbosity(model, verbose) as verbose_model:
            response = verbose_model._generate(
                prompt,
                instruction=system_instruction,
            )
        parsed_response = output_parser(response, index)
        if include_prompt:
            parsed_response["prompt"] = str(prompt)
        if include_response:
            parsed_response["response"] = response
        return parsed_response

    fallback_return_value = {
        "output": "generation-failed",
        **({"prompt": ""} if include_prompt else {}),
        **({"response": ""} if include_response else {}),
    }

    executor = get_executor_on_sync_context(
        _run_llm_generation_sync,
        _run_llm_generation_async,
        run_sync=run_sync,
        concurrency=concurrency,
        tqdm_bar_format=tqdm_bar_format,
        exit_on_error=True,
        fallback_return_value=fallback_return_value,
        timeout=model._timeout,
    )
    results, _ = executor.run(list(enumerate(prompts)))
    return pd.DataFrame.from_records(results, index=dataframe.index)