import base64
import io
import os
import struct
import wave

# SNAC token constants for Maya1
CODE_START_TOKEN_ID = 128257
CODE_END_TOKEN_ID = 128258
CODE_TOKEN_OFFSET = 128266
SNAC_MIN_ID = 128266
SNAC_MAX_ID = 156937
SNAC_TOKENS_PER_FRAME = 7

SOH_ID = 128259
EOH_ID = 128260
SOA_ID = 128261
BOS_ID = 128000
TEXT_EOT_ID = 128009


class EndpointHandler:
    def __init__(self, path=""):
        # allow overriding device and dtype via environment variables for local CPU testing
        # `MAYA_DEVICE`: 'cpu' or 'auto' (default 'auto')
        # `MAYA_TORCH_DTYPE`: 'bf16' or 'fp32' (default 'bf16')
        # `MAYA_USE_FAKE`: if '1', use a tiny fake pipeline for smoke testing (no HF downloads)
        device_override = os.getenv("MAYA_DEVICE", "auto")
        dtype_override = os.getenv("MAYA_TORCH_DTYPE", "bf16")
        use_fake = os.getenv("MAYA_USE_FAKE", "0") == "1"

        # map string to torch dtype will be set after importing torch

        if use_fake:
            # Minimal fake components for quick local smoke tests. Generates a short sine tone.
            self.model = None
            self.tokenizer = None
            # keep device as plain string for fake mode
            self.device = "cpu"
            self.snac = None
            # no external libs required for fake mode
            self.sf = None
            return

        # import heavy inference dependencies lazily so fake-mode doesn't require them
        try:
            import soundfile as sf
            import torch
            from snac import SNAC
            from transformers import AutoModelForCausalLM, AutoTokenizer
        except Exception as e:
            raise RuntimeError(
                f"Failed to import inference dependencies: {e}.\n"
                "Install the packages listed in `inference_endpoints/requirements.txt`."
            )

        # map string to torch dtype
        torch_dtype = torch.bfloat16 if dtype_override == "bf16" else torch.float32

        # load Maya1 text-to-voice model
        # force CPU device_map when requested to avoid trying to use GPUs
        device_map_arg = "auto"
        if device_override == "cpu":
            device_map_arg = "cpu"

        # Allow loading from HuggingFace Hub via environment variable
        # MAYA_MODEL_ID: HuggingFace model ID (e.g., "maya-research/maya1" or similar Maya model)
        model_id = os.getenv("MAYA_MODEL_ID", "")

        # Determine model path: use env var if set, otherwise local path
        local_path = path if path else "/repository"

        # Check if trying to load from self (circular reference)
        # Only block if it's exactly "binhqd/maya1" - other repos ending in /maya1 are fine
        if model_id == "binhqd/maya1":
            print(f"⚠️ MAYA_MODEL_ID is set to {model_id} which is this repository!")
            print("This creates a circular reference. Checking local files instead...")
            model_id = ""  # Force local loading

        if model_id:
            # Load from HuggingFace Hub
            model_path = model_id
            print(f"Loading Maya1 model from HuggingFace Hub: {model_id}")
        else:
            # Load from local repository
            model_path = local_path

            # Check if this is a valid model directory
            config_path = os.path.join(model_path, "config.json")
            if not os.path.exists(config_path):
                raise RuntimeError(
                    f"❌ Model configuration not found at: {config_path}\n\n"
                    f"The repository appears to be missing model weights and configuration.\n\n"
                    f"To fix this, you have 3 options:\n\n"
                    f"1. 🧪 TESTING MODE (Quick Start):\n"
                    f"   Set: MAYA_USE_FAKE=1\n"
                    f"   This generates test audio (sine wave) without requiring model files.\n\n"
                    f"2. 🌐 USE EXISTING MODEL FROM HUGGINGFACE:\n"
                    f"   Set: MAYA_MODEL_ID=<actual-model-id>\n"
                    f"   Example: MAYA_MODEL_ID=maya-research/maya1\n"
                    f"   (Must be a valid HuggingFace model repository with weights)\n\n"
                    f"3. 📦 UPLOAD YOUR MODEL FILES:\n"
                    f"   Upload these files to your repository (binhqd/maya1):\n"
                    f"   - config.json\n"
                    f"   - model.safetensors (or pytorch_model.bin)\n"
                    f"   - tokenizer.json\n"
                    f"   - tokenizer_config.json\n"
                    f"   - generation_config.json (optional)\n\n"
                    f"Current path: {model_path}\n"
                    f"Files found: {os.listdir(model_path) if os.path.exists(model_path) else 'path does not exist'}\n\n"
                    f"⚠️ Note: MAYA_MODEL_ID=binhqd/maya1 won't work (this repository)\n"
                    f"   because it doesn't contain model weights.\n"
                )

        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                model_path, torch_dtype=torch_dtype, device_map=device_map_arg
            )
            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        except Exception as e:
            raise RuntimeError(
                f"Failed to load Maya1 model from {model_path}.\n"
                f"Error: {e}\n\n"
                f"Please ensure:\n"
                f"1. The model repository contains valid model files, OR\n"
                f"2. Set MAYA_MODEL_ID to a valid HuggingFace model ID, OR\n"
                f"3. Set MAYA_USE_FAKE=1 for testing without a real model"
            )

        # determine device from model parameters (safer than using `model.device`)
        try:
            self.device = next(self.model.parameters()).device
        except StopIteration:
            # fallback to CPU if model has no parameters
            self.device = torch.device("cpu")

        # load SNAC model (audio decoder) 24 kHz
        self.snac = (
            SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to(self.device)
        )
        self.sf = sf

    def build_prompt(self, description: str, text: str) -> str:
        """Build formatted prompt for Maya1."""
        soh_token = self.tokenizer.decode([SOH_ID])
        eoh_token = self.tokenizer.decode([EOH_ID])
        soa_token = self.tokenizer.decode([SOA_ID])
        sos_token = self.tokenizer.decode([CODE_START_TOKEN_ID])
        eot_token = self.tokenizer.decode([TEXT_EOT_ID])
        bos_token = self.tokenizer.bos_token
        
        formatted_text = f'<description="{description}"> {text}'
        
        prompt = (
            soh_token + bos_token + formatted_text + eot_token +
            eoh_token + soa_token + sos_token
        )
        
        return prompt

    def extract_snac_codes(self, token_ids: list) -> list:
        """Extract SNAC codes from generated tokens."""
        try:
            eos_idx = token_ids.index(CODE_END_TOKEN_ID)
        except ValueError:
            eos_idx = len(token_ids)
        
        snac_codes = [
            token_id for token_id in token_ids[:eos_idx]
            if SNAC_MIN_ID <= token_id <= SNAC_MAX_ID
        ]
        
        return snac_codes

    def unpack_snac_from_7(self, snac_tokens: list) -> list:
        """Unpack 7-token SNAC frames to 3 hierarchical levels."""
        if snac_tokens and snac_tokens[-1] == CODE_END_TOKEN_ID:
            snac_tokens = snac_tokens[:-1]
        
        frames = len(snac_tokens) // SNAC_TOKENS_PER_FRAME
        snac_tokens = snac_tokens[:frames * SNAC_TOKENS_PER_FRAME]
        
        if frames == 0:
            return [[], [], []]
        
        l1, l2, l3 = [], [], []
        
        for i in range(frames):
            slots = snac_tokens[i*7:(i+1)*7]
            l1.append((slots[0] - CODE_TOKEN_OFFSET) % 4096)
            l2.extend([
                (slots[1] - CODE_TOKEN_OFFSET) % 4096,
                (slots[4] - CODE_TOKEN_OFFSET) % 4096,
            ])
            l3.extend([
                (slots[2] - CODE_TOKEN_OFFSET) % 4096,
                (slots[3] - CODE_TOKEN_OFFSET) % 4096,
                (slots[5] - CODE_TOKEN_OFFSET) % 4096,
                (slots[6] - CODE_TOKEN_OFFSET) % 4096,
            ])
        
        return [l1, l2, l3]

    def __call__(self, data):
        """
        HF Endpoints format:
        {
          "inputs": "text to speak" or {"text": "...", "description": "..."},
          "description": "... optional voice description ...",
          "generation_args": { optional dict for text generation params }
        }

        Returns dict with base64 audio:
        {
          "audio_base64": "<base64-encoded WAV data>",
          "sampling_rate": 24000
        }
        """
        # Extract inputs (HF always provides this key)
        inputs = data.get("inputs", "")

        # Get additional parameters from top level
        description = data.get("description", "")
        generation_args = data.get("generation_args", {})

        # Parse inputs
        if isinstance(inputs, dict):
            # inputs is a dict with text and description
            text = inputs.get("text", "")
            description = inputs.get(
                "description", description
            )  # override if in inputs
        elif isinstance(inputs, str):
            # inputs is just the text string
            text = inputs
        else:
            # Try to convert to string
            text = str(inputs) if inputs else ""

        if not text:
            return {"error": f"No text provided. Received: {data}"}
        
        # Use default description if not provided
        if not description:
            description = "Realistic male voice in the 30s age with american accent. Normal pitch, warm timbre, conversational pacing."

        # If running in fake mode (quick smoke test), synthesize a sine tone
        if getattr(self, "snac", None) is None and getattr(self, "model", None) is None:
            # generate a 1-second 24kHz sine wave and write a 16-bit WAV using stdlib
            sr = 24000
            duration = 1.0
            n_samples = int(sr * duration)
            freq = 220.0
            # generate samples without numpy
            waveform = [
                int(
                    0.1
                    * 32767
                    * __import__("math").sin(
                        2 * __import__("math").pi * freq * (i / sr)
                    )
                )
                for i in range(n_samples)
            ]

            buf = io.BytesIO()
            with wave.open(buf, "wb") as wf:
                wf.setnchannels(1)
                wf.setsampwidth(2)  # 16-bit
                wf.setframerate(sr)
                # pack samples as little-endian signed 16-bit
                frames = struct.pack("<" + ("h" * len(waveform)), *waveform)
                wf.writeframes(frames)
            wav_bytes = buf.getvalue()
            b64 = base64.b64encode(wav_bytes).decode("utf-8")

            return [{"blob": b64, "content-type": "audio/wav"}]

        # Build properly formatted prompt for Maya1
        prompt = self.build_prompt(description, text)

        # Tokenize
        import torch
        tokenizer_inputs = self.tokenizer(prompt, return_tensors="pt")
        if torch.cuda.is_available():
            tokenizer_inputs = {k: v.to(self.device) for k, v in tokenizer_inputs.items()}
        else:
            tokenizer_inputs = tokenizer_inputs.to(self.device)

        # Set default generation args if not provided
        default_gen_args = {
            "max_new_tokens": 2048,
            "min_new_tokens": 28,
            "temperature": 0.4,
            "top_p": 0.9,
            "repetition_penalty": 1.1,
            "do_sample": True,
            "eos_token_id": CODE_END_TOKEN_ID,
            "pad_token_id": self.tokenizer.pad_token_id,
        }
        default_gen_args.update(generation_args)

        # Generate tokens
        with torch.inference_mode():
            outputs = self.model.generate(**tokenizer_inputs, **default_gen_args)

        # Extract generated tokens (everything after the input prompt)
        generated_ids = outputs[0, tokenizer_inputs["input_ids"].shape[1]:].tolist()

        # Extract SNAC audio tokens
        snac_tokens = self.extract_snac_codes(generated_ids)
        
        if len(snac_tokens) < 7:
            return {"error": f"Not enough SNAC tokens generated: {len(snac_tokens)}"}

        # Unpack SNAC tokens to 3 hierarchical levels
        levels = self.unpack_snac_from_7(snac_tokens)
        
        # Convert to tensors
        codes_tensor = [
            torch.tensor(level, dtype=torch.long, device=self.device).unsqueeze(0)
            for level in levels
        ]

        # Generate final audio with SNAC decoder
        with torch.inference_mode():
            z_q = self.snac.quantizer.from_codes(codes_tensor)
            audio = self.snac.decoder(z_q)[0, 0].cpu().numpy()

        # Trim warmup samples (first 2048 samples)
        if len(audio) > 2048:
            audio = audio[2048:]

        # Save audio to WAV
        buf = io.BytesIO()
        self.sf.write(buf, audio, 24000, format="WAV")
        wav_bytes = buf.getvalue()
        b64 = base64.b64encode(wav_bytes).decode("utf-8")

        # Return in HuggingFace Inference Endpoints format
        # The endpoint expects a list with blob data
        return [{"blob": b64, "content-type": "audio/wav"}]
