import asyncio
import uuid
import modal
import json
from pydantic import BaseModel
from typing import Any, Optional, List
import aiohttp
import time
import threading
import subprocess
import os
from suno_utils.audio import Audio
import numpy as np
import random
import subprocess
import os
import queue
import threading
import time
import pygame

STARTER_CLIP = "/Users/zoeshleifer/dj-suno/the-backend/STARTER_CLIP.mp3"
RAND_TAG_LIST = [
    'acoustic',
    'aggressive',
    'anthemic',
    'atmospheric',
    'bouncy',
    'chill',
    'dark',
    'dreamy',
    'electronic',
    'emotional',
    'epic',
    'experimental',
    'futuristic',
    'groovy',
    'heartfelt',
    'infectious',
    'melodic',
    'mellow',
    'powerful',
     'psychedelic',
    'romantic',
    'smooth',
    'syncopated',
    'uplifting',
    'afrobeat',
    'anime',
    'ballad',
    'bedroom pop',
    'bluegrass',
    'blues',
    'classical',
    'country',
    'cumbia',
    'dance',
    'delta blues',
    'electropop',
    'disco',
    'drum and bass',
    'edm',
    'emo',
    'folk',
    'funk',
    'future bass',
    'gospel',
    'grunge',
    'grime',
    'hip hop',
    'house',
    'indie',
    'j-pop',
    'jazz',
    'k-pop',
    'kids music',
    'metal',
    'new jack swing',
    'new wave',
    'opera',
    'punk',
    'raga',
    'rap',
    'reggae',
    'reggaeton',
    'rock',
    'rumba',
    'salsa',
    'samba',
    'sertanejo',
    'soul',
    'synthpop',
    'swing',
     'synthwave',
     'techno',
     'trap',
     'uk garage'
];


async def call_modal_orchestrator(model_name: str, queue_item):
    """Minimal version of the Modal call for extends"""

    env_name = "dev"  # or "dev"
    modal_fn = f"orchestrator-{env_name}/ChatGptStub.moderate_and_generate"

    # Split the modal function name
    modal_stub, modal_call = modal_fn.split("/")
    modal_cls = modal_call.split(".")[0]
    modal_method = modal_call.split(".")[1]

    # Get the Modal function
    f = getattr(modal.Cls.lookup(modal_stub, modal_cls)(), modal_method)

    # Call Modal with the queue item
    fn = await f.spawn.aio(queue_item.model_dump_json())
    return fn


class QueueItem(BaseModel):
    id: str  # clip ID or request ID
    prompt_audio: Optional[str]  # S3 ID of audio to extend from
    prompt_text: Optional[str]  # lyrics/prompt text
    metadata: Any  # All the extend metadata
    callback_url: Optional[str]  # Where to post results
    model_name: Optional[str]  # e.g., "chirp-v4"
    title: Optional[str]  # Clip title
    ids: Optional[List[str]]  # List of clip IDs if batch_all=True


def create_extend_metadata(continue_clip, params):
    """Create the minimal metadata needed for extend"""

    # Copy parent clip metadata, removing certain keys
    metadata = continue_clip.metadata.copy()

    # Remove keys that shouldn't be inherited
    keys_to_remove = [
        "duration",
        "refund_credits",
        "experiment",
        "param_experiment",
        "experiment_version",
    ]
    for key in keys_to_remove:
        metadata.pop(key, None)

    # Add extend-specific metadata
    metadata.update(
        {
            "continued_from_prompt": continue_clip.metadata.get("prompt"),
            "continued_aligned_prompt": params.get("continued_aligned_prompt", ""),
            "history": continue_clip.metadata.get("history", [])
            + [
                {
                    "id": continue_clip.s3_id or str(continue_clip.id),
                    "continue_at": params["continue_at"],  # timestamp in seconds
                    "type": continue_clip.metadata.get("type", "gen"),
                    "source": "web",
                }
            ],
            "continue_at": params["continue_at"],
            "task": "extend",  # or "upload_extend"
            "edited_clip_id": str(continue_clip.id),
            "type": "gen",
            "source": "web",
            "prompt": params["prompt"],
            "tags": params.get("tags"),
            "gpt_description_prompt": params.get("gpt_description_prompt"),
            "user_id": params["user_id"],
        }
    )

    return metadata


