#!/usr/bin/env python3
"""Kitmaker web app — Razzmatazz kit builder + audio converter."""

import asyncio
import contextlib
import hashlib
import io
import json
import logging
import os
import re
import secrets
import shutil
import sqlite3
import subprocess
import tempfile
import threading
import wave
import zipfile
from dataclasses import dataclass, field
from pathlib import Path
from xml.sax.saxutils import escape

log = logging.getLogger(__name__)

import analysis

from fastapi import FastAPI, File, Form, Request, UploadFile
from fastapi.responses import Response
from fastapi.staticfiles import StaticFiles
import uvicorn

# ── Startup checks ───────────────────────────────────────────────────
_ffmpeg_ok = bool(shutil.which('ffmpeg'))
if not _ffmpeg_ok:
    import warnings
    warnings.warn('ffmpeg not found on PATH. Audio conversion and probe will fail.')

AUDIO_EXTENSIONS = {'.wav', '.mp3', '.flac', '.ogg', '.aiff', '.aif', '.m4a', '.aac'}
PAD_COUNT = 8
GRID_WIDTH = 4
ANALYSIS_WORKERS = 2

# ── Analysis queue ───────────────────────────────────────────────────

@dataclass
class AnalysisJob:
    raw: bytes
    ext: str
    fname: str
    force: bool
    cache_key: str
    result: asyncio.Future = field(default_factory=lambda: asyncio.get_event_loop().create_future())

_analysis_queue: asyncio.Queue | None = None
_worker_tasks: list[asyncio.Task] = []


async def _analysis_worker(worker_id: int):
    while True:
        job = await _analysis_queue.get()
        try:
            if not job.force:
                cached = cache_get(job.cache_key)
                if cached:
                    cached['_cached'] = True
                    job.result.set_result(cached)
                    continue
            if job.force:
                db = get_db()
                db.execute('DELETE FROM analysis_cache WHERE file_hash=?', (job.cache_key,))
                db.commit()
            result = await asyncio.to_thread(analysis.analyze_audio, job.raw, job.ext, job.fname)
            result['_cached'] = False
            cache_set(job.cache_key, result)
            job.result.set_result(result)
        except Exception as e:
            job.result.set_exception(e)
        finally:
            _analysis_queue.task_done()


@contextlib.asynccontextmanager
async def lifespan(app):
    global _analysis_queue
    _analysis_queue = asyncio.Queue()
    loop = asyncio.get_running_loop()
    _worker_tasks.clear()
    for i in range(ANALYSIS_WORKERS):
        task = loop.create_task(_analysis_worker(i))
        _worker_tasks.append(task)
    yield
    for task in _worker_tasks:
        task.cancel()
    await asyncio.gather(*_worker_tasks, return_exceptions=True)
    _worker_tasks.clear()
    _analysis_queue = None


app = FastAPI(title="Razzmatazz Kitmaker", lifespan=lifespan)
app.state.ffmpeg_ok = _ffmpeg_ok

# ── Session store ────────────────────────────────────────────────────
_sessions: dict[str, dict] = {}


def _session_id(request) -> str:
    sid = request.cookies.get('kitmaker_session')
    if not sid or sid not in _sessions:
        sid = secrets.token_hex(16)
        _sessions[sid] = {}
    return sid


# ── Analysis cache (SQLite) ──────────────────────────────────────────
# Content-addressed: SHA256(file_bytes) -> analysis result JSON.
# Same file content = same hash = skip re-analysis.
DB_PATH = Path(__file__).with_name('cache.db')

_db: sqlite3.Connection | None = None
_db_lock = threading.Lock()

