# ================================================================
# 🎛️ CIPHER-MD AUDIO ENGINE v3.0
# Requirements: pip install fastapi uvicorn python-multipart aiofiles
# ================================================================

import os
import subprocess
import shutil
import time
import uuid
import asyncio
from concurrent.futures import ThreadPoolExecutor
from fastapi import FastAPI, File, UploadFile, Form, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse, JSONResponse
import aiofiles

app = FastAPI(title="Cipher-MD Audio Engine", version="3.0")

# 📂 Folder Setup
UPLOAD_FOLDER = "uploads"
PROCESSED_FOLDER = "processed"
ASSETS_FOLDER = "assets"

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(PROCESSED_FOLDER, exist_ok=True)
os.makedirs(ASSETS_FOLDER, exist_ok=True)

# Thread pool for blocking ffmpeg calls
executor = ThreadPoolExecutor(max_workers=4)

# ================================================================
# 🛠️ HELPERS
# ================================================================

def get_audio_duration(file_path: str) -> float:
    """Get audio duration in seconds using ffprobe"""
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        file_path
    ]
    try:
        output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL).decode().strip()
        return float(output)
    except Exception:
        return 0.0


def get_audio_info(file_path: str) -> dict:
    """Get detailed audio metadata using ffprobe"""
    cmd = [
        "ffprobe", "-v", "quiet", "-print_format", "json",
        "-show_format", "-show_streams", file_path
    ]
    try:
        import json
        output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL).decode()
        data = json.loads(output)
        fmt = data.get("format", {})
        tags = fmt.get("tags", {})
        streams = data.get("streams", [])
        audio_stream = next((s for s in streams if s.get("codec_type") == "audio"), {})
        return {
            "duration": float(fmt.get("duration", 0)),
            "size_bytes": int(fmt.get("size", 0)),
            "bit_rate": int(fmt.get("bit_rate", 0)),
            "format": fmt.get("format_name", ""),
            "codec": audio_stream.get("codec_name", ""),
            "sample_rate": audio_stream.get("sample_rate", ""),
            "channels": audio_stream.get("channels", 0),
            "title": tags.get("title", ""),
            "artist": tags.get("artist", ""),
        }
    except Exception as e:
        return {"error": str(e)}


def cleanup_file(path: str):
    """Background task to clean up temp files"""
    try:
        if os.path.exists(path):
            os.remove(path)
    except Exception:
        pass


