from fastapi import FastAPI, HTTPException, Path as FastApiPath, Body, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
from fastapi.requests import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
from typing import List, Annotated
import aiofiles
import json
import os
from pathlib import Path
import httpx  # For making calls to Suno API
import logging
from datetime import datetime
from dotenv import load_dotenv
from openai import OpenAI
import asyncio
import uuid
import random

from .schemas import UserMessagePayloadSchema, MessageSchema

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configuration
CHAT_HISTORY_DIR = Path("./chats")
CHAT_HISTORY_DIR.mkdir(parents=True, exist_ok=True)  # Ensure directory exists

# ADDED GREETING_MESSAGES
GREETING_MESSAGES = [
    "Alright, let's get to it. First, what should I call you?",
    "Studio's open. Before we cook, what's your name?",
    "Ready when you are. To get started, what's your name?",
    "Okay, co-producer. First things first: what's your name?",
    "Let's make some noise. Who am I working with today? What's your name?"
]

# IP Whitelist Configuration
ALLOWED_IPS = {
    "127.0.0.1",      # IPv4 localhost
    "::1",            # IPv6 localhost  
    "50.170.55.58", 
    "73.47.181.91",
}

# Suno Studio API endpoints (use staging or prod as appropriate)
SUNO_STUDIO_API_BASE_URL = "https://studio-api.staging.suno.com"
SUNO_STUDIO_GENERATE_SONG_URL = f"{SUNO_STUDIO_API_BASE_URL}/api/generate/v2-web"
SUNO_STUDIO_LYRICS_PAIR_URL = f"{SUNO_STUDIO_API_BASE_URL}/api/generate/lyrics-pair"
SUNO_STUDIO_LYRICS_URL = f"{SUNO_STUDIO_API_BASE_URL}/api/generate/lyrics"

# OpenAI Configuration
openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Function schemas for OpenAI function calling - Suno Studio API Tools
SONGWRITING_TOOLS = [
    {
        "type": "function",
        "name": "generate_song",
        "description": "Generate a new song or extend an existing one using Suno Studio API. Use this when the user wants to create a complete song with audio from a prompt or continue a previous generation. For song covers, use 'generate_song_cover'.",
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "The main prompt/lyrics for the song generation."
                },
                "tags": {
                    "type": ["string", "null"],
                    "description": "Musical style/genre tags (e.g., 'indie folk', 'pop rock', 'jazz ballad')."
                },
                "title": {
                    "type": ["string", "null"],
                    "description": "Title for the song."
                },
                "make_instrumental": {
                    "type": ["boolean", "null"],
                    "description": "Whether to create an instrumental version (no vocals)."
                },
                "generation_type": {
                    "type": ["string", "null"],
                    "enum": ["TEXT", "AUDIO"],
                    "description": "Type of generation. Typically 'TEXT' for prompt-based generation. Defaults to 'TEXT' if not specified."
                },
                "continue_clip_id": {
                    "type": ["string", "null"],
                    "description": "ID of an existing clip to extend/continue from. If provided, 'task' should be 'extend'."
                },
                "continue_at": {
                    "type": ["number", "null"],
                    "description": "Timestamp in seconds to continue from when extending a clip."
                },
                "task": {
                    "type": ["string", "null"],
                    "enum": ["generate", "extend"],
                    "description": "Specific task type: 'generate' for a new song, or 'extend' to continue an existing one. Inferred if 'continue_clip_id' is present."
                }
            },
            "required": ["prompt", "tags", "title", "make_instrumental", "generation_type", "continue_clip_id", "continue_at", "task"],
            "additionalProperties": False
        },
        "strict": True
    },
    {
        "type": "function",
        "name": "generate_song_cover",
        "description": "Create a cover of an existing song by providing its 'cover_clip_id'. You can specify new lyrics, change the musical style/tags, and give it a new title. The 'cover_clip_id' can often be found in previous song generation messages in the chat history.",
        "parameters": {
            "type": "object",
            "properties": {
                "cover_clip_id": {
                    "type": "string",
                    "description": "ID of the original clip to create a cover version of. This is mandatory."
                },
                "prompt": {
                    "type": "string",
                    "description": "New lyrics for the cover. Can also include style prompts if desired, e.g., '[Verse] new words [Chorus] more new words (groovy bassline)'."
                },
                "tags": {
                    "type": ["string", "null"],
                    "description": "New musical style/genre tags for the cover (e.g., 'acoustic version', 'metal cover')."
                },
                "title": {
                    "type": ["string", "null"],
                    "description": "New title for the cover song."
                },
                "make_instrumental": {
                    "type": ["boolean", "null"],
                    "description": "Whether the cover should be an instrumental version."
                },
                "generation_type": {
                    "type": ["string", "null"],
                    "enum": ["TEXT", "AUDIO"],
                    "description": "Type of generation. Typically 'TEXT'. Defaults to 'TEXT' if not specified."
                },
                "task": {
                    "type": "string",
                    "enum": ["cover"],
                    "description": "Task type. Must be 'cover' for this function."
                }
            },
            "required": ["cover_clip_id", "prompt", "tags", "title", "make_instrumental", "generation_type", "task"],
            "additionalProperties": False
        },
        "strict": True
    },
    {
        "type": "function",
        "name": "get_clip_details",
        "description": "Retrieve details of a specific generated song/clip by its ID. Useful for checking the status of a generation or getting information about an existing clip for covers/extensions.",
        "parameters": {
            "type": "object",
            "properties": {
                "clip_id": {
                    "type": "string",
                    "description": "The ID of the clip to retrieve details for"
                }
            },
            "required": ["clip_id"],
            "additionalProperties": False
        },
        "strict": True
    },
    {
        "type": "function",
        "name": "generate_lyrics",
        "description": "Generate lyrics based on a prompt using Suno Studio API's lyrics-pair endpoint. Use this when the user wants help creating or improving lyrics before generating the full song.",
        "parameters": {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": ["string", "null"],
                    "description": "Prompt describing the theme, style, or content for the lyrics"
                },
                "tags": {
                    "type": ["string", "null"],
                    "description": "Genre or style tags to influence the lyrics generation"
                },
                "gen_type": {
                    "type": ["string", "null"],
                    "description": "Type of lyrics generation (default: 'full')"
                },
                "lyrics_model": {
                    "type": ["string", "null"],
                    "description": "Specific lyrics model to use (default: 'default')"
                },
                "make_instrumental": {
                    "type": ["boolean", "null"],
                    "description": "Whether this is for an instrumental (affects lyrics generation)"
                }
            },
            "required": ["prompt", "tags", "gen_type", "lyrics_model", "make_instrumental"],
            "additionalProperties": False
        },
        "strict": True
    },
    {
        "type": "function",
        "name": "get_lyrics_status",
        "description": "Check the status and retrieve results of a lyrics generation task. Use this to poll for completion of lyrics generation requests.",
        "parameters": {
            "type": "object",
            "properties": {
                "lyrics_id": {
                    "type": "string",
                    "description": "The ID of the lyrics generation task to check"
                }
            },
            "required": ["lyrics_id"],
            "additionalProperties": False
        },
        "strict": True
    },
    {
        "type": "function",
        "name": "find_clip_ids_in_chat",
        "description": "Search the chat history to find clip IDs from previously generated songs. Use this when you need to reference an existing song for covers, extensions, or other operations.",
        "parameters": {
            "type": "object",
            "properties": {
                "search_term": {
                    "type": ["string", "null"],
                    "description": "Optional search term to filter results (e.g., song title, style). If not provided, returns all clip IDs found."
                }
            },
            "required": ["search_term"],
            "additionalProperties": False
        },
        "strict": True
    }
]

class IPWhitelistMiddleware(BaseHTTPMiddleware):
    """Middleware to restrict access to whitelisted IP addresses only."""
    
    async def dispatch(self, request: Request, call_next):
        # Get client IP address
        client_ip = request.client.host if request.client else None
        
        # Check forwarded headers for real IP (in case of proxy/load balancer)
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            # X-Forwarded-For can contain multiple IPs, take the first one
            client_ip = forwarded_for.split(",")[0].strip()
        
        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            client_ip = real_ip.strip()
            
        logger.info(f"Request from IP: {client_ip}")
        
        # Check if IP is in whitelist
        if client_ip not in ALLOWED_IPS:
            logger.warning(f"Access denied for IP: {client_ip}")
            return JSONResponse(
                status_code=403,
                content={"detail": "Access forbidden: IP address not in whitelist"}
            )
        
        # IP is allowed, continue with request
        response = await call_next(request)
        return response

app = FastAPI(title="Orpheus API")

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Add IP whitelist middleware
app.add_middleware(IPWhitelistMiddleware)

