import os
import gradio as gr
from gradio import ChatMessage
from typing import Iterator
import google.generativeai as genai
import time # Import time module for potential debugging/delay

print("import library complete")
print("add API key")

# get Gemini API Key from the environ variable
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
genai.configure(api_key=GEMINI_API_KEY)

print("add API key complete ")
print("add model")

used_model = "gemini-2.5-pro-exp-03-25"

# we will be using the Gemini 2.0 Flash model with Thinking capabilities
model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-01-21")

print(f"add model {used_model} complete\n")

def format_chat_history(messages: list) -> list:
    print("\nstart format history")
    """
    Formats the chat history into a structure Gemini can understand
    """
    formatted_history = []
    for message in messages:
        #print(f"t1 {message}")
        # Skip thinking messages (messages with metadata)
        #if not (message.get("role") == "assistant" and "metadata" in message):
        #    print(f"t2 {message}")
        #    formatted_history.append({
        #        "role": "user" if message.get("role") == "user" else "assistant",
        #        "parts": [message.get("content", "")]
        #    })
            
        #print(f"t2 {message}")

        if message.get("role") == "user" :
            formatted_history.append({
            "role": "user",
            "parts": [message.get("content", "")]
            })
        elif message.get("role") == "assistant" :
            formatted_history.append({
            "role": "model",
            "parts": [message.get("content", "")]
            })
        
    #print(f"t3 {formatted_history}")
    print("return formatted history")
    return formatted_history

def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
    print("start model response stream")
    """
    Streams thoughts and response with conversation history support for text input only.
    """
    if not user_message.strip(): # Robust check: if text message is empty or whitespace
        messages.append(ChatMessage(role="assistant", content="Please provide a non-empty text message. Empty input is not allowed.")) # More specific message
        yield messages
        print("Empty text message")
        return

    try:
        print(f"\n=== New Request (Text) ===")
        print(f"User message: {user_message}")

        # Format chat history for Gemini
        chat_history = format_chat_history(messages)

        #print(f"hist {chat_history}")
        
        # Initialize Gemini chat
        print("Chat parameter")
        chat = model.start_chat(history=chat_history)
        print("Start response")
        response = chat.send_message(user_message, stream=True)

        # Initialize buffers and flags
        thought_buffer = ""
        response_buffer = ""
        #thinking_complete = False

        # Add initial thinking message
        #messages.append(
        #    ChatMessage(
        #        role="assistant",
        #        content="",
        #        metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
        #    )
        #)
        
        messages.append(
            ChatMessage(
                role="assistant",
                content=response_buffer
            )
        )
        #print(f"mes {messages} \n\nhis {chat_history}")
        
        thinking_complete = True

        for chunk in response:
            print("chunk start")
            parts = chunk.candidates[0].content.parts
            current_chunk = parts[0].text

            print(f"\n=========\nparts len: {len(parts)}\n\nparts: {parts}\n\ncurrent chunk: {current_chunk}\n=========\n")
            
            if len(parts) == 2 and not thinking_complete:
                # Complete thought and start response
                thought_buffer += current_chunk
                print(f"\n=== Complete Thought ===\n{thought_buffer}")

                messages[-1] = ChatMessage(
                    role="assistant",
                    content=thought_buffer,
                    metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
                )
                yield messages

                # Start response
                response_buffer = parts[1].text
                print(f"\n=== Starting Response ===\n{response_buffer}")

                messages.append(
                    ChatMessage(
                        role="assistant",
                        content=response_buffer
                    )
                )
                thinking_complete = True

            elif thinking_complete:
                # Stream response
                response_buffer += current_chunk
                print(f"\n=== Response Chunk ===\n{current_chunk}")

                messages[-1] = ChatMessage(
                    role="assistant",
                    content=response_buffer
                )

            else:
                # Stream thinking
                thought_buffer += current_chunk
                print(f"\n=== Thinking Chunk ===\n{current_chunk}")

                messages[-1] = ChatMessage(
                    role="assistant",
                    content=thought_buffer,
                    metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
                )
            #time.sleep(0.05) #Optional: Uncomment this line to add a slight delay for debugging/visualization of streaming. Remove for final version
            print("Response end")
            yield messages

        print(f"\n=== Final Response ===\n{response_buffer}")

    except Exception as e:
        print(f"\n=== Error ===\n{str(e)}")
        messages.append(
            ChatMessage(
                role="assistant",
                content=f"I apologize, but I encountered an error: {str(e)}"
            )
        )
        yield messages

def user_message(msg: str, history: list) -> tuple[str, list]:
    """Adds user message to chat history"""
    history.append(ChatMessage(role="user", content=msg))
    return "", history