def build_effect_filter(effect_type: str, level: int) -> str:
    """
    Returns ONLY the audio effect filter string.
    level: 1-20 (bass boost amount, ignored for non-bass effects)
    """
    level = max(1, min(20, level))  # Clamp 1-20

    # Bass level to safe EQ gain (prevent clipping — max usable EQ gain ~15dB)
    # level 1-20 maps to EQ gain 2-15dB with soft ceiling
    bass_gain = round(2 + (level - 1) * (13 / 19), 1)   # 2dB..15dB

    effects = {
        # ─── BASS SERIES ─────────────────────────────────────────
        # ✅ EQ applied to FULL signal (no split/amix) → song never disappears
        # Chain: compressor → bass shelf → sub EQ → treble restore → limiter
        "bass_studio": (
            f"acompressor=threshold=-18dB:ratio=3:attack=5:release=50:makeup=1,"
            f"bass=g={bass_gain}:f=80:w=0.5,"
            f"equalizer=f=45:width_type=h:width=60:g={round(bass_gain*0.5, 1)},"
            f"equalizer=f=200:width_type=h:width=100:g=1,"
            f"alimiter=limit=0.95:level=1.0"
        ),
        "bass_edm": (
            f"acompressor=threshold=-16dB:ratio=4:attack=3:release=40:makeup=1,"
            f"bass=g={bass_gain}:f=90:w=0.6,"
            f"equalizer=f=55:width_type=h:width=80:g={round(bass_gain*0.6, 1)},"
            f"equalizer=f=120:width_type=h:width=100:g=2,"
            f"equalizer=f=8000:width_type=h:width=2000:g=2,"
            f"alimiter=limit=0.95:level=1.0"
        ),
        "bass_car": (
            # Car sub = deep sub (40-60Hz) heavy emphasis
            f"acompressor=threshold=-14dB:ratio=5:attack=3:release=30:makeup=1,"
            f"bass=g={bass_gain}:f=60:w=0.8,"
            f"equalizer=f=40:width_type=h:width=40:g={round(bass_gain*0.7, 1)},"
            f"equalizer=f=80:width_type=h:width=60:g={round(bass_gain*0.4, 1)},"
            f"equalizer=f=250:width_type=h:width=80:g=-2,"
            f"alimiter=limit=0.94:level=1.0"
        ),
        "bass_slap": (
            # Slap bass = punchy mid-bass (100-200Hz) with snap
            f"acompressor=threshold=-18dB:ratio=3:attack=2:release=20:makeup=1,"
            f"bass=g={bass_gain}:f=150:w=1.0,"
            f"equalizer=f=100:width_type=h:width=80:g={round(bass_gain*0.5, 1)},"
            f"treble=g=3,"
            f"alimiter=limit=0.95:level=1.0"
        ),
        "bass_heavy": (
            # Maximum sub-bass — careful with speakers!
            f"acompressor=threshold=-12dB:ratio=6:attack=2:release=25:makeup=1,"
            f"bass=g={bass_gain}:f=70:w=0.7,"
            f"equalizer=f=40:width_type=h:width=50:g={round(bass_gain*0.8, 1)},"
            f"equalizer=f=60:width_type=h:width=60:g={round(bass_gain*0.6, 1)},"
            f"equalizer=f=300:width_type=h:width=100:g=-3,"
            f"alimiter=limit=0.93:level=1.0"
        ),

        # ─── SPATIAL & 8D ────────────────────────────────────────
        "8d_audio": "apulsator=hz=0.125",
        "8d_bass": f"apulsator=hz=0.125,bass=g={level}:f=80:w=0.6",
        "3d_surround": "apulsator=hz=0.5,stereowiden=level=0.4",
        "wide_stereo": "stereowiden=delay=20:feedback=0.3:crossfeed=0.3:drymix=0.8",
        "hall_reverb": "aecho=0.8:0.9:1000:0.3",
        "stadium": "aecho=0.8:0.88:200:0.5,stereowiden=level=0.6,volume=2dB",

        # ─── VOCAL & EQ ──────────────────────────────────────────
        "vocal_boost": "treble=g=5,equalizer=f=1000:width_type=h:width=200:g=4",
        "crystal_clear": "treble=g=4,bass=g=-2,equalizer=f=3000:width_type=h:width=1000:g=2",
        "presence": (
            "equalizer=f=2000:width_type=h:width=500:g=3,"
            "equalizer=f=5000:width_type=h:width=1000:g=2,"
            "treble=g=3"
        ),

        # ─── CREATIVE / FUN ──────────────────────────────────────
        "nightcore": "asetrate=44100*1.25,aresample=44100,equalizer=f=60:width_type=h:width=100:g=5",
        "daycore": "asetrate=44100*0.8,aresample=44100,aecho=0.5:0.7:500:0.2",
        "slow_reverb": "asetrate=44100*0.85,aresample=44100,aecho=0.8:0.9:1000:0.3",
        "vaporwave": "asetrate=44100*0.8,aresample=44100,lowpass=f=2000",
        "lofi_chill": "aresample=22050,asetrate=22050,lowpass=f=3500,highpass=f=200,volume=1.5dB",
        "reverse": "areverse",
        "underwater": "lowpass=f=300,equalizer=f=500:width_type=h:width=200:g=-5,volume=3dB",
        "robot": "asetrate=44100*0.7,aresample=44100,aecho=0.3:0.5:30:0.5",
        "chipmunk": "asetrate=44100*1.5,aresample=44100",
        "telephone": "highpass=f=300,lowpass=f=3400,equalizer=f=1800:width_type=h:width=500:g=5",

        # ─── MASTERING ───────────────────────────────────────────
        "compressor": "acompressor=threshold=-20dB:ratio=4:attack=5:release=50",
        "limiter": "alimiter=limit=0.95:level=1.0",
        "loudness": "loudnorm=I=-16:TP=-1.5:LRA=11",
        "concert": "aecho=0.8:0.88:60:0.4,stereowiden=level=0.4",
        "clean_master": (
            "acompressor=threshold=-18dB:ratio=3:attack=10:release=80,"
            "loudnorm=I=-16:TP=-1.5:LRA=11,"
            "alimiter=limit=0.98"
        ),
    }

    return effects.get(effect_type, effects["bass_studio"])


