import os
from io import BytesIO

from google import genai
from google.genai import types as genai_types
from google.genai.types import (
    CreateBatchJobConfig,
    CreateCachedContentConfig,
    EmbedContentConfig,
    FunctionDeclaration,
    GenerateContentConfig,
    Part,
    SafetySetting,
    Tool,
)
from PIL import Image

os.environ["GOOGLE_API_KEY"] = ""


client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"], http_options=genai_types.HttpOptions(api_version="v1alpha"))
client.models.get(model="gemini-2.5-flash-preview-04-17")


img_path = "./image.png"
image = Image.open(img_path)

model = "gemini-2.5-flash-preview-04-17"
gen_kwargs = {
    "candidate_count": None,
    "max_output_tokens": 800,
    "top_p": None,
    "temperature": None,
    "stop_sequences": None,
    "top_k": 40,
    # "thinking_config": genai_types.ThinkingConfig(include_thoughts=True, thinking_budget=1000),
}


# genai_types.ThinkingConfig(

gen_config = genai_types.GenerateContentConfig(**gen_kwargs)

gen_config.thinking_config = None

# client.models.get(model="gemini-2.5-flash-preview-04-17")

client.models.generate_content(
    model=model,
    contents=[
        "Tell a joke like a pirate.",
    ],
    config=gen_config,
)


safe_settings = [
    SafetySetting(
        category="HARM_CATEGORY_HATE_SPEECH",
        threshold="BLOCK_ONLY_HIGH",
    ),
    SafetySetting(
        category="HARM_CATEGORY_DANGEROUS_CONTENT",
        threshold="BLOCK_ONLY_HIGH",
    ),
    SafetySetting(
        category="HARM_CATEGORY_HARASSMENT",
        threshold="BLOCK_ONLY_HIGH",
    ),
    SafetySetting(
        category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
        threshold="BLOCK_ONLY_HIGH",
    ),
]

gen_config.safety_settings = safe_settings


content = genai_types.Content(
    role="user",
    parts=[
        genai_types.Part.from_text(text="What is in the image? Give a detailed description."),
        genai_types.Part.from_bytes(data=BytesIO(image.tobytes()).read(), mime_type="image/png"),
        # genai_types.Part.from_bytes(data=data, mime_type="image/png"),
        # genai_types.Part.from_uri(file_uri=file_id.uri, mime_type=file_id.mime_type),
    ],
)

sys_content = genai_types.Content(
    role="system",
    parts=[
        "You are a helpful assistant that excels on image analysis.",
        # genai_types.Part.from_text(text="Example Image:"),
        # genai_types.Part.from_bytes(data=BytesIO(image.tobytes()).read(), mime_type="image/png"),
        # genai_types.Part.from_text(text="Example Description:"),
        # genai_types.Part.from_text(text="Let's think step by step. <A detailed description of the image.>"),
    ],
)

sys_content_raw = [
    "You are a helpful assistant that excels on image analysis.",
    "Example Image:",
    # image,
    "Example Description:",
    "Let's think step by step. <A detailed description of the image.>",
]
gen_config.system_instruction = sys_content_raw

first_ans = client.models.generate_content(
    model=model,
    contents=["What is in the image? Give a detailed description.", image],
    config=gen_config,
)
first_ans.text

dir(first_ans)

first_ans.model_dump_json()

content_model = genai_types.Content(
    role="model",
    parts=[
        genai_types.Part.from_text(text=first_ans.text),
    ],
)
# content_model2 = genai_types.Content(
#     role="model",
#     parts=[
#         genai_types.Part.from_text(text="This is a cat."),
#     ],
# )

second_ans = client.models.generate_content(
    model=model,
    contents=[
        content,
        image,
        # content2,
        content_model,
        "Please give me an ordered list of the inputs you received so far. Include any texts and images. Same order as you received them.",
    ],
    config=gen_config,
)

print(second_ans.text)

third_ans = client.models.generate_content(
    model=model,
    contents=[
        content,
        image,
        content_model,
        "Please give me an ordered list of the inputs you received so far. Include any texts and images. Same order as you received them.",
        second_ans.text,
        "Can you describe the crops with as much detail as possible?",
    ],
    config=gen_config,
)
print(third_ans.text)

