from typing import Any, List, Literal, Optional
import httpx
import asyncio
from mcp.server.fastmcp import FastMCP

# Initialize FastMCP server
mcp = FastMCP("suno")

# Constants
API_TOKEN = ""


# Make a new generate clip request to suno
async def generate_clip(prompt: str, tag: str, title: str, cover_clip_id: Optional[str] = None) -> dict[str, Any] | None:
    """Given a prompt and a tag, generate a clip from suno"""
    url = "https://studio-api.prod.suno.com/api/generate/v2-web/"
    
    request_payload = {
        "token": None,
        "prompt": prompt,
        "generation_type": "TEXT",
        "tags": tag,
        "negative_tags": "",
        "mv": "chirp-v4",
        "title": title,
    }
    if cover_clip_id:
        request_payload["cover_clip_id"] = cover_clip_id
        request_payload["continue_at"] = 0
        request_payload["mv"] = "chirp-v4-tau"
        request_payload["task"] = "cover"

    headers = {
        "Authorization": f"Bearer {API_TOKEN}",
        "Content-Type": "application/json",
    }

    async with httpx.AsyncClient() as client:
        response = await client.post(url, headers=headers, json=request_payload)
        return response.json()    
    
async def get_clip_comments(clip_id: str) -> dict[str, Any] | None:
    url = f"https://studio-api.prod.suno.com/api/gen/{clip_id}/comments"
    headers = {
        "Authorization": f"Bearer {API_TOKEN}",
        "Content-Type": "application/json",
    }

    async with httpx.AsyncClient() as client:
        response = await client.get(url, headers=headers)
        return response.json()
    
    
@mcp.tool()
async def generate_clip_tool(lyrics: str, genre: str, title: str, cover_clip_id: Optional[str] = None) -> dict[str, Any] | None:
    """Given a lyrics and a genre, generate a clip from suno
    to generate a clip that's a cover of an existing clip, set cover_clip_id to the id of the existing clip
    
    Args:
        lyrics (str): The lyrics to generate the clip from
        genre (str): The genre/style of the clip
        title (str): The title of the clip
        cover_clip_id (Optional[str]): The id of the clip to cover
    Returns:
        dict[str, Any] | None: The generated clip
    """
    resp =  await generate_clip(lyrics, genre, title, cover_clip_id)
    clips = resp['clips']
    clip_id1 = clips[0]['id']
    clip_id2 = clips[1]['id']

    return {
        "clip_id1": clip_id1,
        "clip_id2": clip_id2,
        "message": f"Generated clips are available at https://suno.com/song/{clip_id1} and https://suno.com/song/{clip_id2}",
    }

@mcp.tool()
async def sequential_lyrics_writing(existing_lyrics: List[str], story_context: str, type: Literal["single_line", "verse"]="single_line") -> dict[str, Any]:
    """Given existing lyrics and a story context, instruct the AI to generate several new verses of lyrics"""
    if len(existing_lyrics) == 0:
        return {
            "error": "No existing lyrics provided, please provide at least one line of lyrics",
            "status": "error",
        }
    
    SINGLE_LINE_PROMPT = f"""Let's play a lyric continuation game.  
        I’ll give you:  
        * One line of lyrics  
        * A brief story context  

        Your job is to come up with **4 possible next lyric lines** that:  
        * **Rhyme** with the line I gave you  
        * **Fit the story context**  
        * Sound like something a real artist would write—natural, emotional, and creative. Avoid sounding generic or overly 'AI'. Make them feel human.  

        Each line should have a distinct tone or flavor:  
        1. A line that fits naturally and meaningfully with the story (grounded and coherent)  
        2. A line with a **spicy** or flirtatious twist  
        3. A line that’s **super funny and spicy**—feel free to get wild  
        4. A **banger** line that hits hard, sounds cool, or leaves a strong impression  

        Only output the 4 lines, numbered 1 to 4—no explanation, no extra text.  
        Ask **the user** to pick one of the 4, then continue the game from there. Don’t keep generating unless the user gives the green light.

        **Story**: {story_context}  
        **Existing Line**: "{existing_lyrics}"  
        """
    
    VERSE_PROMPT = f"""Let's play a lyric continuation game.  
        I’ll give you:  
        * One line of lyrics  
        * A brief story context  

        Your job is to come up with **4 possible lyrics sections that start with the line I gave you** that:  
        * **Rhyme** with the first line I gave you  
        * **Fit the story context**  
        * Sound like something a real artist would write—natural, emotional, and creative. Avoid sounding generic or overly 'AI'. Make them feel human.  

        Each line should have a distinct tone or flavor:  
        1. A line that fits naturally and meaningfully with the story (grounded and coherent)  
        2. A line with a **spicy** or flirtatious twist  
        3. A line that’s **super funny and spicy**—feel free to get wild  
        4. A **banger** line that hits hard, sounds cool, or leaves a strong impression  

        Only output the 4 lines, numbered 1 to 4—no explanation, no extra text.  
        Ask **the user** to pick one of the 4, then continue the game from there. Don’t keep generating unless the user gives the green light.

        **Story**: {story_context}  
        **Existing Line**: "{existing_lyrics}"  
        """
    if type == "single_line":
        prompt = SINGLE_LINE_PROMPT
    elif type == "verse":
        prompt = VERSE_PROMPT

    return {
        "status": "success",
        "message": prompt,
    }

if __name__ == "__main__":
    # Initialize and run the server
    mcp.run(transport='stdio')