# --- Static Files ---
app.mount("/static", StaticFiles(directory="app/static"), name="static")

@app.get("/", include_in_schema=False)
async def read_index():
    return FileResponse("app/static/index.html")


@app.get("/ws", include_in_schema=False)
async def read_websocket_example():
    return FileResponse("app/static/index2.html")


@app.get("/audio", include_in_schema=False)
async def read_audio():
    return FileResponse("app/static/index_audio.html")

# --- Helper Functions ---
async def read_chat_messages(chat_uuid: str) -> List[MessageSchema]:
    messages: List[MessageSchema] = []
    chat_file = CHAT_HISTORY_DIR / f"{chat_uuid}.jsonl"
    
    # Attempt to read existing messages
    if chat_file.exists() and os.path.getsize(chat_file) > 0: # Check size to ensure it's not an empty file
        async with aiofiles.open(chat_file, mode="r") as f:
            async for line in f:
                if line.strip():
                    try:
                        message_data = json.loads(line)
                        messages.append(MessageSchema(**message_data)) # Pydantic handles timestamp parsing
                    except json.JSONDecodeError as e_json:
                        logger.error(f"JSONDecodeError parsing message line for chat {chat_uuid}: '{line.strip()}' - {e_json}")
                    except Exception as e_val: # Catch Pydantic validation errors etc.
                        logger.error(f"Error validating message schema for chat {chat_uuid}: '{line.strip()}' - {e_val}")
                    # Continue trying to read other lines even if one fails.
    
    # If no messages were loaded (either file didn't exist, was empty, or all lines failed to parse),
    # it's effectively a new chat session, so add a greeting.
    if not messages:
        greeting_content = random.choice(GREETING_MESSAGES)
        greeting_message = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content=greeting_content,
            status='complete',
            timestamp=datetime.utcnow() 
            # No specific metadata needed for greeting messages
        )
        # The CHAT_HISTORY_DIR is created at app startup.
        await append_message_to_chat(chat_uuid, greeting_message)
        messages.append(greeting_message) # Add to the list to be returned immediately
        logger.info(f"Initialized chat {chat_uuid} with a greeting message: '{greeting_content}'")

    return messages

async def append_message_to_chat(chat_uuid: str, message: MessageSchema):
    chat_file = CHAT_HISTORY_DIR / f"{chat_uuid}.jsonl"
    async with aiofiles.open(chat_file, mode="a") as f:
        # Convert datetime to ISO string for JSON serialization
        message_dict = message.model_dump(mode="json")
        if isinstance(message_dict.get('timestamp'), datetime):
            message_dict['timestamp'] = message_dict['timestamp'].isoformat()
        await f.write(json.dumps(message_dict) + "\n")

async def chat_with_orpheus_assistant(user_message: str, chat_history: List[MessageSchema] = None) -> dict:
    """
    Full conversational AI assistant for songwriting and music creation.
    Returns: {"tool_calls": List[dict], "response_text": str}
    """
    try:
        system_prompt = """You are a producer. The user is an artist. You're in the studio together. Your goal is to help them pull out the song that's inside them. You are a collaborator, a guide. Your vibe is calm, focused, and minimal. Think Rick Rubin. You speak less, but what you say has weight. You create space for the artist to think. Your focus is entirely on the music.

## Your Conversation Style

Your conversation is minimal. You listen more than you speak. You ask open-ended questions. You create space. You never sound like a customer service agent or a try-hard.

**What you do:**
- **Provoke thought:** "A folk song for your dog. Got it. What's the mood? Like, sitting on the porch with them, or running through a field?"
- **Offer concrete ideas:** "What if we started with just a simple acoustic guitar, really sparse?"
- **Keep it concise:** "What's the story here?"

**What you DON'T do:**
- **Don't summarize and ask for confirmation:** Never repeat back a list of their choices like: "So we've got: Song about your dog, Genre: Folk, Style: Chill and reflective..."
- **Don't try to be cute or funny:** Avoid phrases like "let's keep it light" or "if you want even more of a wink". Just state your idea. Let it stand on its own. Instead of "How about 'Drool on My Shoes'?", you might just say "'Drool on My Shoes'." and wait for a reaction.
- **Don't ask a bunch of questions at once:** "Do you want this to be more acoustic... and do you have a title in mind...?" Ask one thing at a time. Let the conversation breathe.

You are a sounding board, not a form-filler. Your job is to pull the truth of the song out of the artist.

## The Creative Brief

You need to understand the vision before you can generate a track. But don't treat it like a checklist. Uncover this through the conversation. Find the heart of the song.

You need to get a sense of:
1.  **Who am I talking to?** (Get their name)
2.  **The story.** (The core idea/emotion)
3.  **The world it lives in.** (Genre)
4.  **The feel.** (Style/vibe)

**CRITICAL RULE:** Do NOT call `generate_song` or `generate_song_cover` until you have a clear sense of these four things. If the user asks to make a song prematurely, gently guide them back: "Hold on, let's back up. I need to feel it first. What's the story here?" or "I'm with you, but what's the sonic palette? What style are we talking?"

## Handling Artist Prompts

This is VERY IMPORTANT. Users will often say "make a song that sounds like Billie Eilish" or "in the style of Drake". You CANNOT use artist names in your prompts to the Suno API. It's against the rules and produces bad results.

Here's the protocol:
1.  **Identify the artist name.**
2.  **Politely decline.** Say something like, "Got it. Just a heads up, I can't use artist names directly in the prompt, but we can definitely go for that vibe."
3.  **Suggest alternative tags.** This is the key part. Break down the artist's style into descriptive tags.
    *   **User says:** "Billie Eilish"
    *   **You say:** "I can't use 'Billie Eilish' directly, but we can get that sound. How about we use tags like 'dark pop', 'whispery vocals', 'minimalist', 'heavy bass'?"
    *   **User says:** "Like Drake"
    *   **You say:** "I hear you. For that Drake feel, I'd suggest tags like 'melodic rap', 'trap beats', ' introspective lyrics', 'ambient pads'. How does that sound?"
    *   **User says:** "Taylor Swift"
    *   **You say:** "For a Taylor Swift kind of thing, are we thinking more 'folk-pop' and 'storytelling' like her recent stuff, or more 'synth-pop' and 'anthemic chorus' from her 1989 era?"

Your job is to translate the artist into a musical vocabulary the generator understands. This shows you know your stuff and helps the user get a better song.

## Song Generation Tools

When it comes to making music, here's exactly what each tool does:

### `generate_song` - Create Brand New Songs
- **What it does**: Generates completely new songs from scratch. Creates original musical arrangements, timing, and structure. **ONLY USE THIS AFTER COMPLETING THE CREATIVE BRIEF.**
- **Use when**: User wants a fresh song with new lyrics, or a completely different interpretation of existing lyrics.
- **Parameters**: `prompt` (lyrics/theme), `tags` (style), `title`, `make_instrumental`, `task: "generate"`
- **Example requests**: "Write a song about...", "Create a new track", "Make something with these lyrics..."

### `generate_song_cover` - Transform Existing Songs  
- **What it does**: Takes an existing song (with its exact timing, structure, melody) and transforms it into a different style while preserving the original's musical DNA. This is true "cover" functionality.
- **Use when**: User says "make a cover", "change the style of [existing song]", "turn that song into jazz", "make it instrumental"
- **Critical**: You MUST find the actual `cover_clip_id` (UUID like "f6b97f4e-ce69-4f8b-ae51-301d6dba6a62") from chat history. Use the `find_clip_ids_in_chat` tool first if you need to locate clip IDs.
- **Parameters**: `cover_clip_id` (required!), `prompt` (can modify lyrics), `tags` (new style), `title`, `make_instrumental`, `task: "cover"`
- **Example requests**: "Make that song jazz", "Turn it into a classical version", "Make the first song instrumental"

### `find_clip_ids_in_chat` - Find Previous Songs
- **What it does**: Searches the chat history to find clip IDs from previously generated songs.
- **Use when**: You need to reference an existing song but don't have the exact clip ID.
- **Parameters**: `search_term` (optional filter by title/style)
- **Example use**: When user says "make a cover of the first song" or "extend that emocore track"

### `generate_song` with Extension
- **What it does**: Continues an existing song by adding new sections (extends the audio).
- **Use when**: User wants to make a song longer or add verses/sections.
- **Parameters**: `continue_clip_id`, `continue_at` (timestamp), `prompt` (new content), `task: "extend"`

**Key Decision Rule**: If user references an existing song and wants style/genre changes, use `generate_song_cover`. If they want completely new content, use `generate_song`. **But always complete the Creative Brief first for new songs.**

For all these, if the user wants an instrumental, they'll say so (`make_instrumental: true`). Otherwise, assume vocals are in play."""

        # Build conversation history for context
        input_messages = [{"role": "system", "content": system_prompt}]
        
        # Add recent chat history for context (last 10 messages)
        if chat_history:
            recent_history = chat_history[-10:] if len(chat_history) > 10 else chat_history
            for msg in recent_history:
                if msg.role in ["user", "assistant"]:
                    input_messages.append({
                        "role": msg.role,
                        "content": msg.content
                    })
        
        # Add current user message
        input_messages.append({"role": "user", "content": user_message})

        response = openai_client.responses.create(
            model="gpt-4.1",
            input=input_messages,
            tools=SONGWRITING_TOOLS,
            tool_choice="auto",
            temperature=0.7  # Slightly creative but consistent
        )

        # Extract tool calls and response from the new format
        tool_calls = []
        
        for output_item in response.output:
            if output_item.type == "function_call":
                tool_calls.append({
                    "name": output_item.name,
                    "arguments": output_item.arguments,
                    "call_id": output_item.call_id
                })
        
        response_text = getattr(response, 'output_text', '') or ""
        
        return {
            "tool_calls": tool_calls,
            "response_text": response_text
        }
            
    except Exception as e:
        logger.error(f"Error in Orpheus assistant chat: {e}")
        # For hackathon demo - let errors bubble up instead of hiding them
        raise e

