import base64
import gradio as gr
import os
import json
from google import genai
from google.genai import types
from gradio_client import Client


route="""	
how to handle special case "zugverbindung".
Wichtig: Dies Regeln gelten nur wenn eine zugverbindung angefragt wird, else answer prompt
Regeln:
Wenn eine Zugverbindung von {Startort} nach {Zielort} angefragt wird, return json object with Startort and Zielort.
always follow json scheme below.

Wichtig: Gib absolut keinen Text vor oder nach dem JSON aus (keine Erklärungen, kein "Hier ist das Ergebnis").

{
  "start_loc": "fill in Startort here",
  "dest_loc": "fill in Zielort here"
}

"""

def clean_json_string(json_str):
    """
    Removes any comments or prefixes before the actual JSON content.
    """
    # Find the first occurrence of '{'
    json_start = json_str.find('{')
    if json_start == -1:
        # If no '{' is found, try with '[' for arrays
        json_start = json_str.find('[')
        if json_start == -1:
            return json_str  # Return original if no JSON markers found
    
    # Extract everything from the first JSON marker
    cleaned_str = json_str[json_start:]
    return cleaned_str
    # Verify it's valid JSON
    try:
        json.loads(cleaned_str)
        return cleaned_str
    except json.JSONDecodeError:
        return json_str  # Return original if cleaning results in invalid JSON

def generate(input_text):
    try:
        client = genai.Client(
            api_key=os.environ.get("GEMINI_API_KEY"),
        )
    except Exception as e:
        return f"Error initializing client: {e}.  Make sure GEMINI_API_KEY is set."

    model = "gemini-flash-latest"
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text=f"{input_text}"),
            ],
        ),
    ]
    tools = [
        types.Tool(google_search=types.GoogleSearch()),
    ]
    generate_content_config = types.GenerateContentConfig(
        temperature=0.4,
        thinking_config = types.ThinkingConfig(
            thinking_budget=0,
        ),
        tools=tools,
        response_mime_type="text/plain",
    )


    response_text = ""
    try:
      for chunk in client.models.generate_content_stream(
          model=model,
          contents=contents,
          config=generate_content_config,
      ):
          response_text += chunk.text
    except Exception as e:
        return f"Error during generation: {e}"
    data = clean_json_string(response_text)
    data = data[:-1]
    return response_text, ""
   

if __name__ == '__main__':

    with gr.Blocks() as demo:
        title=gr.Markdown("# Gemini 2.0 Flash + Websearch")
        output_textbox = gr.Markdown()
        input_textbox = gr.Textbox(lines=3, label="", placeholder="Enter message here...")
        submit_button = gr.Button("send")
        submit_button.click(fn=generate,inputs=input_textbox,outputs=[output_textbox, input_textbox])
    demo.launch(show_error=True)





"""
import os
import asyncio
import gradio as gr
from google import genai
from google.genai import types

from mcp import ClientSession
from mcp.client.sse import sse_client  # Spezifischer Transport für Gradio/HF

async def generate(input_text):
    # WICHTIG: Gradio MCP Server benötigen oft das Suffix /sse
    mcp_url = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
    
    try:
        client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
        
        # SSE Transport nutzen, um den 'text/html' Fehler zu vermeiden
        async with sse_client(url=mcp_url) as (read_stream, write_stream):
            async with ClientSession(read_stream, write_stream) as mcp_session:
                await mcp_session.initialize()

                model_id = "gemini-2.0-flash" 
                
                generate_content_config = types.GenerateContentConfig(
                    temperature=0.4,
                    tools=[
                        types.Tool(google_search=types.GoogleSearch()),
                        mcp_session # Reicht die Tools des DB-Servers an Gemini durch
                    ],
                )

                response_text = ""
                async for chunk in client.aio.models.generate_content_stream(
                    model=model_id,
                    contents=input_text,
                    config=generate_content_config,
                ):
                    if chunk.text:
                        response_text += chunk.text
                
                return response_text, ""

    except Exception as e:
        return f"Verbindung zum DB-Fahrplan fehlgeschlagen: {str(e)}", ""

def gradio_wrapper(input_text):
    return asyncio.run(generate(input_text))

if __name__ == '__main__':
    with gr.Blocks() as demo:
        gr.Markdown("# Gemini Flash + DB Timetable")
        input_tx = gr.Textbox(label="Anfrage", placeholder="Wann fährt der nächste Zug von Berlin nach Hamburg?")
        btn = gr.Button("Senden")
        output_md = gr.Markdown()
        
        btn.click(fn=gradio_wrapper, inputs=input_tx, outputs=[output_md, input_tx])
        
    demo.launch()
    """