def run_ffmpeg_simple(input_path: str, output_path: str, effect_filter: str, output_format: str) -> tuple:
    """Single file processing (no intro/outro)"""
    codec = "libmp3lame" if output_format == "mp3" else ("libvorbis" if output_format == "ogg" else "pcm_s16le")
    ext_map = {"mp3": "mp3", "ogg": "ogg", "wav": "wav"}

    cmd = [
        "ffmpeg", "-y",
        "-i", input_path,
        "-af", effect_filter,
        "-map_metadata", "0",       # Keep original metadata
        "-b:a", "320k",
        "-threads", "0",
        output_path
    ]

    if output_format == "mp3":
        cmd = [
            "ffmpeg", "-y",
            "-i", input_path,
            "-af", effect_filter,
            "-map_metadata", "0",
            "-c:a", "libmp3lame",
            "-b:a", "320k",
            "-id3v2_version", "3",
            "-threads", "0",
            output_path
        ]

    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return result.returncode, result.stderr.decode()


def run_ffmpeg_with_dj(input_path: str, output_path: str, effect_filter: str,
                        intro_path: str, outro_path: str, duration: float, output_format: str) -> tuple:
    """Processing with DJ intro/outro mixing"""
    outro_delay = max(0, int((duration * 1000) - 5000))

    filter_complex = (
        f"[0:a]volume=1.0[main];"
        f"[1:a]volume=0.8,adelay=0|0[in_v];"
        f"[2:a]volume=0.8,adelay={outro_delay}|{outro_delay}[out_v];"
        f"[main][in_v][out_v]amix=inputs=3:duration=first:dropout_transition=3[mixed];"
        f"[mixed]{effect_filter},alimiter=limit=0.99[out]"
    )

    cmd = [
        "ffmpeg", "-y",
        "-i", input_path,
        "-i", intro_path,
        "-i", outro_path,
        "-filter_complex", filter_complex,
        "-map", "[out]",
        "-map_metadata", "0",
        "-c:a", "libmp3lame",
        "-b:a", "320k",
        "-id3v2_version", "3",
        "-threads", "0",
        output_path
    ]

    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return result.returncode, result.stderr.decode()


# ================================================================
# 📡 ENDPOINTS
# ================================================================

@app.get("/")
def root():
    return {
        "service": "Cipher-MD Audio Engine",
        "version": "3.0",
        "endpoints": ["/process", "/info", "/effects", "/trim"]
    }


@app.get("/effects")
def list_effects():
    """Available audio effects list ගන්න"""
    return {
        "effects": {
            "bass": ["bass_studio", "bass_edm", "bass_car", "bass_slap", "bass_heavy", "8d_bass"],
            "spatial": ["8d_audio", "3d_surround", "wide_stereo", "hall_reverb", "stadium"],
            "vocal": ["vocal_boost", "crystal_clear", "presence"],
            "creative": ["nightcore", "daycore", "slow_reverb", "vaporwave", "lofi_chill",
                         "reverse", "underwater", "robot", "chipmunk", "telephone"],
            "mastering": ["compressor", "limiter", "loudness", "concert", "clean_master"]
        },
        "level_range": "1-20 (bass effects only)",
        "formats": ["mp3", "ogg", "wav"]
    }


@app.post("/info")
async def get_info(file: UploadFile = File(...)):
    """Audio file info/metadata ගන්න"""
    tmp_path = os.path.join(UPLOAD_FOLDER, f"info_{uuid.uuid4()}{os.path.splitext(file.filename)[1]}")
    async with aiofiles.open(tmp_path, "wb") as f:
        content = await file.read()
        await f.write(content)

    try:
        info = get_audio_info(tmp_path)
        return JSONResponse(info)
    finally:
        cleanup_file(tmp_path)


