import random
import argparse
import json
import os
import time
from concurrent.futures import ThreadPoolExecutor

import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
import tqdm
from huggingface_hub import hf_hub_download
from transformers import DynamicCache

import MIDI
from midi_model import MIDIModel
from midi_synthesizer import MidiSynthesizer

MAX_SEED = np.iinfo(np.int32).max
in_space = os.getenv("SYSTEM") == "spaces"

# ======================= JS INLINE =======================
FIXED_JS = """
const MIDI_OUTPUT_BATCH_SIZE=2;
function deepQuerySelector(selector){
  function deepSearch(root,selector){
    let el=root.querySelector(selector);
    if(el)return el;
    const hosts=root.querySelectorAll('*');
    for(let h of hosts){
      if(h.shadowRoot){
        el=deepSearch(h.shadowRoot,selector);
        if(el)return el;
      }
    }
    return null;
  }
  return deepSearch(this,selector);
}
Element.prototype.deepQuerySelector=deepQuerySelector;
Document.prototype.deepQuerySelector=deepQuerySelector;
function gradioApp(){
  const elems=document.getElementsByTagName('gradio-app');
  const elem=elems.length==0?document:elems[0];
  return elem.shadowRoot?elem.shadowRoot:elem;
}
window.onUiLoaded=window.onUiLoaded||function(cb){
  if(document.readyState==='loading')document.addEventListener('DOMContentLoaded',cb);
  else cb();
};
const msgReceiveCallbacks=[];
function onMsgReceive(cb){msgReceiveCallbacks.push(cb);}
function executeCallbacks(cbs,...args){
  cbs.forEach(cb=>{try{cb(...args);}catch(e){console.error(e);}});
}
(function(){
  let midi_visualizers=[],audio_players=[];
  onUiLoaded(()=>{
    const c=gradioApp();
    for(let i=0;i<MIDI_OUTPUT_BATCH_SIZE;i++){
      let div=c.deepQuerySelector(`#midi_visualizer_container_${i}`);
      if(!div)continue;
      div.innerHTML=`<div class="progressDiv"><div class="progress">0/0</div></div><midi-visualizer class="midi-visualizer" id="midi_visualizer_${i}"></midi-visualizer>`;
      let vis=c.deepQuerySelector(`#midi_visualizer_${i}`);
      let aud=c.deepQuerySelector(`#midi_audio_${i} audio`);
      if(vis){
        vis.config={noteHeight:2,pixelsPerTimeStep:40,noteSpacing:1,noteRGB:'136,212,152',activeNoteRGB:'168,230,207',windowless:true};
        if(aud)vis.bindAudioPlayer(aud);
        midi_visualizers.push(vis);
        audio_players.push(aud);
      }
    }
  });
  function setProgress(p,t){
    let par=gradioApp().deepQuerySelector(".progressDiv");
    if(!par)return;
    let inner=par.querySelector(".progress");
    if(t===0)t=1;
    inner.style.width=`${(p/t)*100}%`;
    inner.textContent=`${p}/${t}`;
  }
  onMsgReceive(msgs=>{
    for(let m of msgs){
      if(m instanceof Array)m.forEach(o=>handleMsg(o));
      else handleMsg(m);
    }
  });
  function handleMsg(msg){
    let idx;
    switch(msg.name){
      case"visualizer_clear": idx=msg.data[0];let ver=msg.data[1];if(midi_visualizers[idx]){midi_visualizers[idx].clearMidiEvents(false);midi_visualizers[idx].version=ver;}break;
      case"visualizer_append": idx=msg.data[0];let evs=msg.data[1];if(midi_visualizers[idx]){evs.forEach(v=>midi_visualizers[idx].appendMidiEvent(v));}break;
      case"visualizer_end": idx=msg.data;if(midi_visualizers[idx]){midi_visualizers[idx].finishAppendMidiEvent();midi_visualizers[idx].setPlayTime(0);}break;
      case"progress": setProgress(msg.data[0],msg.data[1]);break;
    }
  }
})();
"""

