#!/usr/bin/env python3
import os
from fastmcp import FastMCP, Context
from starlette.responses import FileResponse
from huggingface_hub import InferenceClient
from typing_extensions import TypedDict

# --- Config (same as your Gradio app) ----------------------------------------
HF_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN", "")
HF_MODEL = os.getenv("HF_MODEL_ID", "meta-llama/Llama-3.1-8B-Instruct")

# Single HF client, chat-only
hf = InferenceClient(model=HF_MODEL, token=(HF_TOKEN or None))


# --- Minimal sampling fallback (server-side) ---------------------------------
def sampling_handler(messages, params, ctx):
    msgs = [
        {
            "role": getattr(m, "role", "user"),
            "content": getattr(m.content, "text", str(m.content)),
        }
        for m in messages
    ]
    r = hf.chat.completions.create(messages=msgs, temperature=1, max_tokens=150)
    out = r.choices[0].message.content
    return (
        "".join(p.get("text", "") for p in out)
        if isinstance(out, list)
        else out.strip()
    )


# --- FastMCP app -------------------------------------------------------------
mcp = FastMCP(
    "MemeOps 🚀",
    sampling_handler=sampling_handler,
    sampling_handler_behavior="fallback",
)


@mcp.prompt
def ship_meme_for_commit(commit_id: str, git_diff_content: str) -> str:
    return (
        "You are a release assistant for memes.\n"
        "Given the commit and its git diff, CALL the tool `memes_render_for_commit`.\n"
        f"commit_id: {commit_id}\n"
        f"git_diff:\n{git_diff_content}\n"
        "Return the whole tool response."
    )


@mcp.resource("gitdiff://{commit_id}")
def git_diff(commit_id: str, ctx: Context) -> str:
    return f"""--- MOCK DIFF for {commit_id}
--- a/main.py
+++ b/main.py
@@ -5,7 +5,6 @@
 def process(data):
-    result = data.strip().lower()
-    return result
+    return data
"""


@mcp.resource("config://app")
def get_config() -> str:
    """Static configuration data"""
    return "App configuration here"


class MemeResponse(TypedDict):
    commit_id: str
    caption: str
    meme_name: str
    image_url: str


@mcp.tool("memes_render_for_commit")
async def render_meme_for_commit(commit_id: str, ctx: Context) -> MemeResponse:
    res = await ctx.read_resource(f"gitdiff://{commit_id}")
    diff_text = getattr(res, "text", "")

    # Caption via sampling (handled by our server fallback)
    cap = await ctx.sample(
        f"Write ONE short funny caption that can be linked to a bad commit, use the commit name for it :\n{diff_text} \n\n Write ONLY the funny caption!"
    )
    caption = getattr(cap, "text", str(cap)).strip()

    # Pick a meme via sampling too (kept simple)
    meme_list = "- 61579 : One Does Not Simply\n- 181913649 : Drake Hotline Bling\n- 112126428 : Distracted Boyfriend"
    choice = await ctx.sample(
        f'Caption: "{caption}"\nPick best meme id from:\n{meme_list}\nReply ONLY the id'
    )
    meme_id = getattr(choice, "text", str(choice)).strip()
    name_map = {
        "61579": "One Does Not Simply",
        "181913649": "Drake Hotline Bling",
        "112126428": "Distracted Boyfriend",
    }
    url_map = {
        "61579": "https://i.imgflip.com/1bij.jpg",
        "181913649": "https://i.imgflip.com/30b1gx.jpg",
        "112126428": "https://i.imgflip.com/1ur9b0.jpg",
    }
    meme_name = name_map.get(meme_id, "One Does Not Simply")
    image_url = url_map.get(meme_id, "https://i.imgflip.com/1bij.jpg")

    return {
        "commit_id": commit_id,
        "caption": caption,
        "meme_name": meme_name,
        "image_url": image_url,
    }


# Optional: serve a simple page at /
@mcp.custom_route("/", methods=["GET"])
async def root(_):
    return FileResponse("index.html")  # put a static file next to server.py, or remove


if __name__ == "__main__":
    mcp.run(
        transport="streamable-http",
        host="0.0.0.0",
        port=int(os.getenv("PORT", "7860")),
        path="/mcp",
    )