async def extend_clip_minimal(continue_clip, extend_params):
    """Minimal extend implementation with correct batch_size"""
    # print(f"DEBUG: Starting extend_clip_minimal")
    request_id = str(uuid.uuid4())
    # print(f"DEBUG: Request ID: {request_id}")

    # Generate 2 clip IDs (Suno always generates 2)
    clip_ids = [str(uuid.uuid4()), str(uuid.uuid4())]
    # print(f"DEBUG: Generated clip IDs: {clip_ids}")

    # 1. Create metadata
    # print(f"DEBUG: Creating metadata...")
    metadata = create_extend_metadata(continue_clip, extend_params)
    print(f"DEBUG: Metadata created: {metadata}")

    # 2. Create QueueItem with 2 IDs
    # print(f"DEBUG: Creating QueueItem...")
    queue_item = QueueItem(
        id=request_id,  # Request ID
        prompt_audio=f"{continue_clip.s3_id}.mp3",
        prompt_text=extend_params["prompt"],
        metadata=metadata,
        callback_url="http://platypus.han-mahi.ts.net",
        model_name="chirp-v4",
        title=extend_params.get("title", "Extended Song"),
        ids=clip_ids,  # ← Must be 2 UUIDs
    )
    # print(f"DEBUG: QueueItem created: {queue_item.model_dump()}")
    
    # 3. Call Modal
    # print(f"DEBUG: About to call modal orchestrator...")
    result = await call_modal_orchestrator("chirp-v4", queue_item)
   # print(f"DEBUG: Modal orchestrator returned: {result}")
    return clip_ids, result  # Return both IDs to track both clips


async def poll_for_mp3(s3_url):
    start_time = time.time()
    attempt_count = 0

    async with aiohttp.ClientSession() as session:
        while True:
            attempt_count += 1
            elapsed_time = time.time() - start_time

            try:
                #print(f"Attempt {attempt_count}: Checking {s3_url}")

                # Try GET request with range to check if file exists without downloading everything
                headers = {"Range": "bytes=0-1023"}  # Just get first 1KB to test
                async with session.get(s3_url, headers=headers) as response:
                    # Check for successful response (200 or 206 for partial content)
                    if response.status in [200, 206]:
                        total_elapsed = time.time() - start_time
                        print(f"MP3 file found at {s3_url}")
                        print(
                            f"Time elapsed: {total_elapsed:.1f} seconds ({attempt_count} attempts)"
                        )

                        # Now download the full file
                        async with session.get(s3_url) as download_response:
                            if download_response.status == 200:
                                content = await download_response.read()

                                # Generate filename based on timestamp
                                timestamp = int(time.time())
                                original_filename = f"extended_song_{timestamp}.mp3"

                                # Save original file temporarily
                                with open(original_filename, "wb") as f:
                                    f.write(content)
                                
                                return original_filename

                            else:
                                print(
                                    f"Failed to download full MP3: {download_response.status}"
                                )
                    # else:
                    #     print(f"File not ready yet (status {response.status})")

            except Exception as e:
                print(
                    f"Error polling URL (attempt {attempt_count}, {elapsed_time:.1f}s elapsed): {e}"
                )
            await asyncio.sleep(5)



SAMPLE_RATE = 24000
MIN_LENGTH = 90
MIN_VOLUME = 0.5
def get_mp3(s3_id):
    return Audio.from_s3(id_1, sample_rate = SAMPLE_RATE)

def first_silence(aud):
    window_size = SAMPLE_RATE
    threshold = MIN_VOLUME
    
    # # Compute the moving average using convolution
    # cumsum = np.cumsum(np.insert(aud.array, 0, 0)) 
    # window_avgs = (cumsum[window_size:] - cumsum[:-window_size]) / window_size
    
    # # Find first index where the average is below the threshold
    # indices = np.where(window_avgs < threshold)[0]
    # first_index = len(aud.array)/SAMPLE_RATE #max gen length
    # if len(indices) > 1:
    #     first_index = indices[1]
    # return first_index
    return len(aud.array)/SAMPLE_RATE
