"""
Generate image with Gemini using prompt, then discuss it in a follow-up turn.

This example:
1. Uses Gemini with prompt to generate an image
2. Records the completion with the generated image
3. Adds the response to history
4. Asks a followup question about the generated image
5. Records the followup conversation with actual LLM responses and start/stop times
"""

import json
import os
import time

import requests

from freeplay import CallInfo, Freeplay, RecordPayload, ResponseInfo

fp_client = Freeplay(
    freeplay_api_key=os.environ["FREEPLAY_API_KEY"],
    api_base=f"{os.environ['FREEPLAY_API_URL']}/api",
)

project_id = os.environ["FREEPLAY_PROJECT_ID"]

print("=" * 80)
print("TURN 1: Generate image with Gemini")
print("=" * 80)

# Create session for multi-turn conversation
session = fp_client.sessions.create()
print(f"Created session ID: {session.session_info.session_id}")

# Turn 1: Generate an image
image_description = "A serene mountain landscape at sunset with a crystal clear lake"
input_variables_1 = {"description": image_description}

template_name = os.environ.get("FREEPLAY_PROMPT_TEMPLATE_NAME", "imagegen")
environment = os.environ.get("FREEPLAY_ENVIRONMENT", "latest")
formatted_prompt_1 = fp_client.prompts.get_formatted(
    project_id=project_id,
    template_name=template_name,
    environment="latest",
    variables=input_variables_1,
)

start_1 = time.time()

# Call Gemini to generate image
model_name = "gemini-2.5-flash-image"
api_key = os.environ.get("GEMINI_API_KEY")

if not api_key:
    raise ValueError("GEMINI_API_KEY environment variable is required")

print(f"Generating image: {image_description}")

request_data = {
    "contents": [{"parts": [{"text": f"Create a picture of {image_description}"}]}]
}

try:
    response = requests.post(
        f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={api_key}",
        headers={"Content-Type": "application/json"},
        data=json.dumps(request_data),
        timeout=60,
    )

    if response.status_code != 200:
        raise ValueError(f"Gemini API request failed: {response.text}")

    response_data = response.json()
    candidate = response_data["candidates"][0]
    assistant_message_1 = candidate["content"]

    if "role" not in assistant_message_1:
        assistant_message_1["role"] = "model"

    # Verify we got an image
    parts = assistant_message_1["parts"]
    has_image = any("inlineData" in part for part in parts)

    if not has_image:
        raise ValueError("No image generated by Gemini")

    print("Image generated successfully")

except requests.RequestException as e:
    raise ValueError(f"Failed to call Gemini API: {e}")

end_1 = time.time()

# Start history from formatted prompt
history = list(formatted_prompt_1.llm_prompt)
history.append(assistant_message_1)

# Record Turn 1
call_info_1 = CallInfo(
    start_time=start_1,
    end_time=end_1,
    provider="gemini",
    model=model_name,
)
response_info_1 = ResponseInfo(is_complete=True)

record_response_1 = fp_client.recordings.create(
    RecordPayload(
        project_id=project_id,
        all_messages=history,
        session_info=session.session_info,
        inputs=input_variables_1,
        prompt_version_info=formatted_prompt_1.prompt_info,
        call_info=call_info_1,
        response_info=response_info_1,
    )
)

print(f"Recorded Turn 1 completion ID: {record_response_1.completion_id}")

# Turn 2: Ask about the generated image
print("\n" + "=" * 80)
print("TURN 2: Discuss the generated image")
print("=" * 80)

input_variables_2 = {"query": "What are the key visual elements in this image?"}

# Add Turn 2 user message to history
history.append({"role": "user", "parts": [{"text": input_variables_2["query"]}]})

print(f"Turn 2: {len(history)} messages (includes generated image from Turn 1)")

# Call Gemini for Turn 2
start_2 = time.time()

request_data_2 = {"contents": history}

try:
    response_2 = requests.post(
        f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={api_key}",
        headers={"Content-Type": "application/json"},
        data=json.dumps(request_data_2),
        timeout=60,
    )

    if response_2.status_code != 200:
        raise ValueError(f"Gemini API request failed: {response_2.text}")

    response_data_2 = response_2.json()
    candidate_2 = response_data_2["candidates"][0]
    assistant_message_2 = candidate_2["content"]

    if "role" not in assistant_message_2:
        assistant_message_2["role"] = "model"

    response_text = assistant_message_2["parts"][0]["text"]
    print(f"Assistant: {response_text[:100]}...")

except requests.RequestException as e:
    raise ValueError(f"Failed to call Gemini API: {e}")

end_2 = time.time()

# Add response to history
history.append(assistant_message_2)

# Record Turn 2 - generated image is now in history
call_info_2 = CallInfo(
    start_time=start_2,
    end_time=end_2,
    provider="gemini",
    model=model_name,
)
response_info_2 = ResponseInfo(is_complete=True)

record_response_2 = fp_client.recordings.create(
    RecordPayload(
        project_id=project_id,
        all_messages=history,
        session_info=session.session_info,
        inputs=input_variables_2,
        prompt_version_info=formatted_prompt_1.prompt_info,
        call_info=call_info_2,
        response_info=response_info_2,
    )
)

print(f"Recorded Turn 2 completion ID: {record_response_2.completion_id}")

print("\n" + "=" * 80)
print("Multi-turn with generated image complete!")
print(f"Session ID: {session.session_info.session_id}")
print(
    f"View in UI: {os.environ.get('FREEPLAY_API_URL', 'http://localhost:8080')}/projects/{project_id}/sessions"
)
print("=" * 80)
