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 = "YOUR_HUGGINGFACE_TOKEN_HERE"  # 🔐 Replace with your real token
login(token=HF_TOKEN)
print("🔑 Logged into Hugging Face successfully.")

# ----------------- 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'
if not os.path.exists(folder_path):
    os.mkdir(folder_path)
    print(f"📦 Folder created: {folder_path}")
else:
    print(f"📦 Folder already exists: {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 (8-bit)...")

try:
    model = AutoModelForCausalLM.from_pretrained(
        "m-a-p/YuE-s1-1B-mini-en-cot",  # ✅ smaller model for CPU
        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)