import json
import os
import pathlib
import string
import subprocess
import tempfile
import threading
import time
import unicodedata
import uuid
from difflib import SequenceMatcher
from pathlib import Path
from typing import List

import boto3
import dotenv
import numpy as np
import requests
from moviepy.editor import AudioFileClip, VideoFileClip
from yt_dlp import YoutubeDL

import modal

image = (
    modal.Image.debian_slim()
    .pip_install_from_pyproject("pyproject.toml")
    .apt_install("ffmpeg")
)
app = modal.App(name="songify-a-video-app", image=image)


# S3 bucket name from environment variable
S3_BUCKET = "suno-data-uploads"


API_BASE = "https://studio-api.staging.suno.com"
FPS = 30


dotenv.load_dotenv()


def upload_video_to_s3(file_path: str, s3_file_name: str = None) -> str:
    s3_client = boto3.client(
        "s3",
        aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
        aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
    )
    bucket = "suno-data-uploads"
    s3_key = f"studio/uploads/{s3_file_name}"

    try:
        # Upload the file
        s3_client.upload_file(file_path, bucket, s3_key)

        # Generate the S3 URL
        s3_url = f"https://cdn1.suno.ai/{s3_file_name}"

        return s3_url
    except Exception as e:
        raise RuntimeError(f"Failed to upload file to S3: {str(e)}")


def _headers():
    token = os.getenv("SUNO_TOKEN")
    if not token:
        raise RuntimeError("SUNO_TOKEN env var missing")
    return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}


def upload_audio(audio_path: Path) -> str:
    # Step 1: Reserve upload slot
    payload = {"filename": audio_path.name, "content_type": "audio/wav"}

    r = requests.post(f"{API_BASE}/api/uploads/audio", json=payload, headers=_headers())

    if not r.ok:
        raise RuntimeError(f"Upload reservation failed: {r.text}")

    upload_data = r.json()
    upload_id = upload_data.get("id") or upload_data.get("upload_id")

    # Step 2: Upload file
    if "fields" in upload_data:
        # Staging flow - multipart POST
        url = upload_data["url"]
        fields = upload_data["fields"]

        with open(audio_path, "rb") as f:
            files = {
                "file": (audio_path.name, f, fields.get("Content-Type", "audio/mpeg"))
            }
            r = requests.post(url, data=fields, files=files)

        if not r.ok:
            raise RuntimeError(f"File upload failed: {r.text}")
    else:
        # Production flow - presigned PUT
        presigned_url = upload_data.get("upload_url")

        with open(audio_path, "rb") as f:
            r = requests.put(
                presigned_url, data=f, headers={"Content-Type": "audio/wav"}
            )

        if not r.ok:
            raise RuntimeError(f"File upload failed: {r.text}")

    # Step 3: Mark upload complete
    finish_payload = {
        "upload_type": "audio",
        "upload_filename": audio_path.name,
        "upload_key": upload_data.get("fields", {}).get(
            "key", f"raw_uploads/{upload_id}.mp3"
        ),
    }

    r = requests.post(
        f"{API_BASE}/api/uploads/audio/{upload_id}/upload-finish",
        json=finish_payload,
        headers=_headers(),
    )

    if not r.ok:
        raise RuntimeError(f"Upload finish failed: {r.text}")

    # Step 4: Poll for processing completion
    poll_start = time.time()

    while True:
        r = requests.get(
            f"{API_BASE}/api/uploads/audio/{upload_id}", headers=_headers()
        )

        if not r.ok:
            raise RuntimeError(f"Status poll failed: {r.text}")

        status_data = r.json()
        status = status_data.get("status", "unknown")

        if status == "complete" or status == "error":
            break

        if time.time() - poll_start > 180:
            raise TimeoutError("Upload processing timed out")

        time.sleep(3)

    if status == "error":
        raise RuntimeError(f"Upload processing failed: {status_data}")

    # Step 5: Initialize clip
    r = requests.post(
        f"{API_BASE}/api/uploads/audio/{upload_id}/initialize-clip",
        json={},
        headers=_headers(),
    )

    if not r.ok:
        raise RuntimeError(f"Clip initialization failed: {r.text}")

    clip_data = r.json()
    clip_id = clip_data.get("clip_id") or clip_data.get("id")

    if not clip_id:
        raise RuntimeError("No clip ID returned from initialization")

    return clip_id


