import os
import struct
import subprocess
import logging
import wave

logger = logging.getLogger(__name__)

SF_TARGET_DIRS = [
    "/opt/daw_engine/soundfonts",
]


def _find_chunk(data, chunk_id, offset=12, max_size=0):
    pos = offset
    end = len(data) if max_size == 0 else offset + max_size
    while pos + 8 <= end:
        ck_id = data[pos:pos + 4]
        ck_size = struct.unpack("<I", data[pos + 4:pos + 8])[0]
        if ck_id == chunk_id:
            return (pos, ck_id, ck_size, pos + 8)
        if ck_id == b"LIST" and pos + 12 <= end:
            list_type = data[pos + 8:pos + 12]
            # Check if we're looking for a specific list type
            if chunk_id == b"smpl" and list_type == b"sdta":
                inner = _find_chunk_in_list(data, chunk_id, pos + 12, ck_size - 4)
                if inner:
                    return inner
            elif chunk_id == b"LIST":
                # When searching for LIST by type, check if it's sdta
                pass
            else:
                inner = _find_chunk_in_list(data, chunk_id, pos + 12, ck_size - 4)
                if inner:
                    return inner
        pos += 8 + ck_size
        if ck_size % 2 == 1:
            pos += 1
    return None


def _find_list_of_type(data, list_type_id, offset=12):
    pos = offset
    while pos + 12 <= len(data):
        ck_id = data[pos:pos + 4]
        ck_size = struct.unpack("<I", data[pos + 4:pos + 8])[0]
        if ck_id == b"LIST":
            form_type = data[pos + 8:pos + 12]
            if form_type == list_type_id:
                return (pos, ck_id, ck_size, pos + 12)
        pos += 8 + ck_size
        if ck_size % 2 == 1:
            pos += 1
    return None


def _find_chunk_in_list(data, chunk_id, list_data_offset, list_data_size):
    pos = list_data_offset
    end = list_data_offset + list_data_size
    while pos + 8 <= end:
        ck_id = data[pos:pos + 4]
        ck_size = struct.unpack("<I", data[pos + 4:pos + 8])[0]
        if ck_id == chunk_id:
            return (pos, ck_id, ck_size, pos + 8)
        pos += 8 + ck_size
        if ck_size % 2 == 1:
            pos += 1
    return None


def _update_size(data, offset, new_size):
    return data[:offset] + struct.pack("<I", new_size) + data[offset + 4:]


