"""LilyScript — a symbolic-music AIGC app built on the LilyletNotaGen model.

Left column:
  (1) generation parameter panel   (2) streaming run log
  (3) .lyl file list (session outputs + built-in examples)   (4) editable lyl editor
Right column:
  sheet-music panel (placeholder for now; later a lyl->MEI->Verovio SVG renderer
  reusing the lilylet-live-editor pipeline).

Generation streams patch-by-patch: raw decoded text (with `[r:x/y]` stream
markers) goes to the run log, while the measure-segmented postprocessed text
fills the editor. The backend is the int8 + two-level KV-cache ONNX generator
(see lilyscript/generator.py); models load from a local dir for now.
"""

import os
import time

import gradio as gr

from lilyscript.generator import StreamingLilyletGenerator

HERE = os.path.dirname(os.path.abspath(__file__))
# TODO: swap for huggingface_hub.snapshot_download(repo_id=...) to pull the int8
# ONNX weights from the hub instead of a local dir.
MODEL_DIR = os.environ.get('LILYSCRIPT_MODEL_DIR', os.path.join(HERE, 'models'))
ASSET_DIR = os.path.join(HERE, 'assets')
EXAMPLES_DIR = os.path.join(HERE, 'examples')
OUTPUT_DIR = os.path.join(HERE, 'outputs')

EXAMPLE_PREFIX = '\U0001F4C4 '		# 📄 examples
OUTPUT_PREFIX = '✨ '			# ✨ session outputs

_GEN = None


def get_generator ():
	'''Lazily build the (heavy) ONNX generator on first use.'''
	global _GEN
	if _GEN is None:
		_GEN = StreamingLilyletGenerator(MODEL_DIR, ASSET_DIR)
	return _GEN


def load_examples ():
	'''Read built-in example .lyl files into a {label: text} dict.'''
	store = {}
	if os.path.isdir(EXAMPLES_DIR):
		for name in sorted(os.listdir(EXAMPLES_DIR)):
			if name.endswith('.lyl'):
				with open(os.path.join(EXAMPLES_DIR, name), encoding='utf-8') as f:
					store[EXAMPLE_PREFIX + name] = f.read()
	return store


def run_generation (prompt, measures, temperature, top_k, top_p, max_patches, seed, store):
	'''Streaming generate callback. Yields updates for (log, editor, file_list, store).

	store: {label: lyl_text} dict held in gr.State; the produced document is added
	to it under a timestamped label once generation finishes.
	'''
	gen = get_generator()
	meas = int(measures) if measures and int(measures) > 0 else None
	store = dict(store or {})

	raw = pretty = ''
	for raw, pretty, done in gen.generate_stream(
		prompt_text=prompt or '', max_patches=int(max_patches), temperature=float(temperature),
		top_k=int(top_k), top_p=float(top_p), measures=meas, seed=int(seed)):
		# stream: update log + editor, keep file list / store unchanged
		yield raw, pretty, gr.update(), store

	# finished: persist the document, refresh the file list, select the new entry
	label = OUTPUT_PREFIX + time.strftime('%H%M%S') + ('_m%d' % meas if meas else '') + '_s%d' % int(seed)
	store[label] = pretty
	os.makedirs(OUTPUT_DIR, exist_ok=True)
	with open(os.path.join(OUTPUT_DIR, label.replace(OUTPUT_PREFIX, '') + '.lyl'), 'w', encoding='utf-8') as f:
		f.write(pretty)
	yield raw, pretty, gr.update(choices=list(store.keys()), value=label), store


def load_file (label, store):
	'''File-list selection -> load that document into the editor.'''
	return (store or {}).get(label, '')


SHEET_PLACEHOLDER = '''
<div style="height:100%;min-height:600px;display:flex;align-items:center;
	justify-content:center;border:1px dashed #c9c9c9;border-radius:8px;
	color:#999;font-family:sans-serif;text-align:center;">
	<div>
		<div style="font-size:42px;margin-bottom:8px;">&#127932;</div>
		<div>Sheet-music preview</div>
		<div style="font-size:12px;margin-top:4px;">
			(coming soon: Lilylet &rarr; MEI &rarr; Verovio)
		</div>
	</div>
</div>
'''


def build_ui ():
	examples = load_examples()

	with gr.Blocks(title='LilyScript') as demo:
		gr.Markdown('## 🎼 LilyScript — symbolic music generation with Lilylet')
		store = gr.State(examples)

		with gr.Row(equal_height=True):
			# ---------------- LEFT ----------------
			with gr.Column(scale=5):
				# top row: (1) params | (2) run log
				with gr.Row(equal_height=True):
					with gr.Group():
						gr.Markdown('**① Parameters**')
						prompt = gr.Textbox(label='Metadata prompt', lines=3, value='',
							placeholder='[composer "..."]\n[genre "..."]\n(optional)')
						measures = gr.Number(label='Measures (0 = let model decide)', value=8, precision=0)
						with gr.Row():
							temperature = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label='temperature')
							top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.01, label='top_p')
						with gr.Row():
							top_k = gr.Number(label='top_k (0 = off)', value=0, precision=0)
							max_patches = gr.Number(label='max patches', value=256, precision=0)
							seed = gr.Number(label='seed', value=0, precision=0)
						with gr.Row():
							gen_btn = gr.Button('Generate', variant='primary')
							stop_btn = gr.Button('Stop', variant='stop')

					log = gr.Textbox(label='② Run log', lines=16, max_lines=16,
						autoscroll=True, interactive=False)

				# bottom row: (3) file list | (4) editor
				with gr.Row(equal_height=True):
					file_list = gr.Radio(label='③ Files', choices=list(examples.keys()),
						value=None, interactive=True)
					editor = gr.Code(label='④ Lilylet editor', language=None, lines=18,
						interactive=True)

			# ---------------- RIGHT ----------------
			with gr.Column(scale=6):
				gr.Markdown('**Sheet music**')
				gr.HTML(SHEET_PLACEHOLDER)

		# ---- wiring ----
		gen_event = gen_btn.click(
			run_generation,
			inputs=[prompt, measures, temperature, top_k, top_p, max_patches, seed, store],
			outputs=[log, editor, file_list, store],
		)
		stop_btn.click(None, None, None, cancels=[gen_event])
		file_list.select(load_file, inputs=[file_list, store], outputs=[editor])

	return demo


if __name__ == '__main__':
	build_ui().queue().launch(theme=gr.themes.Soft())