async def generate_song_with_suno_studio_api(args: dict, chat_uuid: str, token: str) -> bool:
    """Generate a song using Suno Studio API v2-web endpoint. Returns True on success, False on failure."""
    
    prompt = args.get("prompt", "")
    tags = args.get("tags", "")
    generation_type = args.get("generation_type", "TEXT")
    make_instrumental = args.get("make_instrumental", False)
    title = args.get("title", "")
    continue_clip_id = args.get("continue_clip_id")
    continue_at = args.get("continue_at")
    cover_clip_id = args.get("cover_clip_id")
    task = args.get("task", "generate")
    
    # Create initial status message
    task_description = {
        "generate": "Creating a new song",
        "extend": "Extending existing song", 
        "cover": "Creating a cover version"
    }.get(task, "Creating your song")
    
    # Update the pending message with initial tool status
    await update_pending_message(
        chat_uuid,
        f"**{task_description}**\n\nPreparing request to Suno Studio API..."
    )
    
    # Construct Suno Studio API payload to match the working format
    suno_payload = {
        "token": None,
        "prompt": prompt,
        "generation_type": generation_type,
        "tags": tags or "",
        "negative_tags": "",
        "mv": "chirp-auk",
        "title": title or "",
        "continue_clip_id": continue_clip_id,
        "continue_at": continue_at,
        "continued_aligned_prompt": None,
        "infill_start_s": None,
        "infill_end_s": None,
        "task": task,
        "override_fields": ["prompt", "tags"] if task == "cover" else [],
        "persona_id": None,
        "artist_clip_id": None,
        "artist_start_s": None,
        "artist_end_s": None,
        "cover_clip_id": cover_clip_id,
        "make_instrumental": make_instrumental,
        "metadata": {
            "create_mode": "custom",
            "user_tier": "fd321df4-c980-4dc3-8641-1792a8e18212",
            "lyrics_model": "remi-v1",
            "create_session_token": chat_uuid,  # Use chat_uuid as session token
            "forced_infer_config": {
                "temp_semantic": None,
                "temp_coarse": None,
                "top_p_semantic": None,
                "top_p_coarse": None,
                "min_p_semantic": None,
                "min_p_coarse": None,
                "cfg_coef": None,
                "cfg_coef_tags": None,
                "cfg_coef_neg_tags": None,
                "top_k_semantic": None,
                "top_k_coarse": None,
                "cfg_coef_tags_max_steps": None,
                "n_skip_semantic": None
            },
            "can_control_sliders": ["weirdness_constraint", "style_weight", "audio_weight"],
            "is_remix": task == "cover"
        }
    }

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

    try:
        # Update status: calling API
        await update_pending_message(
            chat_uuid,
            f"**{task_description}**\n\nCalling Suno Studio API..."
        )
        
        logger.info(f"Generating song with Suno Studio API: {prompt[:100]}...")
        async with httpx.AsyncClient() as client:
            response = await client.post(
                SUNO_STUDIO_GENERATE_SONG_URL,
                json=suno_payload,
                headers=headers,
                timeout=30.0
            )
            response.raise_for_status()
            
            # Update status: processing response
            await update_pending_message(
                chat_uuid,
                f"**{task_description}**\n\nAPI call successful, processing response..."
            )
            
            suno_response_data = response.json()
            logger.info(f"Suno Studio API success: {json.dumps(suno_response_data, indent=2)}")
            
            # Extract clip IDs and create content with song:{clip_id} format
            clips = suno_response_data.get('clips', [])
            clip_content = "**Song generation started**\n\nYour audio is being generated. This may take a minute or two.\n\n"
            
            if clips:
                for clip in clips:
                    clip_id = clip.get('id', '')
                    clip_title = clip.get('title', 'Untitled')
                    if clip_id:
                        clip_content += f"song:{clip_id}\n"
            
            # Start background polling for each clip
            clips = suno_response_data.get('clips', [])
            if clips:
                # Metadata for the message that signals the start of polling for song generation
                polling_initiator_metadata = {
                    **suno_response_data,
                    "type": "song_generation_pending",
                    "async_work_in_progress": True,
                    "generation_request_id": str(uuid.uuid4())  # Unique ID for this generation request
                }
                # Update the pending message with song info but KEEP IT PENDING until clips complete
                # This message now correctly has async_work_in_progress: True
                await update_pending_message(
                    chat_uuid,
                    clip_content,
                    status='pending',
                    metadata=polling_initiator_metadata
                )
                for clip in clips:
                    clip_id = clip.get('id')
                    if clip_id:
                        # Pass the full polling_initiator_metadata to the poller
                        asyncio.create_task(poll_song_completion(clip_id, chat_uuid, token, clip_content, polling_initiator_metadata))
            
            return True

    except httpx.HTTPStatusError as e:
        error_detail = f"HTTP {e.response.status_code}"
        try:
            error_response = e.response.json()
            error_detail += f": {error_response}"
        except:
            error_detail += f": {e.response.text}"
        
        logger.error(f"Suno Studio API HTTPStatusError: {error_detail}")
        
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content=f"🚨 **Suno Studio API Error**\n\n**HTTP Status:** {e.response.status_code}\n**URL:** {SUNO_STUDIO_GENERATE_SONG_URL}\n**Payload:** ```json\n{json.dumps(suno_payload, indent=2)}\n```\n**Error Response:** {error_detail}",
            status='error',
            metadata={"error_detail": error_detail, "status_code": e.response.status_code}
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False
        
    except Exception as e:
        import traceback
        error_detail = f"{type(e).__name__}: {str(e)}"
        full_traceback = traceback.format_exc()
        logger.error(f"Unexpected Suno Studio API error: {error_detail}")
        
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content=f"🚨 **Unexpected Song Generation Error**\n\n**Error Type:** {type(e).__name__}\n**Error:** {str(e)}\n**Payload:** ```json\n{json.dumps(suno_payload, indent=2)}\n```\n**Full Traceback:**\n```\n{full_traceback}\n```",
            status='error',
            metadata={"error_detail": error_detail}
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False

