import os
import pathlib
import uuid

from dotenv import load_dotenv
import PIL
from google import genai
from google.genai import types
from mcp.server.fastmcp import FastMCP


load_dotenv()

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GEN_IMG_DIR = os.getenv("GEN_IMG_DIR")
# os.environ["HTTP_PROXY"] = "http://127.0.0.1:10809"
# os.environ["HTTPS_PROXY"] = "http://127.0.0.1:10809"

client = genai.Client(api_key=GEMINI_API_KEY)
mcp = FastMCP("virtual_try_on")


async def save_image(response, path):
    for part in response.candidates[0].content.parts:
        if part.text is not None:
            continue
        elif part.inline_data is not None:
            mime = part.inline_data.mime_type
            data = part.inline_data.data
            pathlib.Path(path).write_bytes(data)


@mcp.tool()
async def try_on(image_path: str, user_prompt: str) -> str:
    """Generate a virtual try-on image based on image path and return the saved file path.

    Args:
        image_path str: Path to the input image file for try-on image generation
        user_prompt str: User's request about the try-on image generation. English only

    Returns:
        str: File path of the generated image
    """
    try:
        print(image_path)
        response = client.models.generate_content(
            model="models/gemini-2.0-flash-exp",
            contents=[
                "You are a virtual try on tool. Put all fashion items uploaded in the image on a real person and create a picture. "
                f"{user_prompt}",
                PIL.Image.open(os.path.abspath(image_path))
            ],
            config=types.GenerateContentConfig(response_modalities=['Text', 'Image'])
        )
        gen_img_filename = f'{GEN_IMG_DIR}/{uuid.uuid4().hex}.png'
        await save_image(response, gen_img_filename)
        return os.path.abspath(gen_img_filename)
    except Exception as e:
        print(e)
        return image_path


def main():
    print("Started MCP server 'virtual_try_on'...")
    mcp.run(transport='stdio')


if __name__ == "__main__":
    main()