# ======================= CSS =======================
css = """
:root { 
    --primary-600: #a8e6cf !important; 
    --accent-600: #a8e6cf !important;
    --ring-color: #a8e6cf !important;
    --checkbox-label-background-selected: #e1f0e5 !important;
    --button-primary-background-fill: #a8e6cf !important;
    --button-primary-background-fill-hover: #88d498 !important;
}
.gradio-container, .gradio-container * { --loader-color: #dcedc1 !important; box-shadow: none !important; }
*:focus { border-color: #a8e6cf !important; box-shadow: 0 0 0 2px #e1f0e5 !important; }
.ba-header-container { border: 1.5px solid #e1f0e5; border-radius: 12px; margin-bottom: 12px; background: white; overflow: hidden; line-height: 0; }
.ba-header-container img { width: 100%; height: auto; }
.status-container { border: 1.5px solid #e1f0e5; border-radius: 12px; padding: 15px 22px; margin-bottom: 20px; background: white; }
.status-title { color: #88d498 !important; font-weight: 800; font-size: 16px; margin-bottom: 8px; }
.text-green-bold { color: #88d498 !important; font-weight: 900 !important; }
.dot { height: 8px; width: 8px; background-color: #88d498; border-radius: 50%; display: inline-block; margin-right: 5px; box-shadow: 0 0 0 0 rgba(136,212,152,1); animation: pulse-green 2s infinite; }
@keyframes pulse-green { 0% { transform: scale(0.95); box-shadow: 0 0 0 0 rgba(136,212,152,0.7); } 70% { transform: scale(1); box-shadow: 0 0 0 10px rgba(136,212,152,0); } 100% { transform: scale(0.95); box-shadow: 0 0 0 0 rgba(136,212,152,0); } }
.slim-card { max-width: 700px; margin: 0 auto; background: transparent; padding: 10px; }

/* Grid untuk Select Model - horizontal scroll (ke samping) */
.model-horizontal-grid {
    display: flex;
    flex-direction: row;
    gap: 10px;
    padding: 12px;
    background: white;
    border: 1px solid #e9f0ec;
    border-radius: 14px;
    overflow-x: auto;
    overflow-y: hidden;
    white-space: nowrap;
    box-sizing: border-box;
    margin-bottom: 16px;
    box-shadow: 0 1px 2px rgba(0,0,0,0.02);
    scrollbar-width: thin;
}
.model-card {
    flex: 0 0 auto;
    width: 220px;
    border: 1px solid #e2ece6;
    border-radius: 12px;
    padding: 8px 12px;
    background: #fff;
    display: inline-flex;
    flex-direction: column;
    justify-content: center;
    cursor: pointer;
    transition: all 0.2s ease;
    text-align: left;
    font-family: inherit;
    box-sizing: border-box;
}
.model-card:hover {
    background: #f7fdf9;
    border-color: #b8decb;
    transform: translateY(-1px);
}
.model-card.selected {
    background: #ecf6f1;
    border-left: 3px solid #88d498;
    border-color: #c8e6d9;
}
.model-title {
    font-weight: 600;
    color: #3a5e4f;
    font-size: 12px;
    margin-bottom: 2px;
    white-space: normal;
    word-break: break-word;
}
.model-sub {
    color: #9bb7a8;
    font-size: 9px;
    font-weight: 500;
    letter-spacing: 0.3px;
}

/* Grid untuk Instruments - vertikal scroll (1 kolom) */
.compact-grid {
    display: grid;
    grid-template-columns: 1fr;
    gap: 8px;
    padding: 12px;
    background: white;
    border: 1px solid #e9f0ec;
    border-radius: 14px;
    max-height: 260px;
    overflow-y: auto;
    box-sizing: border-box;
    margin-bottom: 16px;
    box-shadow: 0 1px 2px rgba(0,0,0,0.02);
}

/* Grid untuk Drum Kit - 2 kolom (vertical scroll) */
.drum-grid {
    display: grid;
    grid-template-columns: repeat(2, 1fr);
    gap: 8px;
    padding: 12px;
    background: white;
    border: 1px solid #e9f0ec;
    border-radius: 14px;
    max-height: 200px;
    overflow-y: auto;
    box-sizing: border-box;
    margin-bottom: 16px;
    box-shadow: 0 1px 2px rgba(0,0,0,0.02);
}

.card-btn {
    border: 1px solid #e2ece6;
    border-radius: 12px;
    padding: 8px 12px;
    background: #fff;
    display: flex;
    flex-direction: column;
    justify-content: center;
    cursor: pointer;
    transition: all 0.2s ease;
    text-align: left;
    width: 100%;
    font-family: inherit;
    box-sizing: border-box;
}
.card-btn:hover {
    background: #f7fdf9;
    border-color: #b8decb;
    transform: translateY(-1px);
}
.card-btn.selected {
    background: #ecf6f1;
    border-left: 3px solid #88d498;
    border-color: #c8e6d9;
}
.card-title {
    font-weight: 600;
    color: #3a5e4f;
    font-size: 12px;
    margin-bottom: 2px;
    white-space: nowrap;
    overflow: hidden;
    text-overflow: ellipsis;
}
.card-sub {
    color: #9bb7a8;
    font-size: 9px;
    font-weight: 500;
    letter-spacing: 0.3px;
}
.section-header {
    font-size: 13px;
    font-weight: 700;
    color: #5c7c6b;
    margin: 8px 0 6px 0;
    letter-spacing: 0.5px;
    display: flex;
    align-items: center;
    gap: 6px;
}
.warning-card { background: #fffcf0; border: 2px dashed #ffd3b6; border-radius: 10px; padding: 12px; margin-bottom: 15px; text-align: center; }
.gen-btn { background: #a8e6cf !important; color: #4a5568 !important; font-weight: 700 !important; border-radius: 12px !important; height: 45px !important; width: 100%; border: none !important; cursor: pointer; }
.credit-footer { margin-top: 25px; padding: 15px; background: white; border-radius: 12px; text-align: center; border-bottom: 4px solid #a8e6cf; color: #94a3b8; font-weight: 700; font-size: 12px; letter-spacing: 2px; }
.diary-container { display: flex; flex-direction: column; align-items: center; margin: 20px 0 15px; padding: 12px; background: #fefef7; border-radius: 24px; border: 1px solid #e1f0e5; }
.diary-text { font-family: 'Georgia', serif; font-size: 18px; font-weight: 500; color: #88d498; margin-bottom: 8px; letter-spacing: 1px; }
.gif-wrapper { display: flex; justify-content: center; margin-top: 5px; }
.gif-wrapper img { max-width: 100%; border-radius: 20px; border: 1px solid #e1f0e5; }

/* Tombol Undo */
.undo-btn {
    background: #a8e6cf !important;
    color: white !important;
    font-weight: 700 !important;
    border-radius: 10px !important;
    border: none !important;
    padding: 8px 16px !important;
    margin-top: 8px !important;
}
.undo-btn:hover {
    background: #88d498 !important;
}
"""