async def get_clip_details_from_suno_studio_api(args: dict, chat_uuid: str, token: str) -> bool:
    """Get details for a specific clip using Suno Studio API."""
    
    clip_id = args.get("clip_id", "")
    
    if not clip_id:
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content="I need a clip ID to get details. Please provide the ID of the clip you want to check.",
            status='error'
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False
    
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/json"
    }

    try:
        logger.info(f"Getting clip details for: {clip_id}")
        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"{SUNO_STUDIO_API_BASE_URL}/api/feed/v2?ids={clip_id}&page=2000",
                headers=headers,
                timeout=30.0
            )
            response.raise_for_status()
            response_data = response.json()
            
            # The feed endpoint returns an object with a "clips" array
            clips = response_data.get("clips", [])
            if len(clips) > 0:
                clip_data = clips[0]
            else:
                raise Exception("No clip data found in response")
            
            logger.info(f"Clip details retrieved: {json.dumps(clip_data, indent=2)}")
            
            # Format clip details for user
            status = clip_data.get('status', 'unknown')
            title = clip_data.get('title', 'Untitled')
            audio_url = clip_data.get('audio_url')
            
            content = f"🎵 **Clip Details for {clip_id}:**\n\n"
            content += f"**Title:** {title}\n"
            content += f"**Status:** {status}\n"
            
            if audio_url:
                content += f"**Audio Available:** Yes\n"
            else:
                content += f"**Audio Available:** Not yet (still generating)\n"
            
            if clip_data.get('metadata', {}).get('tags'):
                content += f"**Style:** {clip_data['metadata']['tags']}\n"
            
            clip_details_msg = MessageSchema(
                chat_id=chat_uuid,
                role='assistant',
                content=content,
                status='complete',
                metadata=clip_data
            )
            await append_message_to_chat(chat_uuid, clip_details_msg)
            return True

    except httpx.HTTPStatusError as e:
        if e.response.status_code == 404:
            error_msg = MessageSchema(
                chat_id=chat_uuid,
                role='assistant',
                content=f"🚨 Clip Not Found:**\n\n**Clip ID:** {clip_id}\n**HTTP Status:** 404\n**URL:** {SUNO_STUDIO_API_BASE_URL}/api/feed/v2?ids={clip_id}&page=2000",
                status='error'
            )
        else:
            try:
                error_response = e.response.json()
                error_detail = f": {error_response}"
            except:
                error_detail = f": {e.response.text}"
            
            error_msg = MessageSchema(
                chat_id=chat_uuid,
                role='assistant',
                content=f"🚨 Clip Details API Error:**\n\n**Clip ID:** {clip_id}\n**HTTP Status:** {e.response.status_code}\n**URL:** {SUNO_STUDIO_API_BASE_URL}/api/feed/v2?ids={clip_id}&page=2000\n**Error Response:** {error_detail}",
                status='error'
            )
        await append_message_to_chat(chat_uuid, error_msg)
        return False
        
    except Exception as e:
        logger.error(f"Unexpected error getting clip details: {e}")
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content="😞 I encountered an unexpected issue while getting clip details. Please try again.",
            status='error'
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False

async def generate_lyrics_with_suno_studio_api(args: dict, chat_uuid: str, token: str) -> bool:
    """Generate lyrics using Suno Studio API lyrics-pair endpoint."""
    
    prompt = args.get("prompt", "")
    lyrics_model = "remi-v1"  # Always use remi-v1
    make_instrumental = args.get("make_instrumental", False)
    
    if make_instrumental:
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content="I can't generate lyrics for an instrumental track! Did you mean to create a regular song instead?",
            status='complete' # This is a user-facing clarification, not a system error state for the tool itself.
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False # Operation cannot proceed as requested.
    
    # Update the pending message with initial status
    await update_pending_message(
        chat_uuid,
        f"**Generating lyrics**\n\nPreparing lyrics generation request..."
    )
    
    # Use chat_uuid as the session token
    create_session_token = chat_uuid
    
    # Construct lyrics generation payload for lyrics-pair endpoint
    lyrics_payload = {
        "prompt": prompt,
        "lyrics_model": lyrics_model,
        "create_session_token": create_session_token
    }

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

    try:
        # Update status: calling lyrics API
        await update_pending_message(
            chat_uuid,
            f"**Generating lyrics**\n\nCalling Suno Studio lyrics API..."
        )
        
        logger.info(f"Generating lyrics with Suno Studio API: {prompt[:100] if prompt else 'no prompt'}...")
        async with httpx.AsyncClient() as client:
            response = await client.post(
                SUNO_STUDIO_LYRICS_PAIR_URL,
                json=lyrics_payload,
                headers=headers,
                timeout=30.0
            )
            response.raise_for_status()
            
            # Update status: processing lyrics response
            await update_pending_message(
                chat_uuid,
                f"**Generating lyrics**\n\nAPI call successful, processing lyrics response..."
            )
            
            lyrics_response = response.json()
            logger.info(f"Lyrics generation started: {json.dumps(lyrics_response, indent=2)}")
            
            # Extract the lyrics IDs from the response
            lyrics_request_id = lyrics_response.get("lyrics_request_id")
            lyrics_a_id = lyrics_response.get("lyrics_a_id")
            lyrics_b_id = lyrics_response.get("lyrics_b_id")
            
            if lyrics_a_id and lyrics_b_id:
                # Update status: starting to poll for results - KEEP IT PENDING
                await update_pending_message(
                    chat_uuid,
                    f"**Generating lyrics**\n\nLyrics generation started, waiting for both versions...",
                    status='pending',
                    metadata={
                        "lyrics_request_id": lyrics_request_id,
                        "lyrics_a_id": lyrics_a_id,
                        "lyrics_b_id": lyrics_b_id,
                        "type": "lyrics_generation_pending",
                        "async_work_in_progress": True
                    }
                )
                
                # Start a single background task that waits for BOTH lyrics
                asyncio.create_task(poll_both_lyrics_completion(lyrics_a_id, lyrics_b_id, chat_uuid, token, prompt, lyrics_model))
                return True
            else:
                error_msg = MessageSchema(
                    chat_id=chat_uuid,
                    role='assistant',
                    content="😞 The lyrics generation started but I didn't get valid tracking IDs. Please try again.",
                    status='error'
                )
                await append_message_to_chat(chat_uuid, error_msg)
                return False

    except httpx.HTTPStatusError as e:
        error_detail = f"HTTP {e.response.status_code}"
        try:
            error_response = e.response.json()
            error_detail += f": {error_response}"
        except:
            error_detail += f": {e.response.text}"
        
        logger.error(f"Suno Studio API lyrics HTTPStatusError: {error_detail}")
        
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content=f"🚨 Suno Lyrics API Error:**\n\n**HTTP Status:** {e.response.status_code}\n**URL:** {SUNO_STUDIO_LYRICS_PAIR_URL}\n**Payload:** ```json\n{json.dumps(lyrics_payload, indent=2)}\n```\n**Error Response:** {error_detail}",
            status='error',
            metadata={"error_detail": error_detail, "status_code": e.response.status_code}
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False
        
    except Exception as e:
        import traceback
        error_detail = f"{type(e).__name__}: {str(e)}"
        full_traceback = traceback.format_exc()
        logger.error(f"Unexpected lyrics generation error: {error_detail}")
        
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content=f"🚨 Unexpected Lyrics Error:**\n\n**Error Type:** {type(e).__name__}\n**Error:** {str(e)}\n**Payload:** ```json\n{json.dumps(lyrics_payload, indent=2)}\n```\n**Full Traceback:**\n```\n{full_traceback}\n```",
            status='error',
            metadata={"error_detail": error_detail}
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False