class SoundFontConverter:
    def __init__(self, target_dirs=None):
        self.target_dirs = target_dirs or SF_TARGET_DIRS

    def _check_ffmpeg_ogg(self):
        try:
            r = subprocess.run(["ffmpeg", "-encoders"], capture_output=True, text=True, timeout=5)
            return "libvorbis" in r.stdout
        except Exception:
            return False

    def _encode_sample_ogg(self, pcm: bytes, rate: int, tmp_dir: str) -> bytes:
        """Encode a single mono 16-bit PCM slice to an Ogg Vorbis stream."""
        tmp_wav = os.path.join(tmp_dir, "sample_tmp.wav")
        tmp_ogg = os.path.join(tmp_dir, "sample_tmp.ogg")
        try:
            with open(tmp_wav, "wb") as fw:
                with wave.open(fw, "wb") as w:
                    w.setnchannels(1)
                    w.setsampwidth(2)
                    w.setframerate(rate)
                    w.writeframes(pcm)
            subprocess.run([
                "ffmpeg", "-y", "-i", tmp_wav,
                "-c:a", "libvorbis", "-q:a", "3",
                "-f", "ogg", tmp_ogg
            ], capture_output=True, timeout=600, check=True)
            with open(tmp_ogg, "rb") as fo:
                return fo.read()
        except Exception as e:
            logger.error(f"Sample OGG encode failed: {e}")
            return b""
        finally:
            for p in [tmp_wav, tmp_ogg]:
                try:
                    if os.path.exists(p):
                        os.remove(p)
                except Exception:
                    pass

    def _sf2_to_sf3_python(self, sf2_path: str, sf3_path: str) -> bool:
        has_ogg = self._check_ffmpeg_ogg()
        if not has_ogg:
            logger.warning("ffmpeg with libvorbis not available, cannot convert to SF3")
            return False

        try:
            with open(sf2_path, "rb") as f:
                data = f.read()
        except Exception as e:
            logger.error(f"Cannot read {sf2_path}: {e}")
            return False

        if len(data) < 12 or data[:4] != b"RIFF" or data[8:12] != b"sfbk":
            logger.warning("Not a valid SF2 file")
            return False

        # Find smpl chunk (PCM sample data) recursively
        smpl = _find_chunk(data, b"smpl")
        if smpl is None:
            logger.warning("No smpl chunk found in SF2")
            return False

        smpl_head_off, _, smpl_old_size, smpl_data_off = smpl
        sample_data = data[smpl_data_off:smpl_data_off + smpl_old_size]

        if len(sample_data) < 16:
            logger.warning("Sample data too small")
            return False

        # Locate shdr (sample headers) inside the pdta LIST
        pdta = _find_list_of_type(data, b"pdta", 12)
        if not pdta:
            logger.warning("No pdta LIST found")
            return False
        _, _, pdta_size, pdta_data_off = pdta
        shdr = _find_chunk_in_list(data, b"shdr", pdta_data_off, pdta_size - 4)
        if not shdr:
            logger.warning("No shdr chunk found")
            return False
        _, _, shdr_size, shdr_data_off = shdr
        if shdr_size <= 0 or shdr_size % 46 != 0:
            logger.warning(f"Invalid shdr size {shdr_size}")
            return False
        n_samples = shdr_size // 46

        # Locate ifil (version) inside the INFO LIST
        ifil_abs = None
        info = _find_list_of_type(data, b"INFO", 12)
        if info:
            _, _, info_size, info_data_off = info
            ifil = _find_chunk_in_list(data, b"ifil", info_data_off, info_size - 4)
            if ifil:
                ifil_abs = ifil[3]

        import tempfile
        tmp_dir = tempfile.mkdtemp(prefix="sf3conv_")
        ogg_parts = []
        new_shdr = bytearray()
        byte_offset = 0
        ogg_bytes = 0
        try:
            for i in range(n_samples):
                base = shdr_data_off + i * 46
                # shdr layout: name[20] | start(u32) end(u32) loopstart(u32) loopend(u32) samplerate(i32) ...
                start, end, startloop, endloop, rate = struct.unpack("<IIIIi", data[base + 20:base + 40])
                sampletype = struct.unpack("<H", data[base + 44:base + 46])[0]
                ogg_stream = b""
                if end >= start and start * 2 < len(sample_data):
                    pcm = sample_data[start * 2:(end + 1) * 2]
                    if len(pcm) >= 4:
                        safe_rate = rate if 1000 < rate < 192000 else 44100
                        ogg_stream = self._encode_sample_ogg(pcm, safe_rate, tmp_dir)
                # SF3: start/end are byte offsets into the concatenated OGG stream.
                # FluidSynth treats shdr `end` as EXCLUSIVE (reads [start..end-1]),
                # so end = start + ogg length.
                new_start = byte_offset
                new_end = byte_offset + len(ogg_stream)
                # OGG loop pointers are relative to the individual decompressed sample
                new_sloop = (startloop - start) if (startloop > start and startloop <= end) else 0
                new_eloop = (endloop - start) if (endloop > start and endloop <= end) else 0
                # Mark the sample as Ogg Vorbis compressed (FLUID_SAMPLETYPE_OGG_VORBIS = 0x20)
                new_stype = sampletype | 0x20
                new_shdr += data[base:base + 20]                       # sample name
                new_shdr += struct.pack("<IIIIi", new_start, new_end, new_sloop, new_eloop, rate)
                new_shdr += data[base + 40:base + 44]                  # originalpitch, correction, samplelink
                new_shdr += struct.pack("<H", new_stype)
                ogg_parts.append(ogg_stream)
                ogg_bytes += len(ogg_stream)
                byte_offset = new_start + len(ogg_stream)
                if len(ogg_stream) % 2 == 1:
                    ogg_parts.append(b"\x00")
                    byte_offset += 1

            if not ogg_parts:
                logger.warning("No samples to encode")
                return False

            ogg_padded = b"".join(ogg_parts)
            if len(ogg_padded) % 2 == 1:
                ogg_padded += b"\x00"
            new_smpl_size = len(ogg_padded)
            old_padded = smpl_old_size + (1 if smpl_old_size % 2 == 1 else 0)
            delta = len(ogg_padded) - old_padded

            # Rebuild file: replace smpl chunk and update all sizes
            out = bytearray()
            out.extend(data[:smpl_head_off])           # up to smpl chunk header (excl. id)
            out.extend(b"smpl")                        # chunk id (required for valid SF3)
            out.extend(struct.pack("<I", new_smpl_size))  # new smpl size
            out.extend(ogg_padded)                     # concatenated per-sample OGG streams
            rest_off = smpl_head_off + 8 + old_padded
            rest = bytearray(data[rest_off:])
            # Patch the shdr sample headers inside the pdta copy
            shdr_in_rest = shdr_data_off - rest_off
            if shdr_in_rest < 0 or shdr_in_rest + shdr_size > len(rest):
                logger.warning("shdr not found after smpl chunk")
                return False
            rest[shdr_in_rest:shdr_in_rest + shdr_size] = new_shdr
            out.extend(rest)

            data_out = bytes(out)

            # SF3 requires version 3.0 so FluidSynth treats it as an SF3 file
            if ifil_abs is not None and ifil_abs + 4 <= len(data_out):
                data_out = data_out[:ifil_abs] + struct.pack("<HH", 3, 0) + data_out[ifil_abs + 4:]

            # Find sdta LIST and update its size
            sdta = _find_list_of_type(data_out, b"sdta", 12)
            if sdta:
                lh_off, _, lh_size, ld_off = sdta
                data_out = _update_size(data_out, lh_off + 4, lh_size + delta)

            # Update RIFF root size
            new_total = len(data_out) - 8
            data_out = _update_size(data_out, 4, new_total)

            with open(sf3_path, "wb") as fout:
                fout.write(data_out)

            # Validate: check that RIFF size matches actual size
            written = os.path.getsize(sf3_path)
            parsed_riff = struct.unpack("<I", data_out[4:8])[0]
            if parsed_riff != written - 8:
                logger.warning(f"Size mismatch: RIFF says {parsed_riff}, actual is {written - 8}")
                return False

            logger.info(f"Converted {n_samples} samples -> {ogg_bytes} bytes OGG ({100 * (1 - ogg_bytes / max(len(sample_data), 1)):.0f}% smaller)")
            return True

        except subprocess.TimeoutExpired:
            logger.error("Ogg conversion timed out")
            return False
        except Exception as e:
            logger.error(f"Conversion error: {e}")
            import traceback
            traceback.print_exc()
            return False
        finally:
            try:
                import shutil
                shutil.rmtree(tmp_dir, ignore_errors=True)
            except Exception:
                pass

    @staticmethod
    def _sf3_plays_audio(path: str) -> bool:
        """Verify a SoundFont actually loads and renders audible audio (guards
        against shipping malformed SF3 files that silently play nothing)."""
        if not os.path.exists(path):
            return False
        try:
            import fluidsynth
            import numpy as np
            fl = fluidsynth.Synth()
            try:
                h = fl.sfload(path)
                if h < 0:
                    return False
                fl.program_select(0, h, 0, 0)
                fl.noteon(0, 60, 100)
                frames = 8820  # 0.2s
                buf = np.zeros(frames * 2, dtype=np.float32)
                fluidsynth._fl.fluid_synth_write_float(
                    fl.synth, frames, buf.ctypes.data, 0, 1,
                    buf.ctypes.data + frames * 4, 0, 1
                )
                fl.noteoff(0, 60)
                rms = float(np.sqrt(np.mean(buf ** 2)))
                return rms > 1e-4
            finally:
                try:
                    fl.delete()
                except Exception:
                    pass
        except Exception:
            return False

    @staticmethod
    def _find_sf3_converter():
        for exe in ["fluidsynth", "mscore"]:
            try:
                subprocess.run([exe, "--help"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=2.0)
                return exe
            except (FileNotFoundError, subprocess.TimeoutExpired):
                continue
        return "python"

    def convert_sf2_to_sf3(self, sf2_path: str) -> str:
        if not os.path.exists(sf2_path):
            raise FileNotFoundError(f"Source SF2 file not found: {sf2_path}")

        sf3_path = os.path.splitext(sf2_path)[0] + ".sf3"

        # Reuse a working SF3 if it is newer than the SF2 AND actually plays
        # audio. Malformed SF3s (e.g. produced by an older converter) are
        # re-converted automatically instead of being shipped silently broken.
        if os.path.exists(sf3_path) and os.path.getmtime(sf3_path) >= os.path.getmtime(sf2_path):
            if self._sf3_plays_audio(sf3_path):
                logger.info(f"SF3 already up-to-date: {sf3_path}")
                return sf3_path
            logger.warning(f"Existing SF3 does not play audio, re-converting: {sf3_path}")
            try:
                os.remove(sf3_path)
            except Exception:
                pass

        converter = self._find_sf3_converter()
        try:
            logger.info(f"Converting '{sf2_path}' -> '{sf3_path}' using {converter}...")
            if converter == "fluidsynth":
                converter = "python"

            if converter == "mscore":
                cmd = ["mscore", "-o", sf3_path, sf2_path]
                result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
                if result.returncode == 0 and os.path.exists(sf3_path):
                    logger.info(f"Created SF3 via mscore: {sf3_path} ({os.path.getsize(sf3_path)/1024/1024:.1f}MB)")
                    return sf3_path
                converter = "python"

            if converter == "python":
                if self._sf2_to_sf3_python(sf2_path, sf3_path) and os.path.exists(sf3_path):
                    if self._sf3_plays_audio(sf3_path):
                        size_mb = os.path.getsize(sf3_path) / (1024 * 1024)
                        logger.info(f"Created SF3: {sf3_path} ({size_mb:.2f} MB)")
                        return sf3_path
                    # Conversion produced a broken file — never ship it
                    logger.warning(f"Converted SF3 failed audio verification, removing: {sf3_path}")
                    try:
                        os.remove(sf3_path)
                    except Exception:
                        pass
                logger.warning(f"Python converter failed for {sf2_path}, returning SF2 path")
                return sf2_path

            return sf2_path
        except subprocess.TimeoutExpired:
            logger.error(f"Conversion timed out for {sf2_path}")
            return sf2_path
        except Exception as e:
            logger.error(f"Error converting {sf2_path}: {e}")
            return sf2_path

    def sf3_to_sf2(self, sf3_path: str, sf2_path: str = None) -> str:
        """Convert an SF3 (Ogg Vorbis samples) soundfont into a playable SF2.

        The client FluidSynth WASM cannot decode Ogg Vorbis/SF3 samples, so any
        SF3 soundfont (converted or uploaded) plays silence. Decompressing to
        SF2 makes every instrument audible again.
        """
        if not os.path.exists(sf3_path):
            raise FileNotFoundError(f"Source SF3 file not found: {sf3_path}")

        try:
            with open(sf3_path, "rb") as f:
                data = f.read()
        except Exception as e:
            raise IOError(f"Cannot read {sf3_path}: {e}")

        if len(data) < 12 or data[:4] != b"RIFF" or data[8:12] != b"sfbk":
            raise ValueError(f"Not a valid SoundFont file: {sf3_path}")

        smpl = _find_chunk(data, b"smpl")
        if not smpl:
            raise ValueError("No smpl chunk found")
        smpl_head_off, _, smpl_old_size, smpl_data_off = smpl
        sample_data = data[smpl_data_off:smpl_data_off + smpl_old_size]

        pdta = _find_list_of_type(data, b"pdta", 12)
        if not pdta:
            raise ValueError("No pdta LIST found")
        _, _, pdta_size, pdta_data_off = pdta
        shdr = _find_chunk_in_list(data, b"shdr", pdta_data_off, pdta_size - 4)
        if not shdr:
            raise ValueError("No shdr chunk found")
        _, _, shdr_size, shdr_data_off = shdr
        if shdr_size <= 0 or shdr_size % 46 != 0:
            raise ValueError(f"Invalid shdr size {shdr_size}")
        n_samples = shdr_size // 46

        ifil_abs = None
        info = _find_list_of_type(data, b"INFO", 12)
        if info:
            _, _, info_size, info_data_off = info
            ifil = _find_chunk_in_list(data, b"ifil", info_data_off, info_size - 4)
            if ifil:
                ifil_abs = ifil[3]

        import tempfile
        tmp_dir = tempfile.mkdtemp(prefix="sf2conv_")
        pcm_parts = []
        new_shdr = bytearray()
        frame_offset = 0
        total_pcm_bytes = 0
        try:
            for i in range(n_samples):
                base = shdr_data_off + i * 46
                name = data[base:base + 20]
                start, end, loopstart, loopend, rate = struct.unpack("<IIIIi", data[base + 20:base + 40])
                sampletype = struct.unpack("<H", data[base + 44:base + 46])[0]
                frames = 0
                pcm = b""
                if end > start and start < len(sample_data):
                    # FluidSynth reads the OGG region as [start..end-1]
                    ogg = sample_data[start:min(end, len(sample_data))]
                    if len(ogg) >= 4 and ogg[:4] == b"OggS":
                        pcm = self._decode_sample_ogg(ogg, tmp_dir)
                        frames = len(pcm) // 2
                new_start = frame_offset
                new_end = frame_offset + frames
                # SF3 loop points are relative to the decompressed sample; SF2 needs absolute
                new_loopstart = loopstart + new_start if (loopstart or loopend) else 0
                new_loopend = loopend + new_start if (loopstart or loopend) else 0
                # Clear the Ogg Vorbis flag; keep mono/left/right/linked flags
                new_stype = sampletype & ~0x20
                new_shdr += name
                new_shdr += struct.pack("<IIIIi", new_start, new_end, new_loopstart, new_loopend, rate)
                new_shdr += data[base + 40:base + 44]
                new_shdr += struct.pack("<H", new_stype)
                pcm_parts.append(pcm)
                frame_offset += frames
                if frames:
                    total_pcm_bytes += frames * 2

            if total_pcm_bytes == 0:
                logger.warning("SF3 contained no decodable samples")
                return sf3_path

            new_smpl_size = total_pcm_bytes
            if new_smpl_size % 2 == 1:
                new_smpl_size += 1
            old_padded = smpl_old_size + (1 if smpl_old_size % 2 == 1 else 0)
            delta = new_smpl_size - old_padded

            out = bytearray()
            out.extend(data[:smpl_head_off])
            out.extend(b"smpl")
            out.extend(struct.pack("<I", new_smpl_size))
            for part in pcm_parts:
                out.extend(part)
            if total_pcm_bytes % 2 == 1:
                out.extend(b"\x00")
            rest_off = smpl_head_off + 8 + old_padded
            rest = bytearray(data[rest_off:])
            shdr_in_rest = shdr_data_off - rest_off
            if shdr_in_rest < 0 or shdr_in_rest + shdr_size > len(rest):
                logger.warning("shdr not found after smpl chunk")
                return sf3_path
            rest[shdr_in_rest:shdr_in_rest + shdr_size] = new_shdr
            out.extend(rest)

            data_out = bytes(out)

            # Back to SF2 version 2.01
            if ifil_abs is not None and ifil_abs + 4 <= len(data_out):
                data_out = data_out[:ifil_abs] + struct.pack("<HH", 2, 1) + data_out[ifil_abs + 4:]

            sdta = _find_list_of_type(data_out, b"sdta", 12)
            if sdta:
                lh_off, _, lh_size, ld_off = sdta
                data_out = _update_size(data_out, lh_off + 4, lh_size + delta)

            data_out = _update_size(data_out, 4, len(data_out) - 8)

            sf2_path = sf2_path or (os.path.splitext(sf3_path)[0] + ".sf2")
            with open(sf2_path, "wb") as fout:
                fout.write(data_out)
            logger.info(f"Converted SF3 -> SF2: {sf2_path} ({os.path.getsize(sf2_path) / 1048576:.1f} MB)")
            return sf2_path
        finally:
            try:
                import shutil
                shutil.rmtree(tmp_dir, ignore_errors=True)
            except Exception:
                pass

    def _decode_sample_ogg(self, ogg: bytes, tmp_dir: str) -> bytes:
        """Decode an Ogg Vorbis stream to mono 16-bit PCM; returns PCM bytes."""
        tmp_ogg = os.path.join(tmp_dir, "sample.ogg")
        tmp_pcm = os.path.join(tmp_dir, "sample.pcm")
        try:
            with open(tmp_ogg, "wb") as fo:
                fo.write(ogg)
            r = subprocess.run(
                ["ffmpeg", "-y", "-v", "error", "-i", tmp_ogg, "-f", "s16le", "-ac", "1", tmp_pcm],
                capture_output=True, timeout=600)
            if r.returncode != 0:
                logger.warning(f"OGG decode failed: {r.stderr.decode(errors='replace')[:120]}")
                return b""
            with open(tmp_pcm, "rb") as fp:
                return fp.read()
        except Exception as e:
            logger.warning(f"OGG decode error: {e}")
            return b""
        finally:
            for p in [tmp_ogg, tmp_pcm]:
                try:
                    if os.path.exists(p):
                        os.remove(p)
                except Exception:
                    pass

    def batch_convert_all(self):
        for sdir in self.target_dirs:
            if not os.path.isdir(sdir):
                continue
            for fname in sorted(os.listdir(sdir)):
                if fname.lower().endswith(".sf2"):
                    self.convert_sf2_to_sf3(os.path.join(sdir, fname))