# Module-level convenience functions for hosting platforms (Hugging Face Endpoints)
# The platform typically expects top-level `init` and `predict` (or `run`) callables
# so we provide thin wrappers around the EndpointHandler class.
_HANDLER = None


def init(model_id: str = None):
    """Initialize the module-level handler.

    model_id: optional path or model identifier to pass to EndpointHandler.
    If omitted the handler will look for environment variables `HF_MODEL_ID` or
    `MODEL_ID` and otherwise instantiate with a blank path (which may be
    appropriate when the model files are bundled with the repo).
    """
    global _HANDLER
    if _HANDLER is not None:
        return

    model_path = model_id or os.getenv("HF_MODEL_ID") or os.getenv("MODEL_ID") or ""
    _HANDLER = EndpointHandler(path=model_path)


def predict(payload):
    """Predict/predict wrapper for hosted endpoints.

    Accepts a dict (recommended) or a JSON-ish payload. If a list is passed,
    the first element is used. If a plain string is provided it is treated as
    the `text` field.
    Returns the same dict structure produced by EndpointHandler.__call__.
    """
    global _HANDLER
    if _HANDLER is None:
        init()

    data = payload
    # handle list payloads (common in some platform wrappers)
    if isinstance(payload, (list, tuple)) and len(payload) > 0:
        data = payload[0]

    # allow raw JSON string -> dict
    if isinstance(data, str):
        try:
            import json

            data = json.loads(data)
        except Exception:
            # treat plain string as the text to synthesize
            data = {"text": data}

    if not isinstance(data, dict):
        # best-effort normalization
        data = {"text": str(data)}

    try:
        return _HANDLER(data)
    except Exception as e:
        return {"error": str(e)}


def run(payload):
    """Alias for predict for platforms that expect `run` entrypoint."""
    return predict(payload)