def extend(tags, continue_clip):
    print(f"DEBUG: Starting extend with tags: {tags}")
    
    # Create extend parameters with more complete configuration
    extend_params = {
        "continue_at": 90.0,  # Continue at 90 seconds
        "prompt": "Add a dramatic breakdown with heavy bass drop",
        "user_id": "sunoforever",
        "tags": tags,  # Include both fixed and random tags
        "negative_tags": ["pop", "chill", "synth", "silent"],
        "title": " ".join(tags),
        "continued_aligned_prompt": "The song continues with a powerful breakdown section",
        "gpt_description_prompt": "Create a dramatic breakdown that builds tension before a massive bass drop",
    }

    # print("Mock objects created:")
    # print(f"Continue clip ID: {continue_clip.id}")
    # print(f"Continue clip S3 ID: {continue_clip.s3_id}")
    # print(f"Original prompt: {continue_clip.metadata['prompt']}")
    # print(f"Extend prompt: {extend_params['prompt']}")
    # print(f"Continue at: {extend_params['continue_at']} seconds")
    
   # print(f"DEBUG: About to call extend_clip_minimal...")
    result = asyncio.run(extend_clip_minimal(continue_clip, extend_params))
    print(f"DEBUG: extend_clip_minimal returned: {result[0][0]}")

    # Extract S3 URL from result and poll for MP3
    if result and len(result) > 0:
      #  print(f"DEBUG: Result has {len(result)} items")
        s3_url = f"https://cdn1.suno.ai/{result[0][0]}.mp3"
      #  print(f"Polling for MP3 at: {s3_url}")
        
      #  print(f"DEBUG: About to poll for MP3...")
        downloaded_file = asyncio.run(poll_for_mp3(s3_url))
      #  print(f"DEBUG: poll_for_mp3 returned: {downloaded_file}")
        
        if downloaded_file:
      #      print(f"DEBUG: About to process audio file: {downloaded_file}")
            aud = Audio.from_file(downloaded_file)
            silence_time = first_silence(aud)
      #      print(f"DEBUG: First silence at: {silence_time}, MIN_LENGTH: {MIN_LENGTH}")
            
            if silence_time > MIN_LENGTH:
                print(f"DEBUG: Returning downloaded file: {downloaded_file}")
                return downloaded_file, result[0][0]
            else:
                print("DEBUG: Audio too short, returning starter clip")
                return STARTER_CLIP, continue_clip.s3_id
        else:
      #      print("DEBUG: Downloaded file is None, returning starter clip")
            return STARTER_CLIP, continue_clip.s3_id
    else:
       # print(f"DEBUG: No result or empty result: {result}")
        #print("No result received from extend operation")
        return STARTER_CLIP, continue_clip.s3_id



# Initialize mixer
pygame.mixer.init()

# Queue to hold MP3 file paths
song_queue = queue.Queue()


def read_tags():
    """Read tags from tag.txt file (up to 5 lines), with fallback to random sampling."""
    try:
        with open("/Users/zoeshleifer/dj-suno/the-backend/tag.txt", 'r') as f:
            tag_lines = [line.strip() for line in f.readlines() if line.strip()]
            if tag_lines:
                tags = tag_lines[:5]  # Take up to 5 lines
                print(f"Using tags from tag.txt: {tags}")
                return tags
            else:
                # Fallback to random sampling if file is empty
                tags = random.sample(RAND_TAG_LIST, 5)
                print(f"tag.txt is empty, using random tags: {tags}")
                return tags
    except FileNotFoundError:
        # Fallback to random sampling if file doesn't exist
        tags = random.sample(RAND_TAG_LIST, 5)
        print(f"tag.txt not found, using random tags: {tags}")
        return tags
    except Exception as e:
        # Fallback to random sampling for any other error
        tags = random.sample(RAND_TAG_LIST, 5)
        print(f"Error reading tag.txt ({e}), using random tags: {tags}")
        return tags


def player():
    """Plays the first 90 seconds of each queued MP3 file."""
    while True:
        song_path = song_queue.get()
        if not os.path.isfile(song_path):
            print(f"File not found: {song_path}")
            continue

        try:
            pygame.mixer.music.load(song_path)
            pygame.mixer.music.play()
            print(f"Now playing first 90 seconds of: {song_path}")

            start_time = time.time()
            while pygame.mixer.music.get_busy():
                if time.time() - start_time > 90:
                    pygame.mixer.music.stop()
                    print("Stopped after 90 seconds.")
                    break
                time.sleep(0.5)

        except Exception as e:
            print(f"Playback error: {e}")


if __name__ == "__main__":

    class StarterClip:
        def __init__(self):
            self.id = "clip_123"
            self.s3_id = "798af9d3-f4f8-4c84-b07d-918bc378e119"
            self.metadata = {
                "prompt": "Sax Instrumental",
                "history": [],
                "type": "gen",
                "duration": 120.0,
                "tags": ["jazz", "sax", "instrumental", "beat"],
                "make_instrumental": 1,
            }
    song_queue.put(STARTER_CLIP)
    def clip_generator():
        continue_clip = StarterClip()
        
        def prepare_next_song(tags, clip):
            next_song_file, s3_id = extend(tags, clip)
            song_queue.put(next_song_file)
            return s3_id
        
        while True:
            tags = read_tags()
            print(f"********* next tags: {tags} *********")
            s3_id = prepare_next_song(tags, continue_clip)
            continue_clip.s3_id = s3_id
            continue_clip.metadata["tags"] = tags
            continue_clip.metadata["prompt"] = " ".join(tags)
    
    threading.Thread(target=clip_generator, daemon=True).start()
    player()
   