def download_video_from_tiktok(url: str, out_dir: str) -> str:
    """
    scrapes video from tiktok and downloads it to local file system

    returns: path in local file system
    """
    out_dir = pathlib.Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    tmpl = str(out_dir / "%(id)s.%(ext)s")
    ydl_opts = {
        "outtmpl": tmpl,
        "format": "mp4/best",
        "quiet": True,
        "no_warnings": True,
    }
    with YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        return str(pathlib.Path(ydl.prepare_filename(info)).with_suffix(".mp4"))

def download_video_from_s3(s3_file_name: str, out_dir: str) -> str:
    """
    downloads video from suno s3 bucket to local file system

    returns: path to downloaded video, in local file system
    """

    s3_client = boto3.client(
        "s3",
        aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
        aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
    )

    s3_key = f"studio/uploads/{s3_file_name}"
    output_path = f"{out_dir}/{s3_key.split('/')[-1]}"
    
    # Create the output directory if it doesn't exist
    os.makedirs(out_dir, exist_ok=True)
    
    s3_client.download_file(S3_BUCKET, s3_key, output_path)
    return output_path

def extract_audio(video_path: str, out_dir: str) -> str:
    id = str(uuid.uuid4())
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        video_path,
        "-ac",
        "1",  # mono
        "-ar",
        "16000",  # 16 kHz sample‑rate
        "-c:a",
        "pcm_s16le",  # 16‑bit little‑endian PCM (valid WAV)
        str(f"{out_dir}/{id}.wav"),
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return Path(f"{out_dir}/{id}.wav")


TIMEOUT = 500


def get_lyrics(clip_id: str):
    timer = time.time()
    is_timeout = False
    while True and not is_timeout:
        lyrics = requests.get(
            f"{API_BASE}/api/gen/{clip_id}/aligned_lyrics/v2/", headers=_headers()
        ).json()
        try:
            if lyrics["aligned_lyrics"]:
                return lyrics
        except Exception as e:
            print(e)
            time.sleep(0.5)
            is_timeout = (time.time() - timer) > TIMEOUT
    if is_timeout:
        raise (TimeoutError)


def generate_detailed_tags(style_genre: str) -> str:
    """
    Generate detailed tags that explicitly mention vocals and describe the style

    Args:
        style_genre: Basic genre like "drum and bass", "house music", etc.

    Returns:
        Detailed tag description mentioning vocals
    """
    # Map genres to detailed descriptions that mention vocals
    genre_descriptions = {
        "drum and bass": "This drum n bass cover launches straight into rapid-fire percussion and a driving bassline and clear vocals from the first second. Syncopated synth stabs and crisp hi-hats layer with energetic vocals, while doubled harmonies and subtle FX sweeps keep momentum high.",
        "house": "This house music cover features a four-on-the-floor beat with clear vocals prominently mixed. Deep basslines and warm synth pads support the vocal performance, while classic house piano stabs and filtered sweeps create movement. The vocals remain crisp and upfront throughout.",
        "dubstep": "This dubstep cover opens with clear vocals over sparse atmospherics before the massive sub-bass drops. Syncopated rhythms and wobbling basslines frame the vocal performance, while glitchy effects and dramatic builds maintain energy. The vocals cut through the heavy bass with clarity.",
        "jazz": "This smooth jazz cover features clear vocals supported by warm piano chords and walking basslines. Brushed drums and subtle horn sections complement the vocal melody, while sophisticated harmonies add depth. The intimate vocal delivery remains the centerpiece throughout.",
        "epic orchestral": "This epic orchestral cover features clear vocals supported by sweeping string sections and dramatic brass. Timpani rolls and woodwind flourishes punctuate the vocal phrases, while the full orchestra swells during emotional peaks. The vocals soar above the symphonic arrangement.",
        "classical symphony": "This classical symphony cover presents clear vocals with refined orchestral accompaniment. Elegant string arrangements and delicate woodwinds support the vocal line, while dynamic contrasts highlight the emotional content. The vocals maintain prominence throughout the sophisticated arrangement.",
        "heavy metal": "This heavy metal cover features powerful vocals over crushing guitar riffs and thunderous drums. Palm-muted chugs and soaring lead guitars frame the vocal performance, while double-bass drumming drives the intensity. The vocals cut through the wall of sound with commanding presence.",
        "country": "Country western ballad clear vocals, strummin the guitar",
    }

    # Get description or create a generic one
    style_lower = style_genre.lower().strip()
    if style_lower in genre_descriptions:
        return genre_descriptions[style_lower]
    else:
        # Generic description that mentions vocals
        return f"This {style_genre} cover features clear and prominent vocals throughout. The instrumental arrangement supports and enhances the vocal performance while maintaining the characteristic elements of {style_genre}. The vocals remain crisp and intelligible, perfectly balanced with the backing track."