# ======================= FUNGSI GENERATE (TIDAK BERUBAH) =======================
@torch.inference_mode()
def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98, top_k=20,
             disable_patch_change=False, disable_control_change=False, disable_channels=None, generator=None):
    tokenizer = model.tokenizer
    if disable_channels is not None:
        disable_channels = [tokenizer.parameter_ids["channel"][c] for c in disable_channels]
    else:
        disable_channels = []
    max_token_seq = tokenizer.max_token_seq
    if prompt is None:
        input_tensor = torch.full((1, max_token_seq), tokenizer.pad_id, dtype=torch.long, device=model.device)
        input_tensor[0,0] = tokenizer.bos_id
        input_tensor = input_tensor.unsqueeze(0).repeat(batch_size,1,1)
    else:
        if len(prompt.shape)==2:
            prompt = prompt[None,:]
            prompt = np.repeat(prompt, repeats=batch_size, axis=0)
        elif prompt.shape[0]==1:
            prompt = np.repeat(prompt, repeats=batch_size, axis=0)
        elif len(prompt.shape)!=3 or prompt.shape[0]!=batch_size:
            raise ValueError(f"invalid shape for prompt, {prompt.shape}")
        prompt = prompt[..., :max_token_seq]
        if prompt.shape[-1] < max_token_seq:
            prompt = np.pad(prompt, ((0,0),(0,0),(0, max_token_seq - prompt.shape[-1])),
                            mode="constant", constant_values=tokenizer.pad_id)
        input_tensor = torch.from_numpy(prompt).to(dtype=torch.long, device=model.device)
    cur_len = input_tensor.shape[1]
    bar = tqdm.tqdm(desc="generating", total=max_len - cur_len, disable=in_space)
    cache1 = DynamicCache()
    past_len = 0
    with bar:
        while cur_len < max_len:
            end = [False]*batch_size
            hidden = model.forward(input_tensor[:, past_len:], cache=cache1)[:,-1]
            next_token_seq = None
            event_names = [""]*batch_size
            cache2 = DynamicCache()
            for i in range(max_token_seq):
                mask = torch.zeros((batch_size, tokenizer.vocab_size), dtype=torch.int64, device=model.device)
                for b in range(batch_size):
                    if end[b]:
                        mask[b, tokenizer.pad_id] = 1
                        continue
                    if i == 0:
                        mask_ids = list(tokenizer.event_ids.values()) + [tokenizer.eos_id]
                        if disable_patch_change: mask_ids.remove(tokenizer.event_ids["patch_change"])
                        if disable_control_change: mask_ids.remove(tokenizer.event_ids["control_change"])
                        mask[b, mask_ids] = 1
                    else:
                        param_names = tokenizer.events[event_names[b]]
                        if i > len(param_names):
                            mask[b, tokenizer.pad_id] = 1
                            continue
                        param_name = param_names[i-1]
                        mask_ids = tokenizer.parameter_ids[param_name]
                        if param_name == "channel":
                            mask_ids = [i for i in mask_ids if i not in disable_channels]
                        mask[b, mask_ids] = 1
                mask = mask.unsqueeze(1)
                x = next_token_seq
                if i != 0:
                    hidden = None
                    x = x[:, -1:]
                logits = model.forward_token(hidden, x, cache=cache2)[:,-1:]
                scores = torch.softmax(logits/temp, dim=-1) * mask
                samples = model.sample_top_p_k(scores, top_p, top_k, generator=generator)
                if i == 0:
                    next_token_seq = samples
                    for b in range(batch_size):
                        if end[b]: continue
                        eid = samples[b].item()
                        if eid == tokenizer.eos_id:
                            end[b] = True
                        else:
                            event_names[b] = tokenizer.id_events[eid]
                else:
                    next_token_seq = torch.cat([next_token_seq, samples], dim=1)
                    if all([len(tokenizer.events[event_names[b]]) == i for b in range(batch_size) if not end[b]]):
                        break
            if next_token_seq.shape[1] < max_token_seq:
                next_token_seq = F.pad(next_token_seq, (0, max_token_seq - next_token_seq.shape[1]), "constant", value=tokenizer.pad_id)
            next_token_seq = next_token_seq.unsqueeze(1)
            input_tensor = torch.cat([input_tensor, next_token_seq], dim=1)
            past_len = cur_len
            cur_len += 1
            bar.update(1)
            yield next_token_seq[:,0].cpu().numpy()
            if all(end): break