# Create the Gradio interface
with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", neutral_hue="neutral")) as demo: # Using Soft theme with adjusted hues for a refined look
    gr.Markdown("# Chat with " + used_model)

    
    gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fzelk12%2FGemini-2">
               <img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fzelk12%2FGemini-2&countColor=%23263759" />
               </a>""")

    
    chatbot = gr.Chatbot(
        type="messages",
        label=used_model + " Chatbot (Streaming Output)", #Label now indicates streaming
        render_markdown=True,
        scale=1,
        editable="all",
        avatar_images=(None,"https://lh3.googleusercontent.com/oxz0sUBF0iYoN4VvhqWTmux-cxfD1rxuYkuFEfm1SFaseXEsjjE4Je_C_V3UQPuJ87sImQK3HfQ3RXiaRnQetjaZbjJJUkiPL5jFJ1WRl5FKJZYibUA=w214-h214-n-nu")
    )

    with gr.Row(equal_height=True):
        input_box = gr.Textbox(
            lines=1,
            label="Chat Message",
            placeholder="Type your message here...",
            scale=4
        )

        with gr.Column(scale=1):
            submit_button = gr.Button("Submit", scale=1)
            clear_button = gr.Button("Clear Chat", scale=1)

    with gr.Row(equal_height=True):
        test_button = gr.Button("test", scale=1)
        test1_button = gr.Button("test1", scale=1)
        test2_button = gr.Button("test2", scale=1)
        test3_button = gr.Button("test3", scale=1)

    # Add example prompts - removed file upload examples. Kept text focused examples.
    example_prompts = [
        ["Write a short poem about the sunset."],
        ["Explain the theory of relativity in simple terms."],
        ["If a train leaves Chicago at 6am traveling at 60mph, and another train leaves New York at 8am traveling at 80mph, at what time will they meet?"],
        ["Summarize the plot of Hamlet."],
        ["Write a haiku about a cat."]
    ]

    gr.Examples(
        examples=example_prompts,
        inputs=input_box,
        label="Examples: Try these prompts to see Gemini's thinking!",
        examples_per_page=5 # Adjust as needed
    )

# Created by gemini-2.5-pro-exp-03-25
#def process_message(msg):
#    """Обрабатывает сообщение пользователя: сохраняет, отображает и генерирует ответ."""
#    msg_store_val, _, _ = lambda msg: (msg, msg, "")(msg) # Store message and clear input (inline lambda)
#    input_box_val, chatbot_val = user_message(msg_store_val, chatbot) # Add user message to chat
#    chatbot_val_final = stream_gemini_response(msg_store_val, chatbot_val) # Generate and stream response
#    return msg_store_val, input_box_val, chatbot_val_final
#
#input_box.submit(
#    process_message,
#    inputs=[input_box],
#    outputs=[msg_store, input_box, chatbot],  # Исправлены outputs, чтобы включать chatbot
#    queue=False
#)

#submit_button.click(
#    process_message,
#    inputs=[input_box],
#    outputs=[msg_store, input_box, chatbot],  # Исправлены outputs, чтобы включать chatbot
#    queue=False
#)
    
    # Set up event handlers
    msg_store = gr.State("")  # Store for preserving user message

    input_box.submit(
        lambda msg: (msg, msg, ""),  # Store message and clear input
        inputs=[input_box],
        outputs=[msg_store, input_box, input_box],
        queue=False
    ).then(
        user_message,  # Add user message to chat
        inputs=[msg_store, chatbot],
        outputs=[input_box, chatbot],
        queue=False
    ).then(
        stream_gemini_response,  # Generate and stream response
        inputs=[msg_store, chatbot],
        outputs=chatbot
    )

    submit_button.click(
        lambda msg: (msg, msg, ""),  # Store message and clear input
        inputs=[input_box],
        outputs=[msg_store, input_box, input_box],
        queue=False
    ).then(
        user_message,  # Add user message to chat
        inputs=[msg_store, chatbot],
        outputs=[input_box, chatbot],
        queue=False
    ).then(
        stream_gemini_response,  # Generate and stream response
        inputs=[msg_store, chatbot],
        outputs=chatbot
    )

    clear_button.click(
        lambda: ([], "", ""),
        outputs=[chatbot, input_box, msg_store],
        queue=False
    )

    gr.Markdown(  # Description moved to the bottom - updated for text-only
        """
        <br><br><br>  <!-- Add some vertical space -->
        ---
        ### About this Chatbot
        **Try out the example prompts below to see Gemini in action!**
        **Key Features:**
        *   Powered by Google's **Gemini 2.0 Flash** model.
        *   Supports **conversation history** for multi-turn chats.
        *   Uses **streaming** for a more interactive experience.
        **Instructions:**
        1.  Type your message in the input box below or select an example.
        2.  Press Enter or click Submit to send.
        3.  Observe the chatbot's "Thinking" process followed by the final response.
        4.  Use the "Clear Chat" button to start a new conversation.
        """
    )


# Launch the interface
if __name__ == "__main__":
    demo.launch(debug=True)