# hf_spaces_mcp_server.py
from mcp.server.fastmcp import FastMCP
from gradio_client import Client
import os
import sys
import asyncio
import logging
import uvicorn # Correctly imported

# Configure basic logging to see server activity
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logger = logging.getLogger(__name__)

# The port that Hugging Face Spaces expects our app to listen on.
# It injects this as an environment variable 'PORT'.
# Default to 7860 if PORT is not set (e.g., for local testing outside HF Spaces).
HOST_PORT = int(os.getenv("PORT", 7860))
HOST_ADDRESS = os.getenv("HOST", "0.0.0.0") # Bind to all interfaces

# Initialize FastMCP. The host/port settings passed here will be used by
# mcp.streamable_http_app() to configure its internal server if not mounted to FastAPI.
# However, when passed to uvicorn.run(), uvicorn's host/port will take precedence.
# It's good practice to keep them consistent or omit from here if uvicorn dictates.
# For simplicity, we'll keep them here as they are also part of MCP server's self-description.
mcp = FastMCP(
    name="GradioSpacesTools",
    host=HOST_ADDRESS,
    port=HOST_PORT
)

clients = {}

def get_client(space_id: str) -> Client:
    """Get or create a Gradio client for the specified space."""
    if space_id not in clients:
        logger.info(f"Creating Gradio client for Space ID: {space_id}")
        clients[space_id] = Client(space_id)
    return clients[space_id]

@mcp.tool()
async def generate_image(prompt: str, space_id: str = "ysharma/SanaSprint") -> str:
    """Generate an image using Flux.
    
    Args:
        prompt (str): Text prompt describing the image to generate
        space_id (str): HuggingFace Space ID to use 
    
    Returns:
        str: The URL of the generated image.
    """
    logger.info(f"Generating image for prompt: '{prompt}' using Space: '{space_id}'")
    client = get_client(space_id)
    try:
        result = await client.predict(
            prompt=prompt,
            model_size="1.6B",
            seed=0,
            randomize_seed=True,
            width=1024,
            height=1024,
            guidance_scale=4.5,
            num_inference_steps=2,
            api_name="/infer"
        )
        logger.info(f"Image generation successful. Result: {result}")
        return result
    except Exception as e:
        logger.error(f"Error generating image: {e}")
        return f"Error generating image: {e}"

@mcp.tool()
async def run_dia_tts(prompt: str, space_id: str = "ysharma/Dia-1.6B") -> str:
    """Text-to-Speech Synthesis.
    
    Args:
        prompt (str): Text prompt describing the conversation between speakers S1, S2
        space_id (str): HuggingFace Space ID to use 
    
    Returns:
        str: The URL of the generated audio.
    """
    logger.info(f"Performing TTS for prompt: '{prompt}' using Space: '{space_id}'")
    client = get_client(space_id)
    try:
        result = await client.predict(
            text_input=f"""{prompt}""",
            audio_prompt_input=None, 
            max_new_tokens=3072,
            cfg_scale=3,
            temperature=1.3,
            top_p=0.95,
            cfg_filter_top_k=30,
            speed_factor=0.94,
            api_name="/generate_audio"
        )
        logger.info(f"TTS successful. Result: {result}")
        return result
    except Exception as e:
        logger.error(f"Error performing TTS: {e}")
        return f"Error performing TTS: {e}"

if __name__ == "__main__":
    sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 1)
    logger.info(f"Starting GradioSpacesTools MCP server via Uvicorn on http://{HOST_ADDRESS}:{HOST_PORT}")
    
    # The correct way to get the FastMCP ASGI application for streamable-http
    # as indicated by the article and FastMCP's internal structure.
    # Uvicorn will then run this ASGI application.
    uvicorn.run(
        mcp.streamable_http_app(), # <-- Use this method! It returns the ASGI app.
        host=HOST_ADDRESS,
        port=HOST_PORT,
        log_level="info", # Set logging level for Uvicorn itself
        loop="asyncio"    # Ensure Uvicorn uses the asyncio event loop
    )
    logger.info("GradioSpacesTools MCP server stopped.")