@app.post("/process")
async def process_audio(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    effect: str = Form("bass_studio"),
    level: int = Form(10),
    output_format: str = Form("mp3"),
    use_dj_mix: bool = Form(False),  # True කළොත් intro/outro mix කරනවා
):
    """
    Main audio processing endpoint.
    - effect: effect name (GET /effects for list)
    - level: 1-20 bass boost strength
    - output_format: mp3 / ogg / wav
    - use_dj_mix: true = intro.mp3 + outro.mp3 with the song
    """
    # Validate format
    if output_format not in ("mp3", "ogg", "wav"):
        raise HTTPException(status_code=400, detail="output_format must be mp3, ogg, or wav")

    uid = uuid.uuid4()
    ext = os.path.splitext(file.filename)[1] or ".mp3"
    input_path = os.path.join(UPLOAD_FOLDER, f"in_{uid}{ext}")
    output_path = os.path.join(PROCESSED_FOLDER, f"out_{uid}.{output_format}")

    # Save uploaded file
    async with aiofiles.open(input_path, "wb") as buf:
        content = await file.read()
        await buf.write(content)

    # Build effect filter
    effect_filter = build_effect_filter(effect, level)

    # Background cleanup for input file
    background_tasks.add_task(cleanup_file, input_path)
    background_tasks.add_task(cleanup_file, output_path)

    loop = asyncio.get_event_loop()

    if use_dj_mix:
        intro_path = os.path.join(ASSETS_FOLDER, "intro.mp3")
        outro_path = os.path.join(ASSETS_FOLDER, "outro.mp3")

        if not os.path.exists(intro_path) or not os.path.exists(outro_path):
            raise HTTPException(
                status_code=400,
                detail="DJ Mix requires intro.mp3 and outro.mp3 in the assets folder. "
                       "Set use_dj_mix=false to process without them."
            )

        duration = await loop.run_in_executor(executor, get_audio_duration, input_path)

        returncode, stderr = await loop.run_in_executor(
            executor,
            run_ffmpeg_with_dj,
            input_path, output_path, effect_filter,
            intro_path, outro_path, duration, output_format
        )
    else:
        returncode, stderr = await loop.run_in_executor(
            executor,
            run_ffmpeg_simple,
            input_path, output_path, effect_filter, output_format
        )

    if returncode != 0:
        print("FFmpeg Error:\n", stderr[-3000:])
        raise HTTPException(status_code=500, detail=f"FFmpeg processing failed: {stderr[-500:]}")

    if not os.path.exists(output_path):
        raise HTTPException(status_code=500, detail="Output file was not created")

    media_types = {"mp3": "audio/mpeg", "ogg": "audio/ogg", "wav": "audio/wav"}
    return FileResponse(
        output_path,
        media_type=media_types[output_format],
        filename=f"Cipher_{effect}.{output_format}"
    )


@app.post("/trim")
async def trim_audio(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    start: float = Form(0.0),    # Start time in seconds
    end: float = Form(-1.0),     # End time (-1 = until end)
    output_format: str = Form("mp3"),
    effect: str = Form(""),      # Optional effect after trim
    level: int = Form(10),
):
    """Trim audio + optional effect apply"""
    if output_format not in ("mp3", "ogg", "wav"):
        raise HTTPException(status_code=400, detail="output_format must be mp3, ogg, or wav")

    uid = uuid.uuid4()
    ext = os.path.splitext(file.filename)[1] or ".mp3"
    input_path = os.path.join(UPLOAD_FOLDER, f"trim_in_{uid}{ext}")
    output_path = os.path.join(PROCESSED_FOLDER, f"trim_out_{uid}.{output_format}")

    async with aiofiles.open(input_path, "wb") as buf:
        await buf.write(await file.read())

    background_tasks.add_task(cleanup_file, input_path)
    background_tasks.add_task(cleanup_file, output_path)

    # Build trim + optional effect filter
    trim_filter = f"atrim=start={start}" + (f":end={end}" if end > 0 else "") + ",asetpts=PTS-STARTPTS"
    if effect:
        effect_filter = build_effect_filter(effect, level)
        audio_filter = f"{trim_filter},{effect_filter}"
    else:
        audio_filter = f"{trim_filter},alimiter=limit=0.99"

    cmd = [
        "ffmpeg", "-y",
        "-i", input_path,
        "-af", audio_filter,
        "-map_metadata", "0",
        "-c:a", "libmp3lame" if output_format == "mp3" else ("libvorbis" if output_format == "ogg" else "pcm_s16le"),
        "-b:a", "320k",
        "-threads", "0",
        output_path
    ]

    loop = asyncio.get_event_loop()
    returncode, stderr = await loop.run_in_executor(
        executor,
        lambda: (lambda r: (r.returncode, r.stderr.decode()))(subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE))
    )

    if returncode != 0:
        raise HTTPException(status_code=500, detail=f"Trim failed: {stderr[-500:]}")

    media_types = {"mp3": "audio/mpeg", "ogg": "audio/ogg", "wav": "audio/wav"}
    return FileResponse(output_path, media_type=media_types[output_format], filename=f"Cipher_trimmed.{output_format}")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7862, workers=1)