import os
import gradio as gr
import mimetypes
from google import genai
from google.genai import types

# Client initialisieren
client = genai.Client(
    api_key=os.environ.get("GEMINI_API_KEY"),
)

MAX_FILE_SIZE_MB = 2

def model_chat(message, history):
    try:
        contents = []
        last_role = None
        
        # 1. Historie verarbeiten
        for msg in history:
            if isinstance(msg, dict):
                role = msg.get("role")
                content = msg.get("content", "")
            else:
                role = getattr(msg, "role", "")
                content = getattr(msg, "content", "")
            
            text_val = content if isinstance(content, str) else "[Datei]"
            
            if role == "user":
                if last_role == "user":
                    contents.append(types.Content(role="model", parts=[types.Part.from_text(text="[Keine Antwort erhalten]")]))
                contents.append(types.Content(role="user", parts=[types.Part.from_text(text=text_val)]))
                last_role = "user"
            elif role in ["assistant", "model"]:
                contents.append(types.Content(role="model", parts=[types.Part.from_text(text=text_val)]))
                last_role = "model"

        if last_role == "user":
            contents.append(types.Content(role="model", parts=[types.Part.from_text(text="[Keine Antwort erhalten]")]))

        # 2. Aktuelle Nachricht & Datei-Upload (Universell mit 2MB Limit)
        current_parts = []
        
        # Text hinzufügen
        if message["text"]:
            current_parts.append(types.Part.from_text(text=message["text"]))
        
        # Dateien verarbeiten
        for file_path in message["files"]:
            file_size = os.path.getsize(file_path) / (1024 * 1024) # In MB
            
            if file_size > MAX_FILE_SIZE_MB:
                yield f"⚠️ Datei '{os.path.basename(file_path)}' überspringt das 2 MB Limit ({file_size:.2f} MB)."
                continue

            mime_type, _ = mimetypes.guess_type(file_path)
            mime_type = mime_type or "application/octet-stream"

            # Unterscheidung: Text vs. Binär (Bild, PDF, etc.)
            if mime_type.startswith("text/"):
                try:
                    with open(file_path, "r", encoding="utf-8", errors="replace") as f:
                        content_str = f.read()
                        current_parts.append(types.Part.from_text(text=f"Dateiinhalt ({os.path.basename(file_path)}):\n\n{content_str}"))
                except Exception:
                    # Fallback auf Bytes, falls Text-Lesen scheitert
                    with open(file_path, "rb") as f:
                        current_parts.append(types.Part.from_bytes(data=f.read(), mime_type=mime_type))
            else:
                with open(file_path, "rb") as f:
                    current_parts.append(types.Part.from_bytes(data=f.read(), mime_type=mime_type))

        if not current_parts:
            yield "Bitte gib eine Nachricht ein oder lade eine passende Datei hoch."
            return

        contents.append(types.Content(role="user", parts=current_parts))

        # 3. Konfiguration (Unverändert: gemini-3.1-flash-lite-preview)
        model_id = "gemini-3.1-flash-lite-preview" 
        
        tools = [types.Tool(googleSearch=types.GoogleSearch())]
        
        generate_content_config = types.GenerateContentConfig(
            thinking_config=types.ThinkingConfig(thinking_level="MINIMAL"),
            tools=tools,
        )

        # 4. Stream starten
        response_text = ""
        for chunk in client.models.generate_content_stream(
            model=model_id,
            contents=contents,
            config=generate_content_config,
        ):
            if chunk.text:
                response_text += chunk.text
                yield response_text
                    
    except Exception as e:
        yield f"❌ Fehler: {str(e)}"

# Gradio Interface
demo = gr.ChatInterface(
    fn=model_chat,
    title="Gemini Thinking AI",
    description="KI mit Suche und universellem Datei-Upload (max. 2 MB).",
    multimodal=True,
)

if __name__ == "__main__":
    demo.launch()