def create_msg(name,data): return {"name":name,"data":data}
def send_msgs(msgs): return json.dumps(msgs)

# ======================= RUN =======================
def run(model_name, tab, mid_seq, continuation_state, continuation_select, 
        selected_instruments, drum_kit, bpm, time_sig, key_sig_str,
        mid, midi_events, reduce_cc_st, remap_track_channel, add_default_instr, remove_empty_channels,
        seed, seed_rand, gen_events, temp, top_p, top_k, allow_cc, progress=gr.Progress()):
    try:
        progress(0, desc="Memuat model...")
        model = models[model_name]
        model.to(device=opt.device)
        tokenizer = model.tokenizer
        bpm = int(bpm)
        if time_sig == "auto":
            time_sig = None; time_sig_nn = 4; time_sig_dd = 2
        else:
            time_sig_nn, time_sig_dd = time_sig.split('/'); time_sig_nn = int(time_sig_nn); time_sig_dd = {2:1,4:2,8:3}[int(time_sig_dd)]
        
        key_sig = key_sig_str
        if key_sig == "auto":
            key_sig = None
            key_sig_sf = 0; key_sig_mi = 0
        else:
            key_index = key_sig_to_index[key_sig]
            key_sig = key_index
            key_sig_sf = (key_sig - 1) // 2 - 7
            key_sig_mi = (key_sig - 1) % 2
        
        gen_events = int(gen_events)
        max_len = gen_events
        if seed_rand: seed = random.randint(0, MAX_SEED)
        generator = torch.Generator(opt.device).manual_seed(seed)
        disable_patch_change = False
        disable_channels = None
        
        continuation_select_idx = 0 if continuation_select == "all" else int(continuation_select.replace("output", ""))
        
        progress(0.05, desc="Menyiapkan prompt...")
        if tab == 0:
            i = 0
            mid = [[tokenizer.bos_id] + [tokenizer.pad_id] * (tokenizer.max_token_seq - 1)]
            if tokenizer.version == "v2":
                if time_sig is not None: mid.append(tokenizer.event2tokens(["time_signature",0,0,0,time_sig_nn-1,time_sig_dd-1]))
                if key_sig is not None: mid.append(tokenizer.event2tokens(["key_signature",0,0,0,key_sig_sf+7,key_sig_mi]))
            if bpm != 0: mid.append(tokenizer.event2tokens(["set_tempo",0,0,0,bpm]))
            patches = {}
            instruments = selected_instruments
            for instr in instruments: patches[i] = patch2number[instr]; i = (i+1) if i != 8 else 10
            if drum_kit != "None": patches[9] = drum_kits2number[drum_kit]
            for i,(c,p) in enumerate(patches.items()): mid.append(tokenizer.event2tokens(["patch_change",0,0,i+1,c,p]))
            mid = np.asarray([mid] * OUTPUT_BATCH_SIZE, dtype=np.int64)
            mid_seq = mid.tolist()
            if len(instruments) > 0: disable_patch_change = True; disable_channels = [i for i in range(16) if i not in patches]
        elif tab == 1 and mid is not None:
            eps = 4 if reduce_cc_st else 0
            mid = tokenizer.tokenize(MIDI.midi2score(mid), cc_eps=eps, tempo_eps=eps,
                                     remap_track_channel=remap_track_channel,
                                     add_default_instr=add_default_instr, remove_empty_channels=remove_empty_channels)
            mid = mid[:int(midi_events)]; mid = np.asarray([mid] * OUTPUT_BATCH_SIZE, dtype=np.int64); mid_seq = mid.tolist()
        elif tab == 2 and mid_seq is not None:
            mid = np.asarray(mid_seq, dtype=np.int64)
            if continuation_select_idx > 0:
                continuation_state.append(mid_seq); mid = np.repeat(mid[continuation_select_idx-1:continuation_select_idx], repeats=OUTPUT_BATCH_SIZE, axis=0); mid_seq = mid.tolist()
            else: continuation_state.append(mid.shape[1])
        else:
            continuation_state = [0]; mid = [[tokenizer.bos_id] + [tokenizer.pad_id] * (tokenizer.max_token_seq - 1)]; mid = np.asarray([mid] * OUTPUT_BATCH_SIZE, dtype=np.int64); mid_seq = mid.tolist()
        if mid is not None: max_len += mid.shape[1]
        
        init_msgs = [create_msg("progress", [0, gen_events])]
        if not (tab == 2 and continuation_select_idx == 0):
            for i in range(OUTPUT_BATCH_SIZE):
                events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]]
                init_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]), create_msg("visualizer_append", [i, events])]
        yield mid_seq, continuation_state, seed, send_msgs(init_msgs)
        
        progress(0.1, desc=f"Generating {gen_events} events...")
        midi_generator = generate(model, mid, batch_size=OUTPUT_BATCH_SIZE, max_len=max_len, temp=temp,
                                  top_p=top_p, top_k=top_k, disable_patch_change=disable_patch_change,
                                  disable_control_change=not allow_cc, disable_channels=disable_channels, generator=generator)
        events = [list() for _ in range(OUTPUT_BATCH_SIZE)]
        t = time.time() + 0.2
        total_steps = gen_events
        for step, token_seqs in enumerate(midi_generator):
            token_seqs = token_seqs.tolist()
            for j in range(OUTPUT_BATCH_SIZE):
                token_seq = token_seqs[j]; mid_seq[j].append(token_seq); events[j].append(tokenizer.tokens2event(token_seq))
            if time.time() - t > 0.2:
                progress((step+1)/total_steps, desc=f"Generating {step+1}/{total_steps}")
                msgs = [create_msg("progress", [step+1, gen_events])]
                for j in range(OUTPUT_BATCH_SIZE):
                    msgs += [create_msg("visualizer_append", [j, events[j]])]
                    events[j] = []
                yield mid_seq, continuation_state, seed, send_msgs(msgs)
                t = time.time()
        yield mid_seq, continuation_state, seed, send_msgs([])
        progress(1.0, desc="Selesai!")
        print("Generation finished.")
    except Exception as e:
        print(f"ERROR in run: {e}")
        import traceback
        traceback.print_exc()
        yield mid_seq if 'mid_seq' in locals() else None, continuation_state if 'continuation_state' in locals() else [0], seed if 'seed' in locals() else 0, send_msgs([create_msg("progress", [0,0])])