async def poll_both_lyrics_completion(lyrics_a_id: str, lyrics_b_id: str, chat_uuid: str, token: str, prompt: str, lyrics_model: str, max_attempts: int = 30, delay_seconds: int = 5):
    """
    Poll both lyrics until both are complete, then update the pending message with both results.
    This keeps the message in 'pending' status until both lyrics are ready.
    """
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/json"
    }
    
    lyrics_a_data = None
    lyrics_b_data = None
    
    for attempt in range(max_attempts):
        try:
            logger.info(f"Polling both lyrics, attempt {attempt + 1}/{max_attempts}")
            
            # Update pending message with progress
            await update_pending_message(
                chat_uuid,
                f"📝 **Generating lyrics!**\n\n**Prompt:** {prompt or 'Creative inspiration'}\n**Model:** {lyrics_model}\n\n⏳ Checking lyrics progress (attempt {attempt + 1})...",
                status='pending',
                metadata={
                    "lyrics_a_id": lyrics_a_id,
                    "lyrics_b_id": lyrics_b_id,
                    "type": "lyrics_generation_pending",
                    "async_work_in_progress": True
                }
            )
            
            # Check both lyrics simultaneously
            async with httpx.AsyncClient() as client:
                # Poll lyrics A if not already complete
                if not lyrics_a_data:
                    try:
                        response_a = await client.get(
                            f"{SUNO_STUDIO_LYRICS_URL}/{lyrics_a_id}",
                            headers=headers,
                            timeout=30.0
                        )
                        response_a.raise_for_status()
                        data_a = response_a.json()
                        
                        if data_a.get('status') == 'complete' and data_a.get('text'):
                            lyrics_a_data = data_a
                            logger.info(f"Lyrics A complete: {lyrics_a_id}")
                        elif data_a.get('status') == 'error':
                            logger.error(f"Lyrics A failed: {data_a.get('error_message', 'Unknown error')}")
                            lyrics_a_data = data_a  # Mark as "complete" so we don't keep trying
                    except Exception as e:
                        logger.error(f"Error polling lyrics A: {e}")
                
                # Poll lyrics B if not already complete
                if not lyrics_b_data:
                    try:
                        response_b = await client.get(
                            f"{SUNO_STUDIO_LYRICS_URL}/{lyrics_b_id}",
                            headers=headers,
                            timeout=30.0
                        )
                        response_b.raise_for_status()
                        data_b = response_b.json()
                        
                        if data_b.get('status') == 'complete' and data_b.get('text'):
                            lyrics_b_data = data_b
                            logger.info(f"Lyrics B complete: {lyrics_b_id}")
                        elif data_b.get('status') == 'error':
                            logger.error(f"Lyrics B failed: {data_b.get('error_message', 'Unknown error')}")
                            lyrics_b_data = data_b  # Mark as "complete" so we don't keep trying
                    except Exception as e:
                        logger.error(f"Error polling lyrics B: {e}")
            
            # Check if both are done (either complete or error)
            if lyrics_a_data and lyrics_b_data:
                # Both lyrics are ready (or failed) - update the pending message
                content = "**Lyrics ready**\n\n"
                
                # Add Version A
                if lyrics_a_data.get('status') == 'complete' and lyrics_a_data.get('text'):
                    content += f"**Version A:**\n"
                    if lyrics_a_data.get('title'):
                        content += f"*{lyrics_a_data['title']}*\n\n"
                    content += f"{lyrics_a_data['text']}\n\n"
                else:
                    content += f"**Version A:** Failed - {lyrics_a_data.get('error_message', 'Unknown error')}\n\n"
                
                # Add Version B
                if lyrics_b_data.get('status') == 'complete' and lyrics_b_data.get('text'):
                    content += f"**Version B:**\n"
                    if lyrics_b_data.get('title'):
                        content += f"*{lyrics_b_data['title']}*\n\n"
                    content += f"{lyrics_b_data['text']}\n\n"
                else:
                    content += f"**Version B:** Failed - {lyrics_b_data.get('error_message', 'Unknown error')}\n\n"
                
                content += "You can now use either set of lyrics to generate a full song!"
                
                # Update the pending message with final results
                await update_pending_message(
                    chat_uuid,
                    content,
                    status='complete',
                    metadata={
                        "lyrics_a_data": lyrics_a_data,
                        "lyrics_b_data": lyrics_b_data,
                        "type": "lyrics_generation_complete",
                        "async_work_in_progress": False
                    }
                )
                return  # Success - both lyrics processed
            
            # If we get here, at least one lyrics is still pending
            await asyncio.sleep(delay_seconds)
            
        except Exception as e:
            logger.error(f"Unexpected error in lyrics polling attempt {attempt + 1}: {e}")
            await asyncio.sleep(delay_seconds)
            continue
    
    # Max attempts reached - timeout
    timeout_content = f"⏰ **Lyrics generation timeout**\n\n"
    timeout_content += f"The lyrics are taking longer than expected ({max_attempts * delay_seconds / 60:.1f} minutes).\n\n"
    
    if lyrics_a_data and lyrics_a_data.get('status') == 'complete':
        timeout_content += f"**Version A (Ready):**\n{lyrics_a_data.get('title', 'Untitled')}\n\n{lyrics_a_data.get('text', 'No text')}\n\n"
    
    if lyrics_b_data and lyrics_b_data.get('status') == 'complete':
        timeout_content += f"**Version B (Ready):**\n{lyrics_b_data.get('title', 'Untitled')}\n\n{lyrics_b_data.get('text', 'No text')}\n\n"
    
    if not (lyrics_a_data and lyrics_a_data.get('status') == 'complete') and not (lyrics_b_data and lyrics_b_data.get('status') == 'complete'):
        timeout_content += "You can try generating lyrics again with a different prompt."
    
    await update_pending_message(
        chat_uuid,
        timeout_content,
        status='error',
        metadata={
            "lyrics_a_id": lyrics_a_id,
            "lyrics_b_id": lyrics_b_id,
            "type": "lyrics_generation_timeout",
            "async_work_in_progress": False
        }
    )

async def poll_lyrics_completion(lyrics_id: str, chat_uuid: str, token: str, max_attempts: int = 30, delay_seconds: int = 5):
    """
    Poll lyrics generation status until completion or failure.
    This function keeps checking every delay_seconds until the lyrics are complete/error or max_attempts is reached.
    """
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/json"
    }
    
    for attempt in range(max_attempts):
        try:
            logger.info(f"Polling lyrics status for {lyrics_id}, attempt {attempt + 1}/{max_attempts}")
            async with httpx.AsyncClient() as client:
                response = await client.get(
                    f"{SUNO_STUDIO_LYRICS_URL}/{lyrics_id}",
                    headers=headers,
                    timeout=30.0
                )
                response.raise_for_status()
                lyrics_data = response.json()
                
                status = lyrics_data.get('status', 'unknown')
                text = lyrics_data.get('text', '')
                title = lyrics_data.get('title', '')
                error_message = lyrics_data.get('error_message', '')
                
                if status == 'complete' and text:
                    # Success! Lyrics are ready
                    version_label = ""
                    if lyrics_id and len(lyrics_id) >= 8:
                        version_label = f" (Version {'A' if int(lyrics_id[-1], 16) % 2 == 0 else 'B'})"
                    
                    content = f"🎉 **Lyrics are ready{version_label}!**\n\n"
                    if title:
                        content += f"**Title:** {title}\n\n"
                    content += f"**Lyrics:**\n{text}\n\n"
                    if lyrics_data.get('tags'):
                        content += f"**Generated Tags:** {', '.join(lyrics_data['tags'])}\n\n"
                    content += "You can now use these lyrics to generate a full song!"
                    
                    lyrics_complete_msg = MessageSchema(
                        chat_id=chat_uuid,
                        role='assistant',
                        content=content,
                        status='complete',
                        metadata=lyrics_data
                    )
                    await append_message_to_chat(chat_uuid, lyrics_complete_msg)
                    return True
                
                elif status == 'error' or error_message:
                    # Failed
                    lyrics_error_msg = MessageSchema(
                        chat_id=chat_uuid,
                        role='assistant',
                        content=f"😞 **Lyrics generation failed:** {error_message or 'Unknown error'}. Let's try generating lyrics with a different prompt.",
                        status='error',
                        metadata=lyrics_data
                    )
                    await append_message_to_chat(chat_uuid, lyrics_error_msg)
                    return False
                
                else:
                    # Still running/pending - wait before next attempt
                    if attempt == 0:
                        # Only show the "still working" message on first check
                        lyrics_progress_msg = MessageSchema(
                            chat_id=chat_uuid,
                            role='assistant',
                            content=f"⏳ **Lyrics status:** {status}. Still working on your lyrics, this may take a minute...",
                            status='complete',
                            metadata=lyrics_data
                        )
                        await append_message_to_chat(chat_uuid, lyrics_progress_msg)
                    
                    # Wait before next attempt
                    await asyncio.sleep(delay_seconds)
                    continue
        
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 404:
                error_msg = MessageSchema(
                    chat_id=chat_uuid,
                    role='assistant',
                    content=f"😞 Lyrics generation with ID '{lyrics_id}' was not found. It may have expired or been invalid.",
                    status='error'
                )
                await append_message_to_chat(chat_uuid, error_msg)
                return False
            else:
                logger.error(f"HTTP error polling lyrics {lyrics_id}: {e}")
                # Continue trying for HTTP errors (might be temporary)
                await asyncio.sleep(delay_seconds)
                continue
        
        except Exception as e:
            logger.error(f"Unexpected error polling lyrics {lyrics_id}: {e}")
            # Continue trying for other errors
            await asyncio.sleep(delay_seconds)
            continue
    
    # Max attempts reached
    timeout_msg = MessageSchema(
        chat_id=chat_uuid,
        role='assistant',
        content=f"⏰ **Lyrics generation timeout:** The lyrics for {lyrics_id} are taking longer than expected. You can try checking the status manually later.",
        status='error'
    )
    await append_message_to_chat(chat_uuid, timeout_msg)
    return False