def create_remix(clip_id: str, audio_weight: float = 0.5, genre: str = "dubstep"):
    """create remix/cover from clip_id"""

    lyrics = get_lyrics(clip_id)
    words = [item["word"] for item in lyrics["aligned_words"]]

    joined_text = "".join(words)

    # Enable remixes for the clip
    requests.post(f"{API_BASE}/api/gen/{clip_id}/enable_remixes", headers=_headers())

    # Set remix type
    requests.post(
        f"{API_BASE}/api/gen/{clip_id}/update_remix_type",
        json={"type": "REMIX"},
        headers=_headers(),
    )

    # Generate the remix/cover

    generation_payload = {
        "prompt": joined_text,
        "generation_type": "TEXT",
        "tags": generate_detailed_tags(genre),
        "negative_tags": "instrumental, spoken word",
        "mv": "chirp-bluejay-t2",
        "task": "cover",
        "cover_clip_id": clip_id,
        "metadata": {
            "control_sliders": {
                "audio_weight": audio_weight,
                "style_weight": 0.5,
                "weirdness_constraint": 0.0,
            },
            "is_remix": True,
        },
    }

    r = requests.post(
        f"{API_BASE}/api/generate/v2", json=generation_payload, headers=_headers()
    )

    if not r.ok:
        raise RuntimeError(f"Generation failed: {r.text}")

    gen_data = r.json()
    gen_id = gen_data.get("id")

    # Poll for completion
    final_clips = []
    poll_start = time.time()

    while True:
        time.sleep(5)

        r = requests.get(
            f"{API_BASE}/api/generate/requests?ids={gen_id}", headers=_headers()
        )

        if not r.ok:
            raise RuntimeError(f"Status poll failed: {r.text}")

        status_data = r.json()
        if status_data and len(status_data) > 0:
            gen_status = status_data[0]
            status = gen_status.get("status", "unknown")

            if status in ["complete", "streaming", "error", "failed"]:
                final_clips = gen_status.get("clips", [])
                break

        if time.time() - poll_start > 300:
            raise TimeoutError("Generation timed out")

    if status in ["error", "failed"]:
        raise RuntimeError(f"Generation failed: {gen_status}")

    return final_clips[0]["id"], final_clips[1]["id"]


def bezier_interp(t, xs, ys):
    """
    Piecewise cubic Bézier interpolation through points (xs, ys),
    with endpoint slopes chosen by finite differences for monotonicity.
    """
    # precompute slopes at each sample
    n = len(xs)
    slopes = np.empty(n, float)
    # endpoint slopes
    slopes[0] = (ys[1] - ys[0]) / (xs[1] - xs[0])
    slopes[-1] = (ys[-1] - ys[-2]) / (xs[-1] - xs[-2])
    # interior slopes
    for i in range(1, n - 1):
        slopes[i] = (ys[i + 1] - ys[i - 1]) / (xs[i + 1] - xs[i - 1])

    # find which segment t falls into
    idx = np.searchsorted(xs, t) - 1
    idx = np.clip(idx, 0, n - 2)

    x0, x1 = xs[idx], xs[idx + 1]
    y0, y1 = ys[idx], ys[idx + 1]
    m0, m1 = slopes[idx], slopes[idx + 1]
    h = x1 - x0
    u = (t - x0) / h

    # Bézier control points in y:
    c1 = y0 + m0 * h / 3
    c2 = y1 - m1 * h / 3

    # de Casteljau for cubic Bézier at u
    return (
        (1 - u) ** 3 * y0
        + 3 * (1 - u) ** 2 * u * c1
        + 3 * (1 - u) * u**2 * c2
        + u**3 * y1
    )


def smooth_monotonic(arr: np.ndarray, window: int) -> np.ndarray:
    """Smooth array while maintaining monotonicity."""
    arr = np.maximum.accumulate(arr)
    if window > 1 and len(arr) > window:
        kernel = np.ones(window) / window
        pad = window // 2
        padded = np.pad(arr, (pad, pad), mode="edge")
        arr = np.convolve(padded, kernel, mode="same")[pad:-pad]
    return arr


