import os

import gradio as gr
from dotenv import load_dotenv
from strands import Agent
from strands.models.litellm import LiteLLMModel
from mcp.client.streamable_http import streamablehttp_client
from strands.tools.mcp.mcp_client import MCPClient


load_dotenv()
os.environ["STRANDS_TOOL_CONSOLE_MODE"] = "enabled"


def create_streamable_http_transport():
   return streamablehttp_client(os.environ["MCP_SERVER"])

streamable_http_mcp_client = MCPClient(create_streamable_http_transport)

model = LiteLLMModel(
    client_args={
        "api_key": os.environ["OPENROUTER_API_KEY"],
        "api_base": "https://openrouter.ai/api/v1",
    },
    model_id="openrouter/google/gemini-2.5-flash"
)



SYSTEM_PROMPT = """
You are an Image Agent that allows users to generate an image using prompt or edit images 
using prompt. Always render the images as markdown
"""

def convert_history(history):
    """
    Convert Gradio 6 history format:
    [{"role": "...", "content": [{"text": "..."}]}]

    Into Required format:
    [{"role": "...", "content": "..."}]
    """
    messages = []
    for msg in history:
        # Gradio packs content into list of content blocks
        if isinstance(msg["content"], list):
            # extract just the text pieces
            text = "".join(block.get("text", "") for block in msg["content"])
        else:
            text = msg["content"]

        messages.append({"role": msg["role"], "content": text})
    return messages


def response_generator(message, history):
    """
    Response Generator
    """
    messages = convert_history(history) if history else []
    agent_messages = []
    for msg in messages:
        agent_messages.append({
            "role": msg["role"],
            "content": [{"text": msg["content"]}]
        })
    with streamable_http_mcp_client:
        # Get the tools from the MCP server
        tools = streamable_http_mcp_client.list_tools_sync()
        image_agent = Agent(
            model=model,
            system_prompt=SYSTEM_PROMPT,
            tools=tools,
            messages=agent_messages
        )
        print(agent_messages)
        messages.append({"role": "user", "content": message})
        response = image_agent(message)
        response = response.message["content"][0]["text"]
        response = response.replace("sandbox:/", "")
        yield response

demo = gr.ChatInterface(
    fn=response_generator,
    examples=[["A scenic landscape with mountains, a river, and a clear blue sky", None]],
    title="Image Agent"
)
demo.launch()