def get_db() -> sqlite3.Connection:
    global _db
    if _db is not None:
        return _db
    _db = sqlite3.connect(str(DB_PATH), check_same_thread=False)
    _db.row_factory = sqlite3.Row
    _db.execute('PRAGMA journal_mode=WAL')
    _db.execute('PRAGMA synchronous=FULL')
    _db.execute('PRAGMA busy_timeout=10000')
    _db.execute('PRAGMA wal_autocheckpoint=100')
    _db.execute('''CREATE TABLE IF NOT EXISTS analysis_cache
        (file_hash TEXT PRIMARY KEY, result TEXT, analyzed_at TEXT DEFAULT CURRENT_TIMESTAMP)''')
    _db.commit()
    return _db

def cache_get(file_hash: str) -> dict | None:
    try:
        with _db_lock:
            row = get_db().execute(
                'SELECT result FROM analysis_cache WHERE file_hash=?', (file_hash,)
            ).fetchone()
            return json.loads(row[0]) if row else None
    except Exception as e:
        log.warning('cache_get(%s...) failed: %s', file_hash[:16], e)
        return None

def cache_set(file_hash: str, result: dict):
    try:
        with _db_lock:
            db = get_db()
            db.execute(
                'INSERT OR REPLACE INTO analysis_cache (file_hash, result) VALUES (?,?)',
                (file_hash, json.dumps(result)),
            )
            db.commit()
            db.execute('PRAGMA wal_checkpoint(PASSIVE)')
    except Exception as e:
        log.warning('cache_set(%s...) failed: %s', file_hash[:16], e)

SAMPLE_FMT_TO_BITS = {
    's16': 16, 's16le': 16, 's16be': 16,
    's32': 32, 's32le': 32, 's32be': 32,
    's24': 24, 's24le': 24, 's24be': 24,
    'flt': 32, 'fltle': 32, 'fltbe': 32,
}

ALLOWED_CHANNELS = {1, 2}
ALLOWED_BIT_DEPTHS = {16, 24, 32}


def sanitize_name(name):
    name = name.strip()
    name = re.sub(r'[<>:"/\\|?*\x00-\x1f]', ' ', name)
    name = re.sub(r'\s+', ' ', name)
    name = re.sub(r'\.+$', '', name)
    return name.strip() or 'Unnamed Kit'


def strip_ext(name):
    i = name.rfind('.')
    return name[:i] if i > 0 else name