def warp_video_with_timestamps(
    video_path, output_path, timestamps, audio_path=None, smooth_win=3
):
    """
    Warp video timing based on timestamp pairs and optionally add audio.

    Args:
        video_path: Path to input video
        output_path: Path to output video
        timestamps: List of tuples (original_time, warped_time) in seconds
        audio_path: Optional path to audio file to add (unwarped)
        smooth_win: Window size for smoothing the time mapping
    """
    # Convert timestamps to numpy arrays
    timestamps = sorted(timestamps, key=lambda x: x[1])  # Sort by warped time
    original_times = np.array([t[0] for t in timestamps])
    warped_times = np.array([t[1] for t in timestamps])

    # Remove duplicates in warped times
    mask = np.append([True], np.diff(warped_times) > 1e-6)
    warped_times = warped_times[mask]
    original_times = original_times[mask]

    # Smooth the original times to ensure monotonicity
    original_times = smooth_monotonic(original_times, smooth_win)

    # Load video
    video = VideoFileClip(video_path)

    # Determine output duration
    if audio_path:
        audio = AudioFileClip(audio_path)
        output_duration = audio.duration
    else:
        output_duration = warped_times[-1]

    # Add boundary points if needed
    if warped_times[0] > 0:
        warped_times = np.insert(warped_times, 0, 0.0)
        original_times = np.insert(original_times, 0, original_times[0])

    if warped_times[-1] < output_duration:
        warped_times = np.append(warped_times, output_duration)
        original_times = np.append(original_times, original_times[-1])

    # Define time mapping function
    def map_time(t):
        return float(np.interp(t, warped_times, original_times))

    # Apply time warping
    warped_video = video.fl_time(map_time).set_duration(output_duration)

    # Add audio if provided
    if audio_path:
        warped_video = warped_video.set_audio(audio)

    # Write output video
    warped_video.write_videofile(
        output_path,
        fps=FPS,
        codec="libx264",
        audio_codec="aac" if audio_path else None,
        threads=4,
    )

    # Clean up
    video.close()
    if audio_path:
        audio.close()


def generate_video(
    song_id: str, mp4: str, temp_dir: str, uuid: str, warp_timestamps: list
):
    while True:
        try:
            song_url = f"https://cdn1.suno.ai/{song_id}.mp3"
            file_path = os.path.join(temp_dir.name, f"{uuid}_song_audio.mp3")
            response = requests.get(song_url)
            response.raise_for_status()  # Raises an exception for bad status codes

            with open(file_path, "wb") as f:
                f.write(response.content)
            break
        except Exception as e:
            print(e)
    warp_video_with_timestamps(
        video_path=mp4,
        output_path=f"{temp_dir.name}/songify_video_{uuid}.mp4",
        timestamps=warp_timestamps,
        audio_path=f"{temp_dir.name}/{uuid}_song_audio.mp3",
    )
    upload_video_to_s3(
        f"{temp_dir.name}/songify_video_{uuid}.mp4", f"songify_video_{uuid}.mp4"
    )
    print(f"Uploaded video to S3: {f'songify_video_{uuid}.mp4'}")