crops_text = """Okay, here are the detailed descriptions of the cropped images I received, as separate entries:

**Crop 1 (Top-Left Corner):**

*   **Content:** This crop focuses on the top-left section of the full image.
*   **Logo:** The primary element is the "One Stop Market" logo. The logo consists of text with the words "One Stop" above the word "Market" The image associated with the text has what looks like a shopping basket with a product inside.
*   **Navigation:** Below the logo are more categories, each with sub-categories under them. All text is generated.

**Crop 2 (Top-Center):**

*   **Content:** This crop shows the center section of the top part of the image, emphasizing the product search bar and some of the shoe products.
*   **Search Bar:** A prominent search bar spans the width of the image.
*   **Product Grid:** A row of shoe products are visible, including white heel, white boot, and tan flat-shoes.
*   **Text:** Much of the text in the image is unintelligible gibberish.

**Crop 3 (Top-Right Corner):**

*   **Content:** This crop shows the top right corner of the image.
*   **Shopping Cart:** A shopping cart icon is present.
*   **Menu Items:** There is also a menu item with illegible text to the left of the shopping cart.

**Crop 4 (Bottom-Left Corner):**

*   **Content:** This crop focuses on the left side of the image.
*   **Sidebar:** The categories and filters, such as "Baroness Videos" are displayed.
*   **List of Options:** More lists under the text of video options are displayed as well.

**Crop 5 (Bottom-Right Corner):**

*   **Content:** This crop focuses on the bottom right corner.
*   **Shoes:** A row of shoe products are visible, including tan sandals, black loafers, and white trainers.
*   **Text:** Much of the text in the image is unintelligible gibberish."""


second_ans = client.models.generate_content(
    model="gemini-2.0-flash-exp-image-generation",
    contents=[
        content,
        image,
        # content2,
        content_model,
        # image,
        # "What's your name?",
        # content_model2,
        "Please give me an ordered list of the inputs you received so far. Include any texts and images. Same order as you received them.",
        genai_types.Content(role="model", parts=[genai_types.Part.from_text(text=x)]),
        "Can you describe the crops with as much detail as possible?",
        genai_types.Content(role="model", parts=[genai_types.Part.from_text(text=crops_text)]),
        "Could you please recreate images as close as possible to the crops?",
    ],
    config=genai_types.GenerateContentConfig(response_modalities=["Text", "Image"]),
)

print(second_ans.text)
response = client.models.generate_content(
    model="gemini-2.0-flash-exp-image-generation",
    # contents=f"Could you please generate: an image of a cat and a dog. Please also provide a description in text.",
    # contents=f"Could you please generate: an image of a cat, a description of it, an image of a dog, a description of the dog?",
    contents=f"Could you please generate: an image of a cat and an image of a dog? Please also provide descriptions in text.",
    config=genai_types.GenerateContentConfig(response_modalities=["Text", "Image"]),
)

all_content = []
for i, part in enumerate(response.candidates[0].content.parts):
    if part.text is not None:
        all_content.append(part.text)
    elif part.inline_data is not None:
        image = Image.open(BytesIO((part.inline_data.data)))
        all_content.append(image)

print(all_content)


# LINK: list base models
for model in client.models.list():
    if "gemini-2.0-flash-001" in model.name:
        print(model)


chat = client.chats.create(model="gemini-2.0-flash")


# Case 1: not valid
content = genai_types.Content(
    role="user",
    parts=[
        genai_types.Part.from_text(text="What is in the image?"),
        image,
    ],
)


# Case 2: valid
content = genai_types.Content(
    role="user",
    parts=[
        genai_types.Part.from_text(text="What is in the image?"),
        image,
    ],
)


# response = client.models.generate_content(
#     model=model,
#     contents=[
#         image,
#         "Write a short and engaging blog post based on this picture.",
#     ],
#     config=gen_config,
# )

file_id = client.files.upload(file="examples/bird.png")
file_id.name
file_id.name