def convert_audio(input_bytes: bytes, ext: str) -> bytes:
    """Convert any audio to Razzmatazz-compatible WAV (PCM tag 1, preserve
    sample rate, 24-bit max, 1-2 channels)."""
    with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp_in:
        tmp_in.write(input_bytes)
        inpath = tmp_in.name

    outpath = inpath + '.wav'
    try:
        ch, sr, bd, _dur = _probe(inpath)
        target_ch = ch if ch in ALLOWED_CHANNELS else 2
        target_bd = bd if bd in ALLOWED_BIT_DEPTHS else 24
        if target_bd > 24:
            target_bd = 24
        acodec = {16: 'pcm_s16le', 24: 'pcm_s24le', 32: 'pcm_s24le'}[target_bd]
        raw_fmt = acodec.replace('pcm_', '')

        cmd = ['ffmpeg', '-y', '-i', inpath, '-f', raw_fmt, '-acodec', acodec]
        if target_ch != ch:
            cmd += ['-ac', str(target_ch)]
        cmd += ['pipe:1']

        r = subprocess.run(cmd, capture_output=True, timeout=120)
        if r.returncode != 0:
            err = r.stderr.decode(errors='replace')[:200]
            raise RuntimeError(f'ffmpeg: {err}')

        buf = io.BytesIO()
        with wave.open(buf, 'wb') as wf:
            wf.setnchannels(target_ch)
            wf.setsampwidth(target_bd // 8)
            wf.setframerate(sr)
            wf.writeframes(r.stdout)
        return buf.getvalue()
    finally:
        for p in (inpath, outpath):
            try:
                os.unlink(p)
            except OSError:
                pass


def _probe(path: str):
    r = subprocess.run(
        ['ffprobe', '-v', 'quiet', '-print_format', 'json', '-show_streams',
         '-show_format', path],
        capture_output=True, text=True, timeout=15)
    data = json.loads(r.stdout)
    dur = None
    if 'format' in data:
        d = data['format'].get('duration')
        if d:
            dur = round(float(d), 3)
    for s in data.get('streams', []):
        if s.get('codec_type') == 'audio':
            ch = s.get('channels', 2)
            sr = int(s.get('sample_rate', 44100))
            fmt = s.get('sample_fmt', 's16')
            bd = SAMPLE_FMT_TO_BITS.get(fmt, 16)
            return ch, sr, bd, dur
    return 2, 44100, 16, dur


# ── NNR XML generation ──────────────────────────────────────────────
# Ported from shared/templates.js and shared/export.js

COMMON_ROWS = [
    {'row': 0, 'type': 'samtempl', 'params': {}},
    {'row': 10, 'type': 'seqmod', 'params': {'steplen': 4, 'seqsteps': 16,
     'seqkeytrig': 0, 'seqquant': 0}, 'sequence': True},
    {'row': 12, 'type': 'drumsettings', 'params': {'chokegrp': 0, 'loopmode': 0,
     'maclaunchtogg': 0, 'stackfiltmode': 0, 'macpitchmod': 0, 'playthru': 0, 'fx3enable': 1}},
    {'row': 13, 'type': 'samtempl', 'params': {}},
    {'row': 14, 'type': 'samtempl', 'params': {}},
    {'row': 16, 'type': 'samtempl', 'params': {}},
]

SYNTH_TEMPLATES = json.loads(Path(__file__).with_name('templates.json').read_text())

GLOBAL_FX = json.loads(Path(__file__).with_name('global_fx.json').read_text())

DEFAULT_SEQ_EVENTS = json.loads(Path(__file__).with_name('default_seq.json').read_text())


def render_params(params):
    attrs = ' '.join(f'{k}="{escape(str(v))}"' for k, v in params.items())
    return f'<params {attrs}/>' if attrs else '<params/>'


def render_cell(cell, filename, synth_id, volume, pan, muted):
    p = dict(cell['params'])
    if cell['row'] == 3:
        p['pitchsemi'] = 0
        p['pitchfine'] = 0
    if cell['row'] == 6:
        p['envdecay'] = 1000
    if cell['row'] == 9:
        p['vfxdist'] = 0
        p['fx1send'] = 0
        p['fx2send'] = 0
    if cell['row'] == 11:
        p['selpad'] = synth_id
        p['Mute'] = 1 if muted else 0
        p['gaindb'] = -12000 if muted else round(volume * 10)
        p['panpos'] = 0 if muted else round(pan * 10)
    if cell['row'] == 15:
        for k in ('macpitch',):
            p.pop(k, None)
        p['macdelay'] = 1000
        p['oscwavmix'] = 1000
        p['modeldistamt'] = 0

    fname = filename if cell['row'] == 3 else (cell.get('filename') or '')
    lines = [
        f'        <cell row="{cell["row"]}" column="0" synth="{synth_id}" filename="{escape(fname)}" type="{cell["type"]}">',
        f'            {render_params(p)}',
    ]
    if cell.get('sequence'):
        lines.append('            <sequence/>')
    lines.append('        </cell>')
    return '\n'.join(lines)


def render_seq(seqnum):
    seq = next((s for s in DEFAULT_SEQ_EVENTS if s['seqnum'] == seqnum), None)
    params = seq['params'] if seq else {
        'notesteplen': '10', 'notestepcount': '16',
        'midioutchan': '0', 'seqplayenable': '1', 'seqstepmode': '1'}
    events = seq.get('events', []) if seq else []
    lines = [
        f'        <cell seqnum="{seqnum}" type="noteseq">',
        f'            {render_params(params)}',
        '            <sequence>' if events else '            <sequence/>',
    ]
    for ev in events:
        attrs = ' '.join(f'{k}="{escape(str(v))}"' for k, v in ev.items())
        lines.append(f'                <seqevent {attrs}/>')
    if events:
        lines.append('            </sequence>')
    lines.append('        </cell>')
    return '\n'.join(lines)


def render_global_fx():
    lines = []
    for fx in GLOBAL_FX:
        lines.append(f'        <synth id="{fx["id"]}" name=""/>')
        lines.append(f'        <cell row="0" column="0" synth="{fx["id"]}" filename="" type="{fx["type"]}">')
        lines.append(f'            {render_params(fx["params"])}')
        lines.append('        </cell>')
    return lines


def build_nnr(kit_name, slots):
    safe = sanitize_name(kit_name)
    lines = ['<?xml version="1.0" encoding="UTF-8"?>', '<document>', '    <session version="4">']
    for slot in slots:
        si = 7 - slot['slotIndex']  # reverse to match Razzmatazz pad layout
        has_sample = bool(slot.get('filename'))
        stype = slot.get('type', 'tom')
        template = SYNTH_TEMPLATES.get(stype, SYNTH_TEMPLATES['tom'])
        disp_name = sanitize_name(slot.get('label', f'Pad {slot["slotIndex"] + 1}'))[:24]
        sample_path = f'\\Samples\\{safe}\\{slot.get("filename", "")}' if has_sample else ''
        muted = not has_sample
        lines.append(f'        <synth id="{si}" name="{escape(disp_name)}"/>')
        for cell in template['cells']:
            lines.append(render_cell(cell, sample_path, si, slot.get('volume', 0), slot.get('pan', 0), muted))
    lines.extend(render_global_fx())
    for seqnum in range(16):
        lines.append(render_seq(seqnum))
    lines.append('        <midimap/>')
    lines.append('    </session>')
    lines.append('</document>')
    return '\n'.join(lines) + '\n'


# ── API endpoints ────────────────────────────────────────────────────


@app.get('/api/ping')
def ping():
    return {'ok': True}


@app.post('/api/session/save')
async def api_session_save(request: Request, data: dict = None):
    sid = _session_id(request)
    if data is not None:
        _sessions[sid] = data
    resp = Response(json.dumps({'ok': True}))
    resp.set_cookie('kitmaker_session', sid, max_age=86400 * 30, httponly=True, samesite='lax')
    return resp


@app.get('/api/session/load')
async def api_session_load(request: Request):
    sid = _session_id(request)
    data = _sessions.get(sid, {})
    resp = Response(json.dumps(data))
    resp.set_cookie('kitmaker_session', sid, max_age=86400 * 30, httponly=True, samesite='lax')
    return resp


@app.post('/api/probe')
async def api_probe(file: UploadFile = File(...)):
    ext = os.path.splitext(file.filename or '.wav')[1].lower()
    if ext not in AUDIO_EXTENSIONS:
        return Response(json.dumps({'error': f'Unsupported format: {ext}'}),
                        status_code=400, media_type='application/json')
    raw = await file.read()
    import tempfile
    with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as f:
        f.write(raw)
        path = f.name
    try:
        ch, sr, bd, dur = _probe(path)
        return {
            'channels': ch,
            'sampleRate': sr,
            'bitDepth': bd,
            'duration': dur or 0,
            'ext': ext,
            'size': len(raw),
        }
    except Exception as e:
        return Response(json.dumps({'error': str(e)}),
                        status_code=400, media_type='application/json')
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass


@app.post('/api/lookup')
async def api_lookup(data: dict = None):
    """Bulk lookup: returns {hash -> result_or_null} for known hashes."""
    if data is None:
        return {}
    hashes = data.get('hashes', [])
    if not hashes:
        return {}
    out = {}
    with _db_lock:
        db = get_db()
        for h in hashes:
            row = db.execute(
                'SELECT result FROM analysis_cache WHERE file_hash=?', (h,)
            ).fetchone()
            if row is not None:
                out[h] = json.loads(row[0])
    return out

@app.post('/api/analyze')
async def api_analyze(file: UploadFile = File(...), force: bool = False):
    ext = os.path.splitext(file.filename or '.wav')[1].lower()
    if ext not in AUDIO_EXTENSIONS:
        return Response(json.dumps({'error': f'Unsupported: {ext}'}),
                        status_code=400, media_type='application/json')
    raw = await file.read()
    file_hash = hashlib.sha256(raw).hexdigest()
    cache_key = file_hash
    fname = file.filename or ''
    if fname:
        cache_key = hashlib.sha256((file_hash + fname).encode()).hexdigest()
    job = AnalysisJob(raw=raw, ext=ext, fname=fname, force=force, cache_key=cache_key)
    _analysis_queue.put_nowait(job)
    try:
        result = await asyncio.wait_for(job.result, timeout=120)
        return result
    except asyncio.TimeoutError:
        return Response(json.dumps({'error': 'Analysis timed out'}),
                        status_code=504, media_type='application/json')
    except Exception as e:
        return Response(json.dumps({'error': str(e)}),
                        status_code=400, media_type='application/json')


@app.post('/api/convert')
async def api_convert(file: UploadFile = File(...)):
    ext = os.path.splitext(file.filename or '.wav')[1].lower()
    raw = await file.read()
    try:
        wav = convert_audio(raw, ext)
        return Response(content=wav, media_type='audio/wav',
                        headers={'X-Original-Name': file.filename or 'sample.wav'})
    except Exception as e:
        return Response(json.dumps({'error': str(e)}), status_code=400, media_type='application/json')


@app.post('/api/export')
async def api_export(
    kitName: str = Form(...),
    config: str = Form(...),  # JSON: [{slotIndex, type, label, volume, pan}, ...]
    files: list[UploadFile] = File(default=[]),
):
    name = sanitize_name(kitName) or 'MyKit'
    slots = json.loads(config)

    filled = [s for s in slots if s.get('filename')]
    if not filled:
        return Response(json.dumps({'error': 'No samples assigned to any pad.'}),
                        status_code=400, media_type='application/json')

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
        sample_dir = f'Samples/{name}'
        used_names = set()
        for slot in filled:
            fn = slot['filename']
            uploaded = next((f for f in files if f.filename == fn), None)
            if uploaded is None:
                slot['filename'] = ''
                continue

            raw = await uploaded.read()
            if len(raw) == 0:
                slot['filename'] = ''
                continue

            ext = os.path.splitext(fn)[1].lower()
            try:
                wav = convert_audio(raw, ext)
            except Exception:
                continue

            base = sanitize_name(strip_ext(fn)) or f'pad{slot["slotIndex"] + 1}'
            out_name = base + '.wav'
            if len(out_name) > 200:
                base = base[:180]
                out_name = base + '.wav'
            suffix = 2
            while out_name.lower() in used_names:
                stem = f'{base}-{suffix}'
                if len(stem) > 200:
                    stem = stem[:200]
                out_name = f'{stem}.wav'
                suffix += 1
            used_names.add(out_name.lower())

            slot['filename'] = out_name
            zf.writestr(f'{sample_dir}/{out_name}', wav)

        nnr_xml = build_nnr(name, slots)
        zf.writestr(f'{name}.nnr', nnr_xml)

    return Response(content=buf.getvalue(), media_type='application/zip',
                    headers={'Content-Disposition': f'attachment; filename="{name}.zip"'})


# ── Serve frontend ───────────────────────────────────────────────────

app.mount('/', StaticFiles(directory=Path(__file__).parent / 'static', html=True), name='static')


def main():
    uvicorn.run('server:app', host='0.0.0.0', port=8800, reload=True)


if __name__ == '__main__':
    main()