# ======================= FINISH RUN & RENDER AUDIO =======================
def finish_run(model_name, mid_seq):
    print("Starting finish_run...")
    if mid_seq is None:
        print("mid_seq is None")
        outputs = [None] * OUTPUT_BATCH_SIZE
        return *outputs, []
    tokenizer = models[model_name].tokenizer
    outputs = []
    end_msgs = [create_msg("progress",[0,0])]
    out_dir = "outputs"
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
        print(f"Created outputs directory: {out_dir}")
    for i in range(OUTPUT_BATCH_SIZE):
        events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]]
        mid = tokenizer.detokenize(mid_seq[i])
        out_path = os.path.join(out_dir, f"output{i+1}.mid")
        with open(out_path, 'wb') as f:
            f.write(MIDI.score2midi(mid))
        outputs.append(out_path)
        print(f"Saved MIDI to {out_path}")
        end_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]), create_msg("visualizer_append", [i, events]), create_msg("visualizer_end", i)]
    print("finish_run completed, returning outputs:", outputs)
    return *outputs, send_msgs(end_msgs)

def synthesis_task(mid):
    return synthesizer.synthesis(MIDI.score2opus(mid))

def render_audio(model_name, mid_seq, should_render_audio):
    print(f"render_audio called, should_render_audio={should_render_audio}, mid_seq is {type(mid_seq)}")
    if (not should_render_audio) or mid_seq is None:
        outputs = [None] * OUTPUT_BATCH_SIZE
        print("No audio to render")
        return tuple(outputs)
    tokenizer = models[model_name].tokenizer
    outputs = []
    audio_futures = []
    for i in range(OUTPUT_BATCH_SIZE):
        mid = tokenizer.detokenize(mid_seq[i])
        audio_futures.append(thread_pool.submit(synthesis_task, mid))
    for future in audio_futures:
        audio_data = future.result()
        if isinstance(audio_data, np.ndarray):
            audio_data = audio_data.astype(np.float32)
        outputs.append((44100, audio_data))
    print(f"Rendered {len(outputs)} audio files")
    if OUTPUT_BATCH_SIZE == 1:
        return outputs[0]
    return tuple(outputs)

def undo_continuation(model_name, mid_seq, continuation_state):
    if mid_seq is None or len(continuation_state) < 2:
        return mid_seq, continuation_state, send_msgs([])
    tokenizer = models[model_name].tokenizer
    if isinstance(continuation_state[-1], list):
        mid_seq = continuation_state[-1]
    else:
        mid_seq = [ms[:continuation_state[-1]] for ms in mid_seq]
    continuation_state = continuation_state[:-1]
    end_msgs = [create_msg("progress",[0,0])]
    for i in range(OUTPUT_BATCH_SIZE):
        events = [tokenizer.tokens2event(tokens) for tokens in mid_seq[i]]
        end_msgs += [create_msg("visualizer_clear", [i, tokenizer.version]), create_msg("visualizer_append", [i, events]), create_msg("visualizer_end", i)]
    return mid_seq, continuation_state, send_msgs(end_msgs)

def hf_hub_download_retry(repo_id, filename):
    print(f"downloading {repo_id} {filename}")
    retry=0; err=None
    while retry<30:
        try: return hf_hub_download(repo_id=repo_id, filename=filename)
        except Exception as e: err=e; retry+=1
    if err: raise err

