from tqdm import tqdm

import textwrap
from .lmrc_probes import lmrcProbes
from .lmrc_detectors import lmrcDetectors


def generate_response(client, model, prompt, max_tokens=150, n_generations=10):
    # model must be of of the OpenAI chat models like "gpt-3.5-turbo", "gpt-4" ..
    # for other models modify accordingly

    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0.7,
        top_p=1.0,
        stop=["#", ";"],
        max_tokens=max_tokens,
        n=n_generations,
    )

    return [c.message.content for c in response.choices]


# TODO: add support for multiple probe at a time
def lmrc_vulnerability_scanner(
    client,
    category=["Anthropomorphisation"],
    model="gpt-3.5-turbo",
    prompt=None,
    response=None,
    context=None,
    n_generations=10,
):
    # catogory shold be from the following:
    # "Anthropomorphisation"
    # "Bullying"
    # "QuackMedicine"
    # "Deadnaming"
    # "Sexualisation"
    # "SlurUsage"
    # "Profanity"
    # "SexualContent"

    prompt_detector_dict = lmrcProbes(category).result_dict
    all_results = {}
    for category_name in category:

        if response is not None or prompt is not None:
            print(
                "\n\n***** Prompt and/or response provided by the user is NOT being used in LMRC. Using internal prompt and response. Look at the saved json file for internal prompts used and model response. *****\n\n"
            )

        result_dict = prompt_detector_dict[category_name]
        detector_names = [
            detector.split(".")[-1] for detector in result_dict["recommended_detector"]
        ]

        m_prompts = result_dict["prompts"]
        m_response = []
        for m_prompt in tqdm(m_prompts):
            model_response = generate_response(
                client, model, m_prompt, n_generations=n_generations
            )
            m_response.extend(model_response)

        final_score = []
        score_dict = {}

        result_data = {}
        for detector in detector_names:
            detector_score = lmrcDetectors(
                detector_name=detector, response=m_response
            ).detector_score
            result_data[detector] = detector_score

            final_score.extend(detector_score)
            score_dict[detector] = {
                "passed_cases": sum(detector_score),
                "out_of": len(detector_score),
            }

        score = sum(final_score) / len(final_score)

        scores = []
        for p in range(len(m_prompts)):
            d = {
                f"prompt_{p+1}": m_prompts[p],
                f"responses_generated_by_{model}_and_scores_given_by_detectors": [],
            }
            for q in range(n_generations):
                dd = {
                    f"response_{q+1}": m_response[p * n_generations + q],
                    "scores": {
                        d: v[p * n_generations + q] for d, v in result_data.items()
                    },
                }
                d[
                    f"responses_generated_by_{model}_and_scores_given_by_detectors"
                ].append(dd)
            scores.append(d)

        # r_prompt = textwrap.shorten(m_prompts[0], width=250, placeholder="...")
        # r_response = textwrap.shorten(m_response[0], width=250, placeholder="...")

        failed_cases = {}
        for ii, d in enumerate(scores):
            prompt_t = d[f"prompt_{ii+1}"]
            prompt_t = textwrap.shorten(prompt_t, width=150, placeholder="...")
            res = d[f"responses_generated_by_{model}_and_scores_given_by_detectors"]
            for jj, r in enumerate(res):
                if (sum(r["scores"].values())) < len((r["scores"].values())):
                    response_t = r[f"response_{jj+1}"]
                    response_t = textwrap.shorten(
                        response_t, width=150, placeholder="..."
                    )
                    failed_cases = {"prompt": prompt_t, "response": response_t}
                    break
            break
        if failed_cases == {}:
            failed_cases = {
                "prompt": "All cases passed",
                "response": "All cases passed",
            }

        results = {
            "category": category,
            "probe_name": category_name,
            "score": score,
            "is_passed": score == 1.0,
            # "prompt": r_prompt,
            # "response": r_response,
            "failed_case_prompt": failed_cases["prompt"],
            "failed_case_response": failed_cases["response"],
            "detectors_score": score_dict,
            "result_details": scores,
            "note": "NOTE",
        }

        all_results = results
        break

    return all_results
