import os, sys, torch, gradio as gr, tempfile, uuid, re, numpy as np, gc
import torchaudio
from torchaudio.transforms import Resample
from omegaconf import OmegaConf
from tqdm import tqdm
from einops import rearrange
from transformers import (
    AutoTokenizer, AutoModelForCausalLM,
    LogitsProcessor, LogitsProcessorList
)
from huggingface_hub import snapshot_download, login

# ----------------- Hugging Face Authentication -----------------
HF_TOKEN = os.getenv("HF_TOKEN")  # ✅ Read securely from environment variable
if HF_TOKEN:
    login(token=HF_TOKEN)
    print("🔑 Logged into Hugging Face successfully.")
else:
    print("⚠️ No Hugging Face token found. Please set it in your environment variables.")

# ----------------- Basic Environment Setup -----------------
print("🔧 Initializing environment (CPU mode)...")

os.environ["PYTORCH_JIT"] = "0"
torch.set_num_threads(1)
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.benchmark = False

device = "cpu"

# ----------------- Model Download -----------------
folder_path = './xcodec_mini_infer'
os.makedirs(folder_path, exist_ok=True)
print(f"📦 Folder ready: {folder_path}")

snapshot_download(
    repo_id="m-a-p/xcodec_mini_infer",
    local_dir="./xcodec_mini_infer",
    token=HF_TOKEN
)

sys.path.append(os.path.join(os.path.dirname(__file__), 'xcodec_mini_infer'))
sys.path.append(os.path.join(os.path.dirname(__file__), 'xcodec_mini_infer', 'descriptaudiocodec'))

from codecmanipulator import CodecManipulator
from mmtokenizer import _MMSentencePieceTokenizer
from models.soundstream_hubert_new import SoundStream

# ----------------- Load YuE Model -----------------
print("🧠 Loading YuE model on CPU (small version)...")
try:
    model = AutoModelForCausalLM.from_pretrained(
        "m-a-p/YuE-upsampler",
        torch_dtype=torch.float32,
        attn_implementation="eager",
        low_cpu_mem_usage=True,
        device_map={"": "cpu"},
        token=HF_TOKEN
    )
    model.eval()
    print("✅ YuE model loaded successfully.")
except Exception as e:
    print(f"❌ Model loading failed: {e}")
    sys.exit(1)

# ----------------- Load Codec -----------------
print("🎧 Loading codec model...")
basic_model_config = './xcodec_mini_infer/final_ckpt/config.yaml'
resume_path = './xcodec_mini_infer/final_ckpt/ckpt_00100000.pth'

mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
codectool = CodecManipulator("xcodec", 0, 1)
model_config = OmegaConf.load(basic_model_config)

codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
state_dict = torch.load(resume_path, map_location="cpu")
codec_model.load_state_dict(state_dict['codec_model'])
codec_model.eval()
print("✅ Codec model loaded successfully.")

# ---------------- Utility -----------------
class BlockTokenRangeProcessor(LogitsProcessor):
    def __init__(self, start_id, end_id):
        self.blocked_token_ids = list(range(start_id, end_id))
    def __call__(self, input_ids, scores):
        scores[:, self.blocked_token_ids] = -float("inf")
        return scores

def split_lyrics(lyrics: str):
    pattern = r"\[(\w+)\]\s*(.*?)(?=\s*\n\[|\Z)"
    segments = re.findall(pattern, lyrics, re.DOTALL)
    return [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]

def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
    limit = 0.99
    max_val = wav.abs().max()
    wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
    torchaudio.save(str(path), wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)

# ---------------- Generation -----------------
def generate_music(genre_txt, lyrics_txt, progress=gr.Progress()):
    with tempfile.TemporaryDirectory() as output_dir:
        genres = genre_txt.strip()
        lyrics = split_lyrics(lyrics_txt + "\n")
        prompt_texts = [f"Generate music from lyrics.\n[Genre] {genres}\n" + "\n".join(lyrics)]
        random_id = uuid.uuid4()

        print("🎶 Generating music...")

        start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
        end_of_segment = mmtokenizer.tokenize('[end_of_segment]')

        head_id = mmtokenizer.tokenize(prompt_texts[0])
        prompt_ids = torch.as_tensor(head_id + start_of_segment + [mmtokenizer.soa]).unsqueeze(0).to(device)

        with torch.inference_mode():
            output_seq = model.generate(
                input_ids=prompt_ids,
                max_new_tokens=32,
                top_p=0.9,
                temperature=0.8,
                repetition_penalty=1.1,
                eos_token_id=mmtokenizer.eoa,
                pad_token_id=mmtokenizer.eoa,
                logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002)]),
            )

        ids = output_seq[0].cpu().numpy()
        soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
        eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()

        vocals, instr = [], []
        for i in range(len(soa_idx)):
            codec_ids = ids[soa_idx[i] + 1:eoa_idx[i]]
            codec_ids = codec_ids[:2 * (len(codec_ids) // 2)]
            v = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[0])
            vocals.append(v)
            ins = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[1])
            instr.append(ins)

        vocals = np.concatenate(vocals, axis=1)
        instr = np.concatenate(instr, axis=1)
        mix = (vocals + instr) / 2

        recons_dir = os.path.join(output_dir, "recons")
        os.makedirs(recons_dir, exist_ok=True)
        mix_path = os.path.join(recons_dir, f"mix_{random_id}.wav")

        save_audio(torch.tensor(mix).unsqueeze(0), mix_path, 16000)
        gc.collect()

        print("✅ Generation complete.")
        return mix_path

# ---------------- Gradio Interface -----------------
with gr.Blocks() as demo:
    gr.Markdown("# 🎵 YuE CPU Music Generator (Low-Memory Edition)")
    genre_txt = gr.Textbox(label="Genre", placeholder="Pop, Jazz, Hip-Hop...")
    lyrics_txt = gr.Textbox(label="Lyrics", placeholder="[Verse]\nI walk the night...")
    music_out = gr.Audio(label="Generated Song")
    btn = gr.Button("Generate 🎶")
    btn.click(generate_music, inputs=[genre_txt, lyrics_txt], outputs=[music_out])

demo.queue().launch(show_error=True, share=True)