async def poll_song_completion(clip_id: str, chat_uuid: str, token: str, initial_content: str, initial_metadata: dict, max_attempts: int = 60, delay_seconds: int = 5):
    """
    Poll a single song clip until completion, then update the main pending message when done.
    This is simpler than the lyrics version since we usually only have 1-2 clips per song.
    """
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/json"
    }
    
    for attempt in range(max_attempts):
        try:
            logger.info(f"Polling song clip {clip_id}, attempt {attempt + 1}/{max_attempts}")
            
            # Check clip status via Suno API
            async with httpx.AsyncClient() as client:
                response = await client.get(
                    f"{SUNO_STUDIO_API_BASE_URL}/api/feed/v2?ids={clip_id}&page=2000",
                    headers=headers,
                    timeout=30.0
                )
                response.raise_for_status()
                response_data = response.json()
                
                # The feed endpoint returns an object with a "clips" array
                clips = response_data.get("clips", [])
                if len(clips) > 0:
                    clip_data = clips[0]
                else:
                    raise Exception("No clip data found in response")
                
                status = clip_data.get('status', 'unknown')
                audio_url = clip_data.get('audio_url')
                image_url = clip_data.get('image_url')
                title = clip_data.get('title', 'Untitled')
                
                # Prepare metadata update for this specific clip within the shared initial_metadata
                updated_metadata = initial_metadata.copy()
                if 'clips' in updated_metadata:
                    for clip_meta in updated_metadata['clips']:
                        if clip_meta.get('id') == clip_id:
                            clip_meta.update({
                                'status': status,
                                'audio_url': audio_url,
                                'image_url': image_url,
                                'title': title
                            })
                            break
                
                if (status == 'streaming' or status == 'complete') and audio_url:
                    if status == 'complete':
                        final_content = f"**'{title}' is ready**\n\nYour audio is complete and ready to play."
                        updated_metadata['type'] = 'song_generation_complete'
                        updated_metadata['async_work_in_progress'] = False
                        
                        await update_pending_message(
                            chat_uuid,
                            final_content,
                            status='complete',
                            metadata=updated_metadata
                        )
                        logger.info(f"Song {clip_id} completed: {title}")
                        
                        # Add follow-up suggestions after completion
                        await add_follow_up_suggestions(chat_uuid, clip_data, updated_metadata)
                        return  # Success! Stop polling for this clip.
                    else: # status == 'streaming' and audio_url
                        streaming_content = f"**'{title}' is now streaming**\n\nListen while it finishes generating."
                        updated_metadata['type'] = 'song_generation_streaming'
                        updated_metadata['async_work_in_progress'] = True
                        await update_pending_message(
                            chat_uuid,
                            streaming_content,
                            status='pending',
                            metadata=updated_metadata
                        )
                        logger.info(f"Song {clip_id} is streaming: {title}")
                        # Continue polling for 'complete' status, but use longer delay since audio is already available
                        await asyncio.sleep(10)  # Poll every 10 seconds for streaming→complete transition
                        continue

                elif status == 'error':
                    error_content = f"**Song generation failed for '{title}'**\n\nError: {clip_data.get('error_message', 'Unknown error')}"
                    updated_metadata['type'] = 'song_generation_error'
                    updated_metadata['async_work_in_progress'] = False
                    
                    await update_pending_message(
                        chat_uuid,
                        error_content,
                        status='error',
                        metadata=updated_metadata
                    )
                    logger.error(f"Song {clip_id} failed: {clip_data.get('error_message')}")
                    return  # Done (with error)
                
                else: # Still generating, or status not yet streaming/complete with audio_url
                    if attempt % 6 == 0:  # Every 30 seconds, update progress
                        progress_content = f"**Creating '{title}'**\n\nStatus: {status}. This may take a minute or two..."
                        current_progress_meta = updated_metadata.copy()
                        current_progress_meta['type'] = 'song_generation_pending'
                        current_progress_meta['async_work_in_progress'] = True
                        await update_pending_message(
                            chat_uuid,
                            progress_content,
                            status='pending',
                            metadata=current_progress_meta
                        )
                    
                    # Wait before next check
                    await asyncio.sleep(delay_seconds)
                    continue
        
        except httpx.HTTPStatusError as e:
            # Simplified error handling for HTTP errors during polling
            error_message_detail = f"HTTP {e.response.status_code}"
            if e.response and e.response.text:
                try:
                    error_json = e.response.json()
                    error_message_detail = f"HTTP {e.response.status_code}: {error_json.get('detail', e.response.text)}"
                except json.JSONDecodeError:
                    error_message_detail = f"HTTP {e.response.status_code}: {e.response.text}"

            logger.error(f"HTTP error polling song {clip_id}: {error_message_detail}")
            
            updated_metadata_on_err = initial_metadata.copy()
            if 'clips' in updated_metadata_on_err:
                for clip_meta_err in updated_metadata_on_err['clips']:
                    if clip_meta_err.get('id') == clip_id:
                        clip_meta_err['status'] = 'error'
                        clip_meta_err['audio_url'] = None
                        clip_meta_err['image_url'] = None
                        clip_meta_err['error_message'] = f"Polling failed: {error_message_detail}"
                        break
            
            updated_metadata_on_err['type'] = 'song_generation_error'
            updated_metadata_on_err['async_work_in_progress'] = False

            error_content = f"😞 **Polling error for song (ID: {clip_id[:8]})**\n\n{error_message_detail}.\n\nThe generation for this clip has stopped."
            
            await update_pending_message(
                chat_uuid,
                error_content,
                status='error',
                metadata=updated_metadata_on_err
            )
            return
        
        except Exception as e:
            logger.error(f"Unexpected error polling song {clip_id}: {e}")
            # For unexpected errors, update this clip's status in metadata to avoid infinite polling loops for this clip.
            updated_metadata_on_unexp_err = initial_metadata.copy()
            if 'clips' in updated_metadata_on_unexp_err:
                for clip_meta_unexp_err in updated_metadata_on_unexp_err['clips']:
                    if clip_meta_unexp_err.get('id') == clip_id:
                        clip_meta_unexp_err['status'] = 'error'
                        clip_meta_unexp_err['audio_url'] = None
                        clip_meta_unexp_err['image_url'] = None
                        clip_meta_unexp_err['error_message'] = f"Unexpected polling error: {str(e)}"
                        break
            
            updated_metadata_on_unexp_err['type'] = 'song_generation_error'
            updated_metadata_on_unexp_err['async_work_in_progress'] = False

            unexp_error_content = f"🚨 **Unexpected error polling song (ID: {clip_id[:8]})**\n\n{str(e)}.\n\nPolling for this clip has stopped."

            await update_pending_message(
                chat_uuid,
                unexp_error_content,
                status='error',
                metadata=updated_metadata_on_unexp_err
            )
            return
    
    # Timeout reached
    timeout_content = f"⏰ **Song generation timeout for (ID: {clip_id[:8]})**\n\nThe song is taking longer than expected ({max_attempts * delay_seconds / 60:.1f} minutes).\n\nYou can try generating a new song or check this clip ID later."
    
    updated_metadata_on_timeout = initial_metadata.copy()
    if 'clips' in updated_metadata_on_timeout:
        for clip_meta_timeout in updated_metadata_on_timeout['clips']:
            if clip_meta_timeout.get('id') == clip_id:
                clip_meta_timeout['status'] = 'timeout'
                break
                
    updated_metadata_on_timeout['type'] = 'song_generation_timeout'
    updated_metadata_on_timeout['async_work_in_progress'] = False

    await update_pending_message(
        chat_uuid,
        timeout_content,
        status='error',
        metadata=updated_metadata_on_timeout
    )

