"""Built-in code templates shown in the playground editor."""

TEMPLATES = {
    "Custom": '''# Custom template — write any Python code here
# Helpers available to every template:
#   load_text_model(id)      -> (model, tokenizer, device)     cached
#   load_image_model(id)     -> diffusers pipeline             cached
#   load_tts_model(id)       -> (pipe, tokenizer, extra)       cached
#   preview_text(content, clear=False)  -> render text/markdown in the output panel
#   preview(file_or_path)               -> show an image / audio / video / file
#   preview([{role, content}, ...])     -> interactive chat view, returns the user text
#   _cancel_event.is_set()              -> True while the Stop button is pressed
import os

print("Workspace files:")
for f in sorted(os.listdir(".")):
    print("  ", f)''',
    "💬 Chat": '''# Interactive chat — type in the chat box below (type 'quit' to end)
import threading
from transformers import TextIteratorStreamer

MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
model, tokenizer, device = load_text_model(MODEL_ID)

history = [{"role": "system", "content": "You are a helpful assistant."}]
print("Chat started — type in the chat box below, or 'quit' to end.")

while True:
    user_msg = preview(history)
    if user_msg == "__CANCELLED__" or user_msg.lower() in ("quit", "exit", "bye"):
        preview_text("\\n\\n*Chat ended.*")
        break
    if not user_msg.strip():
        continue

    history.append({"role": "user", "content": user_msg})
    preview_text(f"\\n\\n**User:** {user_msg}\\n\\n**Assistant:** ")

    if tokenizer.chat_template:
        prompt = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
    else:
        prompt = "\\n".join(f"{m['role']}: {m['content']}" for m in history[-2:])
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    gen_kwargs = dict(max_new_tokens=512, temperature=0.7, do_sample=True, streamer=streamer)
    thread = threading.Thread(target=model.generate, args=(inputs.input_ids,), kwargs=gen_kwargs)
    thread.start()

    reply = ""
    for text in streamer:
        if text:
            preview_text(text)
            reply += text
    history.append({"role": "assistant", "content": reply})
print("Done")''',
    "✏️ Completion": '''# Text completion — output streams into the main panel
import threading
from transformers import TextIteratorStreamer

MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"
model, tokenizer, device = load_text_model(MODEL_ID)

prompt = "Once upon a time, in a land far away,"
print(f"Model: {MODEL_ID}\\nPrompt: {prompt}\\n")
preview_text(prompt, clear=True)

inputs = tokenizer(prompt, return_tensors="pt").to(device)
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
gen_kwargs = dict(max_new_tokens=300, temperature=0.8, do_sample=True, streamer=streamer)
thread = threading.Thread(target=model.generate, args=(inputs.input_ids,), kwargs=gen_kwargs)
thread.start()

for text in streamer:
    if text:
        preview_text(text)
print("\\nDone")''',
    "🧠 Reasoning": '''# Reasoning with step-by-step thinking — output in the main panel
import threading
from transformers import TextIteratorStreamer

MODEL_ID = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
model, tokenizer, device = load_text_model(MODEL_ID)

questions = [
    "How many r's are in the word 'strawberry'? Think step by step.",
    "Alice has 3 apples. Bob gives her 5 more. She eats 2. How many are left?",
]
preview_text("", clear=True)
for q in questions:
    preview_text(f"\\n**Q:** {q}\\n\\n**A:** ")
    messages = [{"role": "user", "content": q}]
    if tokenizer.chat_template:
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    else:
        prompt = q
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    gen_kwargs = dict(max_new_tokens=512, temperature=0.6, do_sample=True, streamer=streamer)
    thread = threading.Thread(target=model.generate, args=(inputs.input_ids,), kwargs=gen_kwargs)
    thread.start()

    for text in streamer:
        if text:
            preview_text(text)
    preview_text("\\n\\n")
print("Done")''',
    "🌐 Translate": '''# Translation — output in the main panel
import threading
from transformers import TextIteratorStreamer

MODEL_ID = "google/flan-t5-base"
model, tokenizer, device = load_text_model(MODEL_ID)

pairs = [
    ("English", "Spanish", "Hello, how are you today?"),
    ("English", "French", "Where is the nearest train station?"),
    ("Spanish", "English", "El gato está durmiendo en el sofá."),
]
preview_text("", clear=True)
for src, tgt, text in pairs:
    prompt = f"Translate {src} to {tgt}: {text}"
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    gen_kwargs = dict(max_new_tokens=128, do_sample=False, streamer=streamer)
    thread = threading.Thread(target=model.generate, args=(inputs.input_ids,), kwargs=gen_kwargs)
    thread.start()

    preview_text(f"\\n**[{src} → {tgt}]** {text}\\n\\n**Translation:** ")
    for t in streamer:
        if t:
            preview_text(t)
    preview_text("\\n\\n")
print("Done")''',
    "📝 Summarize": '''# Summarization — output in the main panel
import threading
from transformers import TextIteratorStreamer

MODEL_ID = "facebook/bart-large-cnn"
model, tokenizer, device = load_text_model(MODEL_ID)

doc = ("The Industrial Revolution was a period of major industrialization and innovation "
       "during the late 1700s and early 1800s. It began in Great Britain and quickly spread "
       "worldwide. Steam power replaced manual labor, factories sprang up in cities, and "
       "transportation was transformed by railways and steamships. While it brought enormous "
       "economic growth, it also created harsh working conditions and rapid urban expansion. "
       "The innovations of this era laid the foundation for modern industry and changed the "
       "shape of society forever.")
print(f"Input: {len(doc)} chars")
inputs = tokenizer(doc, return_tensors="pt", truncation=True, max_length=512).to(device)
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
gen_kwargs = dict(max_new_tokens=150, min_new_tokens=30, do_sample=False, streamer=streamer)
thread = threading.Thread(target=model.generate, args=(inputs.input_ids,), kwargs=gen_kwargs)
thread.start()

preview_text("**Summary:**\\n", clear=True)
for text in streamer:
    if text:
        preview_text(text)
print("Done")''',
    "🎨 Image Gen": '''# Image generation — live preview of each diffusion step
import os, glob

MODEL_ID = "runwayml/stable-diffusion-v1-5"
pipe = load_image_model(MODEL_ID)

prompt = "A majestic dragon soaring through clouds at sunset"
height, width, steps, seed = 384, 384, 12, 42

for f in glob.glob("_step_*.png"):  # clean up previous step files
    try:
        os.remove(f)
    except OSError:
        pass

def decode_latents(latents):
    sf = pipe.vae.config.scaling_factor
    sh = getattr(pipe.vae.config, "shift_factor", None) or 0
    l = latents.to(pipe.vae.dtype)
    l = (l / sf) + sh
    with torch.no_grad():
        return pipe.image_processor.postprocess(pipe.vae.decode(l, return_dict=False)[0], output_type="pil")[0]

has_vae = hasattr(pipe, "vae") and hasattr(pipe, "image_processor")
step_cnt = {"i": 0}

def cb(pipe_obj, i, t, kw):
    if _cancel_event.is_set():
        return kw
    step_cnt["i"] += 1
    if has_vae and "latents" in kw:
        try:
            img = decode_latents(kw["latents"])
            img.save(f"_step_{step_cnt['i']:03d}.png")
            print(f"Step {step_cnt['i']}/{steps}")
            preview(f"_step_{step_cnt['i']:03d}.png")
        except Exception:
            pass
    return kw

gen = torch.Generator(device="cpu").manual_seed(seed)
kw = dict(prompt=prompt, height=height, width=width, num_inference_steps=steps, generator=gen)
if has_vae:
    kw.update(output_type="latent", callback_on_step_end=cb, callback_on_step_end_tensor_inputs=["latents"])

print(f"Generating {width}x{height} x {steps} steps...")
result = pipe(**kw)
final = decode_latents(result.images) if has_vae else result.images[0]
final.save("output.png")
preview("output.png")
print("Done")''',
    "🔊 TTS": '''# TTS — plays audio in the output panel
import os, scipy.io.wavfile as wav

MODEL_ID = "suno/bark-small"
pipe_or_model, tok, extra = load_tts_model(MODEL_ID)

text = "Hello! Welcome to the Hugging Face Playground. This is a test of text to speech."
print(f"Generating speech for {len(text)} chars...")

if tok is not None:
    speaker = extra if extra is not None else torch.zeros((1, 512))
    inputs = tok(text, return_tensors="pt")
    with torch.no_grad():
        speech = pipe_or_model.generate(input_ids=inputs["input_ids"], speaker_embeddings=speaker)
    audio, sr = speech[0].cpu().numpy(), 16000
else:
    result = pipe_or_model(text)
    sr = result.get("sampling_rate", 22050)
    audio = result["audio"]
    if isinstance(audio, list):
        audio = audio[0]

wav.write("output.wav", sr, audio)
preview("output.wav")
print(f"Done — {len(audio)/sr:.1f}s of audio")''',
    "🎙️ Voice Clone": '''# Voice Cloning — clone any voice, speak any text in it
#   If reference.wav exists in the workspace it is used as the source voice.
#   Otherwise a demo reference voice is auto-generated so this runs end-to-end.
import os, sys, subprocess

REF_AUDIO = "reference.wav"   # your reference voice (upload to workspace to use your own)
TARGET_TEXT = "Hello, this is a cloned voice speaking English and Vietnamese. Xin chào!"

# 1) Obtain a reference voice (user file, or auto-generated demo voice)
if not os.path.exists(REF_AUDIO):
    print("No reference.wav found — auto-generating a demo reference voice...")
    pipe_or_model, tok, extra = load_tts_model("suno/bark-small")
    result = pipe_or_model("Hello, I am the reference voice for this cloning demo. "
                           "Please repeat after me. One, two, three, four, five, six.")
    sr = result.get("sampling_rate", 22050)
    audio = result["audio"]
    if isinstance(audio, list):
        audio = audio[0]
    import scipy.io.wavfile as wav
    wav.write(REF_AUDIO, sr, audio)
    print("Reference voice saved to reference.wav")

# 2) Clone the reference voice with XTTS v2
try:
    from TTS.api import TTS
except ImportError:
    print("Installing TTS package (coqui-tts)... this may take a moment.")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--no-input", "TTS"])
    from TTS.api import TTS

print("Loading XTTS v2 and generating cloned speech (CPU, may take a minute)...")
pipe = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2").to("cpu")
wav_data = pipe.tts(text=TARGET_TEXT, speaker_wav=REF_AUDIO, language="auto")

import scipy.io.wavfile as wav
wav.write("output_voice.wav", 22050, (wav_data * 32767).astype("int16"))
preview("output_voice.wav")
print("Done — voice cloned successfully")''',
    "📂 Read File": '''# Read a file and summarize it with an LLM — output in the main panel
import os, threading
from transformers import TextIteratorStreamer

MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
model, tokenizer, device = load_text_model(MODEL_ID)

path = "app.py"  # change to any file in the workspace
if not os.path.exists(path):
    print(f"File not found: {path}")
else:
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        content = f.read()
    preview_text(f"**File: `{path}`** ({len(content)} chars)\\n\\n", clear=True)
    preview_text(content[:300])
    print(f"Read {len(content)} chars; first 300 shown in preview")

    chunk = content[:1000]
    if tokenizer.chat_template:
        prompt = tokenizer.apply_chat_template(
            [{"role": "user", "content": f"Summarize this Python code:\\n\\n{chunk}"}],
            tokenize=False, add_generation_prompt=True)
    else:
        prompt = f"Summarize this Python code:\\n\\n{chunk}"
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    gen_kwargs = dict(max_new_tokens=128, do_sample=False, streamer=streamer)
    thread = threading.Thread(target=model.generate, args=(inputs.input_ids,), kwargs=gen_kwargs)
    thread.start()

    preview_text("\\n\\n--- **Summary** ---\\n")
    for text in streamer:
        if text:
            preview_text(text)
print("Done")''',
    "#! Shell": '''# Shell + Python — output goes to the Console
#!pwd
#!ls -la *.py 2>/dev/null || echo "no .py files"
#!wc -l app.py 2>/dev/null

import os
files = sorted(os.listdir("."))
print(f"Workspace files ({len(files)}):")
for f in files:
    print("  ", f)''',
}