import os
import sys

# Ensure both current folder and absolute folder paths are visible to Python
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, "/home/user/app")
os.environ["GRADIO_SSR_MODE"] = "0"
os.environ["PYTHONPATH"] = f".:{os.environ.get('PYTHONPATH', '')}"

# Mock distribution versions to satisfy importlib check on HF
try:
    import importlib.metadata
    orig_version = importlib.metadata.version
    def fake_version(pkg_name):
        if pkg_name == "omnivoice":
            return "0.1.3"
        return orig_version(pkg_name)
    importlib.metadata.version = fake_version
except Exception:
    pass

# Patch HfFolder back into huggingface_hub before gradio/spaces import
import huggingface_hub
if not hasattr(huggingface_hub, 'HfFolder'):
    class HfFolder:
        @staticmethod
        def get_token():
            return huggingface_hub.get_token() if hasattr(huggingface_hub, 'get_token') else None
        @staticmethod
        def save_token(token):
            pass
        @staticmethod
        def delete_token():
            pass
    huggingface_hub.HfFolder = HfFolder

import torch
import spaces  # Import Hugging Face ZeroGPU SDK

# We load the model inside a class wrapper to lazy-load or handle global imports safely
model = None

def get_model():
    global model
    if model is None:
        print("Loading OmniVoice model globally inside ZeroGPU context...")
        from omnivoice.models.omnivoice import OmniVoice
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model = OmniVoice.from_pretrained(
            "k2-fsa/OmniVoice",
            device_map=device,
            dtype=torch.float16 if device == "cuda" else torch.float32,
            load_asr=True
        )
    return model

# Define the generator function decorated with @spaces.GPU
@spaces.GPU
def gpu_generate_fn(
    text,
    language,
    ref_audio,
    instruct,
    num_step,
    guidance_scale,
    denoise,
    speed,
    duration,
    preprocess_prompt,
    postprocess_output,
    mode,
    ref_text=None,
):
    from omnivoice.models.omnivoice import OmniVoiceGenerationConfig
    import time
    import numpy as np
    import scipy.io.wavfile as wavfile

    # Load model instance
    m = get_model()

    # Ensure model is on CUDA inside the ZeroGPU container
    if hasattr(m, "to") and torch.cuda.is_available():
        m.to("cuda")

    gen_config = OmniVoiceGenerationConfig(
        num_step=int(num_step or 32),
        guidance_scale=float(guidance_scale) if guidance_scale is not None else 2.0,
        denoise=bool(denoise) if denoise is not None else True,
        preprocess_prompt=bool(preprocess_prompt),
        postprocess_output=bool(postprocess_output),
    )

    lang = language if (language and language != "Auto") else None
    kw = dict(text=text.strip(), language=lang, generation_config=gen_config)

    if speed is not None and float(speed) != 1.0:
        kw["speed"] = float(speed)
    if duration is not None and float(duration) > 0:
        kw["duration"] = float(duration)

    try:
        if mode == "clone":
            if not ref_audio:
                return None, "Please upload a reference audio.", None
            kw["voice_clone_prompt"] = m.create_voice_clone_prompt(
                ref_audio=ref_audio,
                ref_text=ref_text,
            )

        if instruct and instruct.strip():
            kw["instruct"] = instruct.strip()

        audio = m.generate(**kw)
        waveform = audio[0].squeeze(0).cpu().numpy()
        waveform = (waveform * 32767).astype(np.int16)
        
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        save_path = os.path.join("outputs", f"RJD_{timestamp}.wav")
        os.makedirs("outputs", exist_ok=True)
        wavfile.write(save_path, m.sampling_rate, waveform)
        
        return (m.sampling_rate, waveform), f"Done. Saved to: {save_path}", save_path
    except Exception as e:
        import traceback
        traceback.print_exc()
        return None, f"Error: {type(e).__name__}: {e}", None

def main():
    print("Initializing Gradio Interface...")
    from omnivoice.cli.demo import build_demo
    # Ensure model is initialized for building Gradio components
    m = get_model()
    demo = build_demo(m, checkpoint="k2-fsa/OmniVoice", generate_fn=gpu_generate_fn)
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)

if __name__ == "__main__":
    main()