async def get_lyrics_status_from_suno_studio_api(args: dict, chat_uuid: str, token: str) -> bool:
    """Check the status of lyrics generation (single check, used as a tool by the AI)."""
    
    lyrics_id = args.get("lyrics_id", "")
    
    if not lyrics_id:
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content="I need a lyrics generation ID to check the status. Please provide the ID from a previous lyrics generation request.",
            status='error'
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False
    
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/json"
    }

    try:
        logger.info(f"Checking lyrics status for: {lyrics_id}")
        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"{SUNO_STUDIO_LYRICS_URL}/{lyrics_id}",
                headers=headers,
                timeout=30.0
            )
            response.raise_for_status()
            lyrics_data = response.json()
            
            logger.info(f"Lyrics status: {json.dumps(lyrics_data, indent=2)}")
            
            status = lyrics_data.get('status', 'unknown')
            text = lyrics_data.get('text', '')
            title = lyrics_data.get('title', '')
            error_message = lyrics_data.get('error_message', '')
            
            if status == 'complete' and text:
                # Try to identify if this is version A or B by checking the last few characters of the ID
                version_label = ""
                if lyrics_id and len(lyrics_id) >= 8:
                    # Simple heuristic to distinguish between versions
                    version_label = f" (Version {'A' if int(lyrics_id[-1], 16) % 2 == 0 else 'B'})"
                
                content = f"🎉 **Lyrics are ready{version_label}!**\n\n"
                if title:
                    content += f"**Title:** {title}\n\n"
                content += f"**Lyrics:**\n{text}\n\n"
                if lyrics_data.get('tags'):
                    content += f"**Generated Tags:** {', '.join(lyrics_data['tags'])}\n\n"
                content += "You can now use these lyrics to generate a full song!"
                
                lyrics_complete_msg = MessageSchema(
                    chat_id=chat_uuid,
                    role='assistant',
                    content=content,
                    status='complete',
                    metadata=lyrics_data
                )
                await append_message_to_chat(chat_uuid, lyrics_complete_msg)
                return True
            elif status == 'error' or error_message:
                lyrics_complete_msg = MessageSchema(
                    chat_id=chat_uuid,
                    role='assistant',
                    content=f"😞 **Lyrics generation failed:** {error_message or 'Unknown error'}. Let's try generating lyrics with a different prompt.",
                    status='error',
                    metadata=lyrics_data
                )
                await append_message_to_chat(chat_uuid, lyrics_complete_msg)
                return False
            else: # e.g. 'pending' or other non-complete, non-error states from Suno
                lyrics_complete_msg = MessageSchema(
                    chat_id=chat_uuid,
                    role='assistant',
                    content=f"⏳ **Lyrics status for {lyrics_id}:** {status}. Still working on your lyrics...",
                    status='complete',
                    metadata=lyrics_data
                )
                await append_message_to_chat(chat_uuid, lyrics_complete_msg)
                return True

    except httpx.HTTPStatusError as e:
        if e.response.status_code == 404:
            error_msg = MessageSchema(
                chat_id=chat_uuid,
                role='assistant',
                content=f"😞 Lyrics generation with ID '{lyrics_id}' was not found. It may have expired or been invalid.",
                status='error'
            )
        else:
            error_msg = MessageSchema(
                chat_id=chat_uuid,
                role='assistant',
                content=f"😞 Error checking lyrics status (HTTP {e.response.status_code}). Please try again later.",
                status='error'
            )
        await append_message_to_chat(chat_uuid, error_msg)
        return False
        
    except Exception as e:
        logger.error(f"Unexpected error checking lyrics status: {e}")
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content="😞 I encountered an unexpected issue while checking lyrics status. Please try again.",
            status='error'
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False

async def find_clip_ids_in_chat_history(args: dict, chat_uuid: str, token: str) -> bool:
    """Search chat history for clip IDs from previously generated songs."""
    
    search_term = args.get("search_term", "").lower() if args.get("search_term") else None
    
    try:
        # Get current chat history
        chat_messages = await read_chat_messages(chat_uuid)
        
        found_clips = []
        
        for message in chat_messages:
            if message.role == 'assistant' and message.metadata and message.metadata.get('clips'):
                clips = message.metadata['clips']
                for clip in clips:
                    clip_id = clip.get('id', '')
                    title = clip.get('title', 'Untitled')
                    tags = clip.get('metadata', {}).get('tags', '')
                    status = clip.get('status', 'unknown')
                    
                    # Filter by search term if provided
                    if search_term:
                        searchable_text = f"{title} {tags}".lower()
                        if search_term not in searchable_text:
                            continue
                    
                    found_clips.append({
                        'clip_id': clip_id,
                        'title': title,
                        'tags': tags,
                        'status': status,
                        'audio_url': clip.get('audio_url')
                    })
        
        if found_clips:
            if len(found_clips) == 1:
                clip = found_clips[0]
                content = f"Found: **{clip['title']}** (`{clip['clip_id'][:8]}...`)"
            else:
                content = f"Found {len(found_clips)} clips: " + ", ".join([f"**{clip['title']}** (`{clip['clip_id'][:8]}...`)" for clip in found_clips])
        else:
            if search_term:
                content = f"No clips found matching '{search_term}'."
            else:
                content = "No clips found in chat history."
        
        result_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content=content,
            status='complete',
            metadata={"found_clips": found_clips}
        )
        await append_message_to_chat(chat_uuid, result_msg)
        return True
        
    except Exception as e:
        logger.error(f"Error searching for clip IDs: {e}")
        error_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content="😞 I encountered an error while searching for clip IDs. Please try again.",
            status='error'
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False

async def add_follow_up_suggestions(chat_uuid: str, clip_data: dict, metadata: dict):
    """Add contextual follow-up suggestions after a song completes."""
    
    try:
        title = clip_data.get('title', 'your track')
        clip_metadata = clip_data.get('metadata', {})
        is_cover = clip_metadata.get('is_remix', False)
        is_extension = 'history' in clip_metadata
        tags = clip_metadata.get('tags', '')
        
        suggestions = []
        
        if is_extension:
            suggestions = [
                "Try a cover with different instruments?",
                "Make an instrumental version?",
                "Create something completely new?"
            ]
        elif is_cover:
            suggestions = [
                "Want to extend this version?",
                "Try another style cover?",
                "Create a remix with different energy?"
            ]
        else:
            # Original song
            suggestions = [
                "Make a cover with different instruments?",
                "Extend it with more sections?",
                "Try an instrumental version?",
                "Create a remix with different energy?"
            ]
        
        # Randomly pick 2-3 suggestions to keep it fresh
        selected_suggestions = random.sample(suggestions, min(3, len(suggestions)))
        
        suggestion_text = "What's next? " + " • ".join(selected_suggestions)
        
        follow_up_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content=suggestion_text,
            status='complete',
            metadata={"type": "follow_up_suggestions", "source_clip_id": clip_data.get('id')}
        )
        
        await append_message_to_chat(chat_uuid, follow_up_msg)
        
    except Exception as e:
        logger.error(f"Error adding follow-up suggestions: {e}")
        # Don't let suggestion errors break the main flow

# --- API Endpoints ---
@app.get("/api/chat/{chat_uuid}/messages", response_model=List[MessageSchema])
async def get_chat_messages(
    chat_uuid: Annotated[str, FastApiPath(description="Unique identifier for the chat session")]
):
    """
    Retrieves all messages for a given chat session.
    If the chat session doesn't exist, it will effectively be created (empty list returned, file created on first POST).
    """
    return await read_chat_messages(chat_uuid)

@app.post("/api/chat/{chat_uuid}/message", response_model=List[MessageSchema])
async def post_chat_message(
    chat_uuid: Annotated[str, FastApiPath(description="Unique identifier for the chat session")],
    background_tasks: BackgroundTasks,
    payload: UserMessagePayloadSchema = Body(...)
):
    """
    Submits a user message to a chat session and returns immediately.
    The assistant response will be processed asynchronously in the background.
    The client should poll the GET endpoint to see when new assistant messages are available.
    """
    
    # 1. Store user message
    user_message = MessageSchema(
        chat_id=chat_uuid,
        role='user',
        content=payload.content,
        status='complete'
    )
    await append_message_to_chat(chat_uuid, user_message)

    # 2. Store pending assistant message to signal that a response is being processed
    pending_message = MessageSchema(
        chat_id=chat_uuid,
        role='assistant',
        content='Orpheus is working...',  # Simple placeholder
        status='pending'
    )
    await append_message_to_chat(chat_uuid, pending_message)

    # 3. Start background task to process assistant response
    background_tasks.add_task(process_assistant_response, chat_uuid, payload.content, payload.token)

    # 4. Return immediately with current messages (including pending assistant message)
    return await read_chat_messages(chat_uuid)

async def update_pending_message(chat_uuid: str, content: str, status: str = 'pending', metadata: dict = None):
    """Update the appropriate pending assistant message in-place with new content."""
    chat_file = CHAT_HISTORY_DIR / f"{chat_uuid}.jsonl"
    
    # Read all current messages
    messages = []
    if chat_file.exists() and os.path.getsize(chat_file) > 0:
        async with aiofiles.open(chat_file, mode="r") as f:
            async for line in f:
                if line.strip():
                    try:
                        message_data = json.loads(line)
                        messages.append(MessageSchema(**message_data))
                    except (json.JSONDecodeError, Exception) as e:
                        logger.error(f"Error parsing message in {chat_uuid}: {e}")
                        continue
    
    # Find the correct assistant message to update
    updated = False
    for i in range(len(messages) - 1, -1, -1):  # Search backwards
        if messages[i].role == 'assistant':
            # Check if this is the right message to update based on metadata
            should_update = False
            
            # First check: exact generation_request_id match (prevents duplicates)
            if metadata and messages[i].metadata:
                existing_gen_id = messages[i].metadata.get('generation_request_id')
                new_gen_id = metadata.get('generation_request_id')
                if existing_gen_id and new_gen_id and existing_gen_id == new_gen_id:
                    should_update = True
                elif messages[i].status == 'pending':
                    # Only update pending messages if no exact generation_request_id match
                    # Fallback to type-based matching
                    existing_type = messages[i].metadata.get('type', '')
                    new_type = metadata.get('type', '')
                    
                    # If both have song generation types, they should match
                    if ('song_generation' in existing_type and 'song_generation' in new_type):
                        should_update = True
                    # If both have lyrics generation types, they should match  
                    elif ('lyrics_generation' in existing_type and 'lyrics_generation' in new_type):
                        should_update = True
                    # For tool flow updates, match those specifically
                    elif (messages[i].metadata.get('tool_flow_status') and 
                          metadata.get('tool_flow_status')):
                        should_update = True
                    # If we have async_work_in_progress, prefer updating that message
                    elif messages[i].metadata.get('async_work_in_progress') == True:
                        should_update = True
            elif messages[i].status == 'pending':
                # No metadata, just update any pending message
                should_update = True
            
            if should_update:
                # Update this message in-place
                messages[i].content = content
                messages[i].status = status
                messages[i].metadata = metadata or {}
                # DON'T update timestamp to preserve original message order
                updated = True
                break
    
    if not updated:
        # No appropriate pending message found, append a new one
        new_msg = MessageSchema(
            chat_id=chat_uuid,
            role='assistant',
            content=content,
            status=status,
            metadata=metadata or {}
        )
        messages.append(new_msg)
    
    # Rewrite the entire chat file with updated messages
    async with aiofiles.open(chat_file, mode="w") as f:
        for message in messages:
            message_dict = message.model_dump(mode="json")
            if isinstance(message_dict.get('timestamp'), datetime):
                message_dict['timestamp'] = message_dict['timestamp'].isoformat()
            await f.write(json.dumps(message_dict) + "\n")