def create_remix_pairs(
    original_lyrics: dict,
    clip_id: str,
    mp4: str,
    temp_dir: str,
    uuids: List[str],
    genre: str,
):
    print(f"Creating remix pairs for genre: {genre}")
    remix1, remix2 = create_remix(clip_id, genre=genre)
    lyrics_remix1 = get_lyrics(remix1)
    lyrics_remix2 = get_lyrics(remix2)
    print(f"Lyrics for remix1: {lyrics_remix1}")
    print(f"Lyrics for remix2: {lyrics_remix2}")
    RATIO_MIN, RATIO_MAX = 0.3, 3.0

    def _norm(txt: str) -> str:
        tbl = str.maketrans("", "", string.punctuation + '‘’“”"…—–')
        return unicodedata.normalize("NFKD", txt).translate(tbl).lower().strip()

    def get_anchors(old_alignment, new_alignment):
        v_words = [w["text"] for w in old_alignment]

        s_words = [w["text"] for w in new_alignment]
        sm = SequenceMatcher(a=v_words, b=s_words, autojunk=False)

        anchors = []
        prev_vs = prev_ss = -1.0
        for tag, i1, i2, j1, j2 in sm.get_opcodes():
            if tag != "equal":
                continue
            for vi, sj in zip(range(i1, i2), range(j1, j2)):
                vw, sw = old_alignment[vi], new_alignment[sj]
                vs0, ve0 = vw["start_s"], vw["end_s"]
                ss0, se0 = sw["start_s"], sw["end_s"]
                if vs0 <= prev_vs or ss0 <= prev_ss:
                    continue
                ratio = (se0 - ss0) / max(ve0 - vs0, 1e-6)
                if RATIO_MIN <= ratio <= RATIO_MAX:
                    anchors.append((vs0, ss0))
                    anchors.append((ve0, se0))
                    prev_vs, prev_ss = vs0, ss0
        if not anchors:
            raise RuntimeError("No anchors aligned – check transcripts.")
        return anchors

    def make_cleaned_lyrics(lyrics: dict):
        return [
            (word | {"text": _norm(word["text"])})
            for section in lyrics["aligned_lyrics"]
            for word in section["words"]
        ]

    clean_original_lyrics = make_cleaned_lyrics(original_lyrics)
    clean_remix_lyrics1 = make_cleaned_lyrics(lyrics_remix1)
    clean_remix_lyrics2 = make_cleaned_lyrics(lyrics_remix2)
    deduped1 = get_anchors(clean_original_lyrics, clean_remix_lyrics1)
    deduped2 = get_anchors(clean_original_lyrics, clean_remix_lyrics2)
    threads = []
    threads.append(
        threading.Thread(
            target=generate_video, args=(remix1, mp4, temp_dir, uuids[0], deduped1)
        )
    )
    threads.append(
        threading.Thread(
            target=generate_video, args=(remix2, mp4, temp_dir, uuids[1], deduped2)
        )
    )
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()


@app.function(
    secrets=[
        modal.Secret.from_name("aws-bucket"),
        modal.Secret.from_name("jason-backend-api-key"),
    ]
)
@modal.fastapi_endpoint(
    docs=True  # adds interactive documentation in the browser
)
def songify_a_video(
    url: str = "https://www.tiktok.com/@miryummi/video/7532607674991512862",
    url_type: str = "tiktok",
    genres: str = None,
    uuids: str = None,
):
    """
    Generate two songify videos per user-selected genre.
    Returns list of video uuids when generation is complete.

    Args:
        url: URL of user provided video 
        url_type: "tiktok" or "s3"
        genres: List of user-selected genres
        uuids: List of uuid per video, pre-generated in frontend for continuous status polling
    """
    # MODAL logging
    print("Starting songify_a_video")
    print(f"url: {url}")
    print(f"url_type: {url_type}")
    print(f"genres: {genres}")
    print(f"uuids: {uuids}")

    # pre-process inputs, TODO: add error handling for missing uuids and genres
    uuids = json.loads(uuids) if uuids else None 
    if genres:
        genres = json.loads(genres)
        assert len(genres) * 2 == len(uuids)
    else:
        genres = ["dubstep"]
    print(f"uuids: {uuids}")

    # Create a temporary directory
    temp_dir = tempfile.TemporaryDirectory()
    print(f"Temporary directory created at: {temp_dir.name}")

    # Branched download logic for tiktok and s3
    if url_type == "tiktok":
        print("Downloading tiktok video")
        mp4 = download_video_from_tiktok(url, temp_dir.name)
        print(f"Video downloaded to: {mp4}")
    elif url_type == "s3":
        print("Downloading s3 video")
        mp4 = download_video_from_s3(url, temp_dir.name)
        print(f"Video downloaded to: {mp4}")
    else:
        raise ValueError(f"Invalid url_type: {url_type}")

    print("Extracting audio")
    audio = extract_audio(mp4, temp_dir.name)
    print(f"Audio extracted to: {audio}")

    print("Uploading audio")
    clip_id = upload_audio(Path(audio))
    print(f"Audio uploaded to: {clip_id}")

    print("Getting original lyrics")
    original_lyrics = get_lyrics(clip_id)
    print(f"Original lyrics: {original_lyrics}")

    print("Creating remix pairs")
    all_videos = []

    threads = []
    print(genres)
    print(uuids)
    for i, genre in enumerate(genres):
        if not uuids:
            this_uuids = [str(uuid.uuid4()) for _ in range(2)]
        else:
            this_uuids = uuids[i * 2 : (i + 1) * 2]
        all_videos += this_uuids
        threads.append(
            threading.Thread(
                target=create_remix_pairs,
                args=(original_lyrics, clip_id, mp4, temp_dir, this_uuids, genre),
            )
        )
    for thread in threads:
        thread.start()

    for thread in threads:
        thread.join()

    return all_videos