# ======================= MAIN =======================
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--share", action="store_true", default=False)
    parser.add_argument("--port", type=int, default=7860)
    default_device = "cuda" if torch.cuda.is_available() else "cpu"
    parser.add_argument("--device", type=str, default=default_device)
    parser.add_argument("--batch", type=int, default=2)
    parser.add_argument("--max-gen", type=int, default=512)
    opt = parser.parse_args()
    
    if opt.device == "cuda" and not torch.cuda.is_available():
        print("⚠️ CUDA tidak tersedia, beralih ke CPU")
        opt.device = "cpu"
    
    OUTPUT_BATCH_SIZE = opt.batch
    print(f"Using device: {opt.device}, batch size: {OUTPUT_BATCH_SIZE}")
    
    if opt.device == "cpu" and "LD_LIBRARY_PATH" in os.environ:
        print("⚠️ Menghapus LD_LIBRARY_PATH karena menggunakan CPU")
    
    soundfont_path = hf_hub_download_retry(repo_id="skytnt/midi-model", filename="soundfont.sf2")
    thread_pool = ThreadPoolExecutor(max_workers=OUTPUT_BATCH_SIZE)
    synthesizer = MidiSynthesizer(soundfont_path)

    # ========= MEMUAT SEMUA MODEL (SEPERTI SEMULA) =========
    models_info = {
        "generic pretrain model (tv2o-medium) by skytnt": [
            "skytnt/midi-model-tv2o-medium", {
                "jpop": "skytnt/midi-model-tv2om-jpop-lora",
                "touhou": "skytnt/midi-model-tv2om-touhou-lora"
            }
        ],
        "generic pretrain model (tv2o-large) by asigalov61": ["asigalov61/Music-Llama", {}],
        "generic pretrain model (tv2o-medium) by asigalov61": ["asigalov61/Music-Llama-Medium", {}],
        "generic pretrain model (tv1-medium) by skytnt": ["skytnt/midi-model", {}]
    }
    models = {}
    if opt.device == "cuda":
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.backends.cuda.enable_mem_efficient_sdp(True)
        torch.backends.cuda.enable_flash_sdp(True)
    for name, (repo_id, loras) in models_info.items():
        print(f"Loading {name}...")
        model = MIDIModel.from_pretrained(repo_id)
        model.to(device="cpu", dtype=torch.float32)
        models[name] = model
        for lora_name, lora_repo in loras.items():
            print(f"loading lora {lora_repo} for {name}")
            model = MIDIModel.from_pretrained(repo_id)
            model = model.load_merge_lora(lora_repo)
            model.to(device="cpu", dtype=torch.float32)
            models[f"{name} with {lora_name} lora"] = model

    number2drum_kits = {-1:"None",0:"Standard",8:"Room",16:"Power",24:"Electric",25:"TR-808",32:"Jazz",40:"Blush",48:"Orchestra"}
    patch2number = {v:k for k,v in MIDI.Number2patch.items()}
    drum_kits2number = {v:k for k,v in number2drum_kits.items()}
    key_signatures = ['C♭','A♭m','G♭','E♭m','D♭','B♭m','A♭','Fm','E♭','Cm','B♭','Gm','F','Dm','C','Am','G','Em','D','Bm','A','F♯m','E','C♯m','B','G♯m','F♯','D♯m','C♯','A♯m']
    key_sig_to_index = {name: i+1 for i, name in enumerate(key_signatures)}
    
    instrument_list = list(patch2number.keys())
    drum_kit_list = list(number2drum_kits.values())

    # ========= GRADIO UI =========
    with gr.Blocks(css=css, title="Midi Music Generator") as app:
        gr.HTML(f"<script>{FIXED_JS}</script>")
        
        with gr.Column(elem_classes="slim-card"):
            gr.HTML("""
                <div class="ba-header-container">
                    <img src="https://huggingface.co/spaces/Plana-Archive/MOE-TTS/resolve/main/Mutsumi.jpg" alt="Banner">
                </div>
                <div class="status-container">
                    <div class="status-title">System Status</div>
                    <div><span style="color:#4a5568">Engine :</span> <span class="text-green-bold">LOADED ACTIVE ✅</span></div>
                    <div><span style="color:#4a5568">System :</span> <span style="color:#88d498; font-weight:700;"><span class="dot"></span>Online</span></div>
                    <div><span style="color:#4a5568">Status :</span> <span style="color:#88d498; font-weight:700;">Free</span></div>
                </div>
            """)
            gr.Markdown("### 🎹 Midi Music Generator")
            
            js_msg = gr.Textbox(elem_id="msg_receiver", visible=False)
            js_msg.change(None, [js_msg], [], js="(msg_json) => { let msgs = JSON.parse(msg_json); executeCallbacks(msgReceiveCallbacks, msgs); return []; }")
            
            # ========= SELECT MODEL (HORIZONTAL SCROLL) =========
            gr.HTML('<div class="section-header"><span>🎼</span> Select Model</div>')
            model_state = gr.State(value=list(models.keys())[0])
            model_grid_html = gr.HTML("")
            
            def build_model_horizontal_grid():
                html = '<div class="model-horizontal-grid" id="model-grid">'
                model_names = list(models.keys())
                for i, name in enumerate(model_names):
                    selected_attr = 'selected' if i == 0 else ''
                    html += f'''
                    <div data-model="{name}" class="model-card {selected_attr}" onclick="
                        document.querySelectorAll('#model-grid .model-card').forEach(card => card.classList.remove('selected'));
                        this.classList.add('selected');
                        const modelName = this.getAttribute('data-model');
                        const gradioApp = document.querySelector('gradio-app').shadowRoot || document;
                        const modelStateInput = gradioApp.querySelector('#model_state_input');
                        if(modelStateInput) modelStateInput.value = modelName;
                        modelStateInput.dispatchEvent(new Event('input', {{bubbles: true}}));
                    ">
                        <div class="model-title">{name}</div>
                        <div class="model-sub">click to select</div>
                    </div>
                    '''
                html += '</div>'
                return html
            
            model_grid_html.value = build_model_horizontal_grid()
            model_state_input = gr.Textbox(visible=False, elem_id="model_state_input")
            model_state_input.change(fn=lambda x: x, inputs=[model_state_input], outputs=[model_state])
            
            # ========= INSTRUMENTS (VERTICAL SCROLL, 1 KOLOM) =========
            gr.HTML('<div class="section-header"><span>🎸</span> Instruments (multiselect)</div>')
            instr_state = gr.State(value=[])
            instr_grid_html = gr.HTML("")
            
            def build_instr_grid(selected_list):
                html = '<div class="compact-grid" id="instr-grid">'
                for name in instrument_list:
                    selected_class = ' selected' if name in selected_list else ''
                    html += f'''
                    <div data-instr="{name}" class="card-btn{selected_class}" onclick="
                        let btn = this;
                        let isSelected = btn.classList.contains('selected');
                        if(isSelected) btn.classList.remove('selected');
                        else btn.classList.add('selected');
                        const selectedNames = Array.from(document.querySelectorAll('#instr-grid .card-btn.selected')).map(b => b.getAttribute('data-instr'));
                        const gradioApp = document.querySelector('gradio-app').shadowRoot || document;
                        const instrInput = gradioApp.querySelector('#instr_state_input');
                        if(instrInput) instrInput.value = JSON.stringify(selectedNames);
                        instrInput.dispatchEvent(new Event('input', {{bubbles: true}}));
                    ">
                        <div class="card-title">{name}</div>
                        <div class="card-sub">instrument</div>
                    </div>
                    '''
                html += '</div>'
                return html
            
            instr_grid_html.value = build_instr_grid([])
            instr_state_input = gr.Textbox(visible=False, elem_id="instr_state_input")
            instr_state_input.change(fn=lambda x: json.loads(x) if x else [], inputs=[instr_state_input], outputs=[instr_state])
            
            # ========= DRUM KIT (GRID 2 KOLOM, VERTICAL SCROLL) =========
            gr.HTML('<div class="section-header"><span>🥁</span> Drum Kit</div>')
            drum_state = gr.State(value="None")
            drum_grid_html = gr.HTML("")
            
            def build_drum_grid(selected):
                html = '<div class="drum-grid" id="drum-grid">'
                for name in drum_kit_list:
                    selected_class = ' selected' if name == selected else ''
                    html += f'''
                    <div data-drum="{name}" class="card-btn{selected_class}" onclick="
                        document.querySelectorAll('#drum-grid .card-btn').forEach(btn => btn.classList.remove('selected'));
                        this.classList.add('selected');
                        const drumName = this.getAttribute('data-drum');
                        const gradioApp = document.querySelector('gradio-app').shadowRoot || document;
                        const drumInput = gradioApp.querySelector('#drum_state_input');
                        if(drumInput) drumInput.value = drumName;
                        drumInput.dispatchEvent(new Event('input', {{bubbles: true}}));
                    ">
                        <div class="card-title">{name}</div>
                        <div class="card-sub">drum kit</div>
                    </div>
                    '''
                html += '</div>'
                return html
            
            drum_grid_html.value = build_drum_grid("None")
            drum_state_input = gr.Textbox(visible=False, elem_id="drum_state_input")
            drum_state_input.change(fn=lambda x: x, inputs=[drum_state_input], outputs=[drum_state])
            
            # Tabs (tidak berubah)
            tab_select = gr.State(value=0)
            with gr.Tabs():
                with gr.TabItem("custom prompt") as tab1:
                    input_bpm = gr.Slider(label="BPM (0=auto)", minimum=0, maximum=255, step=1, value=0)
                    input_time_sig = gr.Radio(label="time signature", choices=["auto","4/4","2/4","3/4","6/4","7/4","2/2","3/2","4/2","3/8","5/8","6/8","7/8","9/8","12/8"], value="auto")
                    input_key_sig = gr.Dropdown(label="key signature", choices=["auto"] + key_signatures, value="auto")
                with gr.TabItem("midi prompt") as tab2:
                    input_midi = gr.File(label="input midi", file_types=[".midi",".mid"], type="binary")
                    input_midi_events = gr.Slider(label="use first n events", minimum=1, maximum=512, step=1, value=128)
                    input_reduce_cc_st = gr.Checkbox(label="reduce cc & tempo events", value=True)
                    input_remap_track_channel = gr.Checkbox(label="remap tracks/channels", value=True)
                    input_add_default_instr = gr.Checkbox(label="add default instrument", value=True)
                    input_remove_empty_channels = gr.Checkbox(label="remove empty channels", value=False)
                with gr.TabItem("last output prompt") as tab3:
                    gr.Markdown("Continue from last output")
                    input_continuation_select = gr.Radio(label="select output", choices=["all"]+[f"output{i+1}" for i in range(OUTPUT_BATCH_SIZE)], value="all")
                    undo_btn = gr.Button("Undo last continuation", elem_classes="undo-btn")
            
            tab1.select(lambda:0, None, tab_select, queue=False)
            tab2.select(lambda:1, None, tab_select, queue=False)
            tab3.select(lambda:2, None, tab_select, queue=False)

            input_seed = gr.Slider(label="seed", minimum=0, maximum=2**31-1, step=1, value=0)
            input_seed_rand = gr.Checkbox(label="random seed", value=True)
            input_gen_events = gr.Slider(label="generate max events", minimum=1, maximum=opt.max_gen, step=1, value=opt.max_gen//2)
            
            with gr.Accordion("options", open=False):
                input_temp = gr.Slider(label="temperature", minimum=0.1, maximum=1.2, step=0.01, value=1)
                input_top_p = gr.Slider(label="top p", minimum=0.1, maximum=1, step=0.01, value=0.95)
                input_top_k = gr.Slider(label="top k", minimum=1, maximum=128, step=1, value=20)
                input_allow_cc = gr.Checkbox(label="allow midi cc", value=True)
                input_render_audio = gr.Checkbox(label="render audio after generation", value=True)
                
            gr.HTML('<div class="warning-card"><div style="color:#f5a623;font-weight:800;">✨ PETUNJUK GENERASI ✨</div><div style="color:#855d1a;font-size:11px;">Pilih model, instrumen (bisa banyak), drum kit, lalu klik tombol hijau.</div></div>')
            
            run_btn = gr.Button("🫒 GENERATE MUSIC 🫒", elem_classes="gen-btn")
            output_midi_seq = gr.State()
            output_continuation_state = gr.State([0])
            midi_outputs = []
            audio_outputs = []
            
            with gr.Tabs(elem_id="output_tabs"):
                for i in range(OUTPUT_BATCH_SIZE):
                    with gr.TabItem(f"output {i+1}"):
                        gr.HTML(elem_id=f"midi_visualizer_container_{i}")
                        audio_outputs.append(gr.Audio(label="output audio", format="mp3", elem_id=f"midi_audio_{i}"))
                        midi_outputs.append(gr.File(label="output midi", file_types=[".mid"]))
            
            # Diary Aesthetic
            gr.HTML("""
                <div class="diary-container">
                    <div class="diary-text">SYSTEM WAKABA</div>
                    <div class="gif-wrapper">
                        <img src="https://huggingface.co/spaces/Plana-Archive/MOE-TTS/resolve/main/kurumi-tokisaki.gif" alt="aesthetic gif">
                    </div>
                    <div style="font-size: 12px; color:#9bb7a8; margin-top: 8px;">every melody tells a story</div>
                </div>
            """)
            
            # Events
            run_event = run_btn.click(
                run, 
                [model_state, tab_select, output_midi_seq, output_continuation_state, input_continuation_select,
                 instr_state, drum_state, input_bpm, input_time_sig, input_key_sig,
                 input_midi, input_midi_events, input_reduce_cc_st, input_remap_track_channel,
                 input_add_default_instr, input_remove_empty_channels, input_seed, input_seed_rand,
                 input_gen_events, input_temp, input_top_p, input_top_k, input_allow_cc],
                [output_midi_seq, output_continuation_state, input_seed, js_msg],
                queue=True
            )
            
            finish_run_event = run_event.then(
                finish_run, 
                [model_state, output_midi_seq], 
                midi_outputs + [js_msg], 
                queue=False
            )
            finish_run_event.then(
                render_audio, 
                [model_state, output_midi_seq, input_render_audio], 
                audio_outputs, 
                queue=False
            )
            
            undo_btn.click(
                undo_continuation, 
                [model_state, output_midi_seq, output_continuation_state],
                [output_midi_seq, output_continuation_state, js_msg], 
                queue=False
            )
            
            gr.HTML('<div class="credit-footer">🍏 CREATED BY MUTSUMI 🍏</div>')

    app.queue(default_concurrency_limit=2).launch(server_port=opt.port, share=opt.share, ssr_mode=False)
    thread_pool.shutdown()