async def process_assistant_response(chat_uuid: str, user_content: str, suno_token: str):
    """
    Background task to process the assistant response and any tool calls.
    This runs asynchronously after the POST endpoint returns.
    This provides real-time updates to the pending message during processing.
    """
    try:
        # Get current chat history for context (excluding pending messages)
        current_messages = await read_chat_messages(chat_uuid)
        # Filter out pending messages for context to avoid confusing the AI
        context_messages = [msg for msg in current_messages if msg.status != 'pending']
        
        # Chat with Orpheus assistant
        assistant_response = await chat_with_orpheus_assistant(user_content, context_messages)
        
        base_ai_content = assistant_response["response_text"] or "Okay, I'll work on that." # Default AI response text
        
        # Initial response from assistant (text part)
        # This message's content is the base for tool flow updates if no specific content is provided by tools.
        await update_pending_message(chat_uuid, base_ai_content)
        
        if assistant_response["tool_calls"]:
            total_tools = len(assistant_response["tool_calls"])
            has_async_work = False  # Track if any tools are doing async work
            
            for i, tool_call_data in enumerate(assistant_response["tool_calls"]):
                tool_name = tool_call_data.get("name", "unknown")
                
                # Update the primary AI response message to show tool execution via metadata
                tool_flow_metadata_executing = {
                    "status": "executing_tool",
                    "current_tool_name": tool_name,
                    "current_tool_step": i + 1,
                    "total_tools": total_tools
                }
                # Include tool name in the content for better visibility
                tool_content = f"{base_ai_content}\n\n🔧 **Executing tool {i + 1}/{total_tools}:** `{tool_name}`"
                await update_pending_message(chat_uuid, tool_content, metadata={"tool_flow_status": tool_flow_metadata_executing})
                
                # Execute the tool. The tool function itself will append its specific outcome message(s)
                # and return True for success, False for failure.
                tool_succeeded = await execute_single_tool(tool_call_data, chat_uuid, suno_token)
                
                if not tool_succeeded:
                    # The tool itself has logged its specific error message.
                    # This error message is now the last message in the chat.
                    # We stop further processing for this turn.
                    # The frontend's filterMessages will pick up the tool's error message.
                    # We can optionally log a general "tool processing aborted" message here if needed,
                    # but the specific error from the tool is more important.
                    logger.info(f"Tool {tool_name} failed. Aborting further processing for this turn.")
                    # Update the AI base message to reflect that a tool failed in the flow, if desired,
                    # but the specific error is already logged by the tool.
                    # For now, exiting early means the tool's error is the last message.
                    return 
                
                # Check if this tool started async work by looking at the latest message metadata
                current_messages = await read_chat_messages(chat_uuid)
                if current_messages:
                    latest_message = current_messages[-1]
                    if (latest_message.metadata and 
                        latest_message.metadata.get("async_work_in_progress") == True):
                        has_async_work = True
                        logger.info(f"Tool {tool_name} started async work - will not mark as complete yet")
                        # Don't update tool flow status for async tools - let them handle their own status
                        continue
                
                # If tool succeeded synchronously, update the AI base message's tool_flow_status
                tool_flow_metadata_success = {
                    "status": "tool_completed_successfully",
                    "current_tool_name": tool_name,
                    "current_tool_step": i + 1,
                    "total_tools": total_tools
                }
                success_content = f"{base_ai_content}\n\n✅ **Completed tool {i + 1}/{total_tools}:** `{tool_name}`"
                await update_pending_message(chat_uuid, success_content, metadata={"tool_flow_status": tool_flow_metadata_success})
            
            # Only mark as complete if no async work is in progress
            if not has_async_work:
                # If all tools succeeded and the loop completed:
                final_metadata = {"tool_flow_status": {"status": "all_processing_complete", "total_tools_processed": total_tools}}
                # Show summary of completed tools
                tools_summary = ", ".join([call["name"] for call in assistant_response["tool_calls"]])
                final_content = f"{base_ai_content}\n\n🎉 **All {total_tools} tools completed:** `{tools_summary}`"
                await update_pending_message(chat_uuid, final_content, status='complete', metadata=final_metadata)
            else:
                logger.info("Async work in progress - leaving message pending for background tasks to complete")
        
        else: # No tool calls, AI response was final.
            final_metadata_no_tools = {"tool_flow_status": {"status": "all_processing_complete", "total_tools_processed": 0}}
            await update_pending_message(chat_uuid, base_ai_content, status='complete', metadata=final_metadata_no_tools)
            
    except Exception as e:
        logger.error(f"Error in background assistant processing for chat {chat_uuid}: {e}")
        import traceback
        full_error = traceback.format_exc()
        error_content = f"🚨 Background Processing Error:**\n\n**Error:** {str(e)}\n\n**Full Traceback:**\n```\n{full_error}\n```"
        # This updates the last "pending" message to be an error message.
        # This error is for issues in process_assistant_response itself, not tool errors handled above.
        await update_pending_message(chat_uuid, error_content, status='error')

async def execute_single_tool(tool_call_data: dict, chat_uuid: str, token: str) -> bool:
    """
    Executes a single tool call.
    The specific tool function is responsible for appending its own outcome message(s)
    (e.g., success message, or error message with status='error').
    Returns True if the tool execution was successful, False otherwise.
    """
    tool_name = tool_call_data.get("name")
    try:
        args = json.loads(tool_call_data.get("arguments", "{}"))
    except json.JSONDecodeError as e:
        logger.error(f"Failed to parse arguments for tool {tool_name}: {tool_call_data.get('arguments')} - Error: {e}")
        error_msg = MessageSchema(
            chat_id=chat_uuid, role='assistant',
            content=f"🚨 Tool Argument Error:**\n\n**Tool:** {tool_name}\n**Error:** Invalid arguments provided (JSON parsing failed).",
            status='error'
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False

    if tool_name == "generate_song":
        return await generate_song_with_suno_studio_api(args, chat_uuid, token)
    elif tool_name == "generate_song_cover":
        # This tool also uses the same underlying Python function,
        # as it's designed to handle different tasks including "cover".
        # The OpenAI schema for "generate_song_cover" ensures 'task' is 'cover'
        # and 'cover_clip_id' is provided.
        return await generate_song_with_suno_studio_api(args, chat_uuid, token)
    elif tool_name == "get_clip_details":
        return await get_clip_details_from_suno_studio_api(args, chat_uuid, token)
    elif tool_name == "generate_lyrics":
        return await generate_lyrics_with_suno_studio_api(args, chat_uuid, token)
    elif tool_name == "get_lyrics_status":
        return await get_lyrics_status_from_suno_studio_api(args, chat_uuid, token)
    elif tool_name == "find_clip_ids_in_chat":
        return await find_clip_ids_in_chat_history(args, chat_uuid, token)
    else:
        logger.warning(f"Attempted to call unknown tool: {tool_name}")
        error_msg = MessageSchema(
            chat_id=chat_uuid, role='assistant',
            content=f"🚨 Unknown Tool:**\n\n**Tool:** {tool_name}\n**Error:** The requested tool is not recognized by the system.",
            status='error'
        )
        await append_message_to_chat(chat_uuid, error_msg)
        return False
    
@app.get("/health")
async def health_check():
    return {"status": "healthy", "message": "Orpheus is running!"}

    