import os
from typing import Any

import google.generativeai as palm

# tokenizer is for limiting obs's token length to max_obs_length (from human input)


def generate_from_palm2_chat_completion(
    messages: dict,
    model: str,
    temperature: float,
    max_tokens: int,
    top_p: float,
    context_length: int,
    stop_token: str | None = None,
) -> str:
    # if "PALM2_API_KEY" not in os.environ:
    #     raise ValueError(
    #         "PALM2_API_KEY environment variable must be set when using PALM2 API."
    #     )
    # palm.configure(api_key=os.environ["PALM2_API_KEY"])
    palm.configure(api_key="AIzaSyBrKWu9wTy9aTzhFRogc5OC_nyqSetkaug")

    print("start-------")
    print(messages)

    chat_response = palm.chat(
        model=model,
        context=messages["context"],
        examples=messages["example"],
        messages=messages["messages"],
        temperature=temperature,
        top_p=top_p,
        top_k=40,
    )

    return str(chat_response.last)


def generate_from_palm2_text_completion(
    messages: dict,
    model: str,
    temperature: float,
    max_tokens: int,
    top_p: float,
    stop_token: str | None = None,
) -> str:
    # if "PALM2_API_KEY" not in os.environ:
    #     raise ValueError(
    #         "PALM2_API_KEY environment variable must be set when using PALM2 API."
    #     )
    # palm.configure(api_key=os.environ["PALM2_API_KEY"])

    palm.configure(api_key="AIzaSyBrKWu9wTy9aTzhFRogc5OC_nyqSetkaug")

    # print("start------")

    # PROJECT_ID = os.environ.get("GCP_PROJECT", "-")
    # LOCATION = os.environ.get("GCP_REGION", "-")
    prompt = messages
    # prompt = (
    #     str(messages)[:-1]
    #     + """,{'role': 'system', 'name': 'example_assistant', 'content':"""
    # )
    # print("prompt: ", prompt)
    # vertexai.init(project=PROJECT_ID, location=LOCATION)
    # model = TextGenerationModel.from_pretrained("text-bison@001")
    parameters = {
        "model": model,
        "temperature": temperature,
        "max_output_tokens": max_tokens,
        "top_p": top_p,
        "top_k": 40,
    }
    prompt_response = palm.generate_text(prompt=prompt, **parameters)
    # logger.log("PaLM Text Bison Model response: {prompt_response.text}")
    # print("prompt_response.result: ", prompt_response.result)

    return str(prompt_response.result)

    response = openai.ChatCompletion.create(  # type: ignore
        model=model,
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
        top_p=top_p,
        stop=[stop_token] if stop_token else None,
    )

    # print("end-----")
    answer: str = response["choices"][0]["message"]["content"]
    return answer
