# Copyright: Shayekh Bin Islam. KAIST, South Korea. 2026.

MAX_TEXT_CHAR = 1500

# model_id = "Qwen/Qwen3.5-9B"
model_id = "Qwen/Qwen3.5-2B"

try:
    import spaces
    IS_HF = True
except ImportError:
    IS_HF = False


if not IS_HF:
    class spaces:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(func):
                return func
            if len(args) == 1 and callable(args[0]) and not kwargs:
                return args[0]
            return decorator
else:
    import os, sys, subprocess
    os.environ['SUPERTONIC_CACHE_DIR']  = '/home/user/huggingface'
    os.environ["HF_HOME"] = "/home/user/huggingface"
    os.environ['XDG_CACHE_HOME'] = "/home/user/huggingface"
    
    os.environ['PLAYWRIGHT_BROWSERS_PATH'] = "/home/user/huggingface/ms-playwright"
    # os.system("playwright install chromium")
    result = subprocess.run(
        ["python", "-m", "playwright", "install", "chromium"],
        env={**os.environ},
        check=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )

import gradio as gr
import fitz  # PyMuPDF
from PIL import Image
import io
import json
import base64
import soundfile as sf
import torch
import os
import tempfile
import re as re_module

from supertonic import TTS
from transformers import AutoProcessor, AutoModelForImageTextToText

# model = None
# processor = None
# tts = None
# voice_style = None

global_stop_thinking = [False]
global_kill_threads = [False]

def set_stop_thinking():
    global_stop_thinking[0] = True
    print(f"[STOP-THINK] set_stop_thinking CALLED! Flag is now: {global_stop_thinking[0]}")
    return gr.update(value="⚡ Forcing generation...")

def set_kill_threads():
    global_kill_threads[0] = True
    print(f"[KILL] set_kill_threads CALLED! Flag is now: {global_kill_threads[0]}")
    return gr.update(value="🛑 Stopping...")

def reset_generation_flags():
    """Reset all generation control flags at the start of a new generation."""
    global_stop_thinking[0] = False
    global_kill_threads[0] = False
    print("[FLAGS] Reset stop_thinking and kill_threads to False")


def extract_pdf_content(pdf_path, max_pages=2):
    """Extract text and images from up to max_pages of a PDF."""
    doc = fitz.open(pdf_path)
    text = ""
    images = []
    for i in range(min(max_pages, len(doc))):
        page = doc[i]
        text += page.get_text() + "\n"
        pix = page.get_pixmap(dpi=150)
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        images.append(img)
    return text, images

def is_youtube_url(url):
    """Check if a URL is a YouTube link."""
    if not url:
        return False
    youtube_patterns = [
        r'(https?://)?(www\.)?youtube\.com/watch',
        r'(https?://)?(www\.)?youtube\.com/shorts/',
        r'(https?://)?youtu\.be/',
        r'(https?://)?(www\.)?youtube\.com/embed/',
        r'(https?://)?m\.youtube\.com/',
    ]
    for pattern in youtube_patterns:
        if re_module.search(pattern, url):
            return True
    return False

def extract_youtube_audio(url, max_duration_sec=300, cookiefile=None):
    """Extract audio from YouTube video (first max_duration_sec seconds).
    Returns path to the downloaded audio file.
    cookiefile: optional path to a Netscape-format cookies.txt to bypass bot filtering.
    """
    import yt_dlp
    from yt_dlp.utils import download_range_func
    
    os.makedirs("log", exist_ok=True)
    output_path = os.path.join("log", "yt_audio")
    
    ydl_opts = {
        'format': 'bestaudio/best',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'wav',
            'preferredquality': '192',
        }],
        'download_ranges': download_range_func(None, [(0, max_duration_sec)]),
        'force_keyframes_at_cuts': True,
        'outtmpl': output_path + '.%(ext)s',
        'quiet': True,
        'no_warnings': True,
    }
    
    if cookiefile:
        ydl_opts['cookiefile'] = cookiefile
    
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        title = info.get('title', 'YouTube Video')
    
    # The output file will be output_path.wav after postprocessing
    wav_path = output_path + '.wav'
    if not os.path.exists(wav_path):
        # Sometimes the extension might differ, search for it
        for ext in ['wav', 'mp3', 'opus', 'webm', 'm4a']:
            candidate = output_path + '.' + ext
            if os.path.exists(candidate):
                wav_path = candidate
                break
    
    if not os.path.exists(wav_path):
        raise FileNotFoundError(f"Could not find downloaded audio file at {output_path}.*")
    
    return wav_path, title

@spaces.GPU(duration=180)
def transcribe_audio_with_asr(audio_path):
    """Transcribe audio file using Cohere ASR model via transformers."""
    global asr_model, asr_processor
    from transformers.audio_utils import load_audio
    
    audio = load_audio(audio_path, sampling_rate=16000)
    
    inputs = asr_processor(audio, language="ko", sampling_rate=16000, return_tensors="pt")
    inputs = inputs.to(asr_model.device, dtype=asr_model.dtype)
    
    outputs = asr_model.generate(**inputs, max_new_tokens=2048)
    texts = asr_processor.decode(outputs, skip_special_tokens=True)
    # text = texts[0] if isinstance(texts, list) else texts
    # join texts
    # Filter the lines in texts which are english only and no korean
    if isinstance(texts, list):
        # Filter out lines that are purely English/symbols (no Korean characters)
        # Korean Unicode range: AC00-D7A3 (Syllables), 1100-11FF (Jamo), 3130-318F (Compatibility Jamo)
        korean_re = re_module.compile(r'[가-힣ㄱ-ㅎㅏ-ㅣ]')
        texts = [line for line in texts if korean_re.search(line)]
        

    text = "\n".join(texts) if isinstance(texts, list) else texts
    
    
    return text.strip()

def extract_website_content(url, max_images=2):
    """Extract text and images from a website URL."""
    import requests
    from bs4 import BeautifulSoup
    import io
    
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
    }
    
    html_content = ""
    try:
        from playwright.sync_api import sync_playwright
        with sync_playwright() as p:
            browser = p.chromium.launch(headless=True)
            page = browser.new_page(user_agent=headers['User-Agent'])
            # Wait until there are no network connections for at least 500 ms (so JS can finish)
            page.goto(url, timeout=30000, wait_until="networkidle")
            html_content = page.content()
            browser.close()
    except Exception as e:
        print(f"Playwright headless fetch failed: {e}. Falling back to requests...")
        response = requests.get(url, headers=headers, timeout=10)
        response.raise_for_status()
        html_content = response.content
        
    soup = BeautifulSoup(html_content, 'html.parser')
    
    for script in soup(["script", "style", "nav", "footer", "header", "noscript"]):
        script.extract()
        
    text = soup.get_text(separator='\n')
    lines = (line.strip() for line in text.splitlines())
    chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
    text = '\n'.join(chunk for chunk in chunks if chunk)
    
    images = []
    img_tags = soup.find_all('img')
    for img in img_tags:
        if len(images) >= max_images:
            break
        src = img.get('src') or img.get('data-src')
        if src:
            if src.startswith('//'):
                src = 'https:' + src
            elif src.startswith('/'):
                from urllib.parse import urljoin
                src = urljoin(url, src)
                
            try:
                img_resp = requests.get(src, headers=headers, timeout=5)
                if img_resp.status_code == 200:
                    pil_img = Image.open(io.BytesIO(img_resp.content))
                    if pil_img.mode != 'RGB':
                        pil_img = pil_img.convert('RGB')
                    if pil_img.width >= 100 and pil_img.height >= 100:
                        images.append(pil_img)
            except Exception as e:
                print(f"Failed to load image {src}: {e}")
                
    return text, images

def get_base64_image(image):
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/jpeg;base64,{img_str}"

@spaces.GPU(duration=180)
def extract_vocabulary(pdf_text, images, translit_lang, translit_format, target_lang, max_text_char=1500, repetition_penalty_val=1.1, partial_assistant_text=None, auto_force_chars=1000):
    """Use Transformers to extract vocabulary from text and images."""
    global model, processor
    
    os.makedirs("log", exist_ok=True)
    
    if len(pdf_text.strip()) == 0:
        pdf_text = '''"No Text available, see provided Images only."'''
    
    no_img = ""
    if len(images) == 0:
        no_img = '''\n"No Images available, see provided Text only."'''

    non_english = ""
    if translit_lang.upper() != "ENGLISH":
        non_english = f" CRITICAL: You MUST use the native alphabet/script of {translit_lang.upper()}, do NOT use English letters unless requested."
    
    prompt_text = f"""Text:

<scrpated-content>
{pdf_text[:int(max_text_char)]}
</scrpated-content>{no_img}

Extract at least 10 key Korean words or phrases from the following text and images.
Focus on meaningful vocabulary that is highly helpful for a new language learner (e.g., common nouns, verbs, adjectives, or useful expressions).
CRITICAL: Do NOT extract website template words, navigation menus, boilerplate text, UI elements, or titles like 'Home page', 'News', 'Menu'.

Return ONLY a valid JSON list of dictionaries, where each dictionary has four keys:
- 'korean' (the Korean text)
- 'transliteration' (the pronunciation transliterated into {translit_lang.upper()} script/characters, formatted as {translit_format}.{non_english})
- 'translation' (the brief translation into {target_lang.upper()})
- 'explanation' (a brief grammar or context note in {target_lang.upper()}).

Just output raw JSON with ```json and ``` markers, as the user will load in python. Example:

```json
[
  {{
    "korean": "날씨",
    "transliteration": "nal-ssi",
    "translation": "weather",
    "explanation": "Common noun used to describe weather conditions."
  }},
  {{
    "korean": "맛있다",
    "transliteration": "ma-sit-da",
    "translation": "to be delicious",
    "explanation": "Descriptive verb. Polite form: 맛있어요. Used to compliment food."
  }}
]
```

CRITICAL: Do NOT overthink. Do NOT deliberate over conditions, edge cases, or reasoning. Keep your thinking extremely brief (a few words at most). Output the JSON array IMMEDIATELY without lengthy analysis.


"""
    
    # DEBUG: Log prompt text
    with open("log/debug_vlm_prompt.txt", "w", encoding="utf-8") as f:
        f.write(prompt_text)

    content = []
    pil_images = []
    
    for i, img in enumerate(images):
        # DEBUG: Log images
        img.save(f"log/debug_image_{i}.png", format="PNG")
        pil_images.append(img)
        
        content.append({
            "type": "image",
        })
        
    content += [{"type": "text", "text": prompt_text}]

    messages = [
        {
            "role": "user",
            "content": content
        }
    ]

    try:
        model.to("cuda")
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        if partial_assistant_text:
            text += partial_assistant_text + "\nReady to generate.\n</think>\n\n```json\n[\n"
            
        inputs = processor(
            text=[text],
            images=pil_images if pil_images else None,
            return_tensors="pt",
            padding=True
        ).to("cuda")
        
        from transformers import TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
        from threading import Thread
        import queue

        local_stop = [False]
        
        class LocalKillCriteria(StoppingCriteria):
            def __call__(self, input_ids, scores, **kwargs):
                return local_stop[0] or global_kill_threads[0]

        def run_generation(cur_inputs, cur_streamer, cur_local_stop):
            """Run model.generate in a thread, always calling streamer.end() on exit."""
            kill_criteria = StoppingCriteriaList([LocalKillCriteria()])
            gen_kwargs = dict(
                **cur_inputs,
                streamer=cur_streamer,
                max_new_tokens=2048*16,
                do_sample=True,
                repetition_penalty=repetition_penalty_val,
                stopping_criteria=kill_criteria
            )
            if len(images) > 0:
                gen_kwargs.update(dict(temperature=0.6, top_p=0.95, top_k=20, min_p=0.0))
            else:
                gen_kwargs.update(dict(temperature=1.0, top_p=0.95, top_k=20, min_p=0.0))
            try:
                model.generate(**gen_kwargs)
            except Exception as e:
                import traceback
                print(f"\n[THREAD ERROR] model.generate crashed: {e}")
                traceback.print_exc()
            finally:
                try:
                    cur_streamer.end()
                except Exception:
                    pass

        output_text = partial_assistant_text + "\n</think>\n\n```json\n[\n" if partial_assistant_text else ""

        streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
        thread = Thread(target=run_generation, args=(inputs, streamer, local_stop))
        thread.start()
        
        force_triggered = False
        AUTO_FORCE_CHARS = auto_force_chars
        for new_text in streamer:
            output_text += new_text
            yield output_text, None
            
            # Auto-force JSON if thinking exceeds 300 chars without producing JSON
            should_auto_force = (
                not force_triggered
                and not partial_assistant_text
                and len(output_text) > AUTO_FORCE_CHARS
                and '```json' not in output_text
            )
            
            # Check if user clicked "Stop thinking" OR auto-force threshold reached
            if (global_stop_thinking[0] or should_auto_force) and not force_triggered:
                force_triggered = True
                reason = f"auto-force (>{AUTO_FORCE_CHARS} chars)" if should_auto_force else "user clicked stop"
                print(f"[STOP-THINK] Force triggered ({reason})! Killing current generation...")
                
                # 1. Kill the current generation thread
                local_stop[0] = True
                # Drain queue so the thread can exit
                while not streamer.text_queue.empty():
                    try:
                        streamer.text_queue.get_nowait()
                    except queue.Empty:
                        break
                thread.join(timeout=5)
                print("[STOP-THINK] Old thread joined. Starting forced JSON generation...")
                
                # 2. Reset flags
                global_stop_thinking[0] = False
                local_stop[0] = False
                
                # 3. Append the think-closing + JSON prefix
                output_text += "\nReady to generate.\n</think>\n\n```json\n[\n"
                yield output_text, None
                
                # 4. Build new prompt with partial assistant text
                text2 = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
                text2 += output_text
                inputs2 = processor(
                    text=[text2],
                    images=pil_images if pil_images else None,
                    return_tensors="pt",
                    padding=True
                ).to("cuda")
                
                # 5. Start new generation thread with force-JSON context
                # This loop also monitors stop_thinking so user can force again if model keeps thinking
                streamer2 = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
                thread2 = Thread(target=run_generation, args=(inputs2, streamer2, local_stop))
                thread2.start()
                
                for new_text2 in streamer2:
                    output_text += new_text2
                    yield output_text, None
                    
                    # Allow user to force again if model still isn't producing JSON
                    if global_stop_thinking[0] or global_kill_threads[0]:
                        print("[STOP-THINK] Flag detected in forced generation loop! Killing...")
                        local_stop[0] = True
                        while not streamer2.text_queue.empty():
                            try:
                                streamer2.text_queue.get_nowait()
                            except queue.Empty:
                                break
                        thread2.join(timeout=5)
                        global_stop_thinking[0] = False
                        local_stop[0] = False
                        
                        # Force JSON prefix again
                        output_text += "\nReady to generate.\n</think>\n\n```json\n[\n"
                        yield output_text, None
                        
                        text3 = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
                        text3 += output_text
                        inputs3 = processor(
                            text=[text3],
                            images=pil_images if pil_images else None,
                            return_tensors="pt",
                            padding=True
                        ).to("cuda")
                        
                        streamer3 = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
                        thread3 = Thread(target=run_generation, args=(inputs3, streamer3, local_stop))
                        thread3.start()
                        
                        for new_text3 in streamer3:
                            output_text += new_text3
                            yield output_text, None
                        
                        thread3.join(timeout=10)
                        break
                else:
                    thread2.join(timeout=10)
                break  # Exit the outer streamer loop
        
        if not force_triggered:
            thread.join()
        
        # Reset flag in case it was set but generation finished naturally
        global_stop_thinking[0] = False

        # DEBUG: Log raw output text
        with open("log/debug_vlm_output.txt", "w", encoding="utf-8") as f:
            f.write(output_text)
            
    except Exception as e:
        print(f"Error during Transformers inference: {e}")
        yield f"Error during Transformers inference: {e}", []
        return

    try:
        import re
        # Extract JSON from markdown code fences or raw output
        json_matches = list(re.finditer(r'```(?:json)?\s*([\s\S]*?)```', output_text))
        if json_matches:
            clean_text = json_matches[-1].group(1).strip()
        else:
            # Fallback: find last [ ... ] or { ... } block
            json_matches = list(re.finditer(r'(\[[\s\S]*\]|\{[\s\S]*\})', output_text))
            clean_text = json_matches[-1].group(1).strip() if json_matches else output_text.strip()
        try:
            data = json.loads(clean_text)
        except:
            import jiter
            # Get bytes from string
            data = jiter.from_json(clean_text.encode("utf-8"), partial_mode=True)
        
        if not isinstance(data, list):
            data = [data]
        yield output_text, data
    except Exception as e:
        print(f"Error parsing JSON: {e}\nRaw output: {output_text}")
        yield output_text, []

def translate_vocabulary(korean_words, translit_lang, translit_format, target_lang, repetition_penalty_val=1.1):
    """Use Transformers text-only inference to translate/transliterate Korean words."""
    global model, processor
    
    non_english = ""
    if translit_lang.upper() != "ENGLISH":
        non_english = f" CRITICAL: You MUST use the native alphabet/script of {translit_lang.upper()}, do NOT use English letters unless requested."
    
    words_str = ", ".join(korean_words)
    prompt_text = f"""Translate and transliterate the following Korean words.
Return ONLY a valid JSON list of dictionaries, where each dictionary has four keys:
- 'korean' (the original Korean text)
- 'transliteration' (the pronunciation transliterated into {translit_lang.upper()} script/characters, formatted as {translit_format}.{non_english})
- 'translation' (the translation into {target_lang.upper()})
- 'explanation' (a brief grammar or context note in {target_lang.upper()}).
No markdown formatting, just raw JSON with ```json and ``` markers.
CRITICAL: Do NOT overthink. Do NOT deliberate over conditions, edge cases, or reasoning. Keep your thinking extremely brief (5 paragraphs at most). Output the JSON array IMMEDIATELY without lengthy analysis.

Korean words:
{words_str}
"""

    # DEBUG: Log translation prompt text
    with open("log/debug_translate_prompt.txt", "w", encoding="utf-8") as f:
        f.write(prompt_text)

    messages = [
        {
            "role": "user",
            "content": [{"type": "text", "text": prompt_text}]
        }
    ]

    try:
        model.to("cuda")
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(
            text=[text],
            images=None,
            return_tensors="pt",
            padding=True
        ).to("cuda")

        generated_ids = model.generate(
            **inputs,
            # max_new_tokens=2048*16,
            max_new_tokens=2048*2,
            # temperature=1.0,
            # top_p=0.95,
            temperature=1.0, top_p=0.95, top_k=20, min_p=0.0, 
            # presence_penalty=1.5, 
            repetition_penalty=repetition_penalty_val,
            do_sample=True
        )
        
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
        ]
        output_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        
        # DEBUG: Log raw translation output text
        with open("log/debug_translate_output.txt", "w", encoding="utf-8") as f:
            f.write(output_text)
            
    except Exception as e:
        print(f"Error during Transformers text inference: {e}")
        return []

    try:
        import re
        json_matches = list(re.finditer(r'```(?:json)?\s*([\s\S]*?)```', output_text))
        if json_matches:
            clean_text = json_matches[-1].group(1).strip()
        else:
            json_matches = list(re.finditer(r'(\[[\s\S]*\]|\{[\s\S]*\})', output_text))
            clean_text = json_matches[-1].group(1).strip() if json_matches else output_text.strip()
            
        try:
            data = json.loads(clean_text)
        except:
            import jiter
            data = jiter.from_json(clean_text.encode("utf-8"), partial_mode=True)
            
        if not isinstance(data, list):
            data = [data]
        return data
    except Exception as e:
        print(f"Error parsing JSON: {e}\nRaw output: {output_text}")
        return []

def numpy_to_base64_audio(wav, sample_rate):
    wav = wav.squeeze()
    buffer = io.BytesIO()
    sf.write(buffer, wav, sample_rate, format='WAV')
    buffer.seek(0)
    audio_base64 = base64.b64encode(buffer.read()).decode('utf-8')
    return f"data:audio/wav;base64,{audio_base64}"

import hashlib

def hash_file(filepath):
    with open(filepath, 'rb') as f:
        return hashlib.md5(f.read(1024*1024)).hexdigest()

@spaces.GPU(duration=180)
def process_pdf(pdf_file, url_input, audio_file_input, yt_url_input, yt_cookies_file, translit_lang, translit_format, target_lang, max_text_char, repetition_penalty_val, auto_force_chars_val, last_source_hash, last_korean_words, active_tab, progress=gr.Progress()):
    global tts, voice_style
    
    # Clean language choices from "Family - Language" to just "Language"
    if " - " in translit_lang:
        translit_lang = translit_lang.split(" - ")[-1]
    if " - " in target_lang:
        target_lang = target_lang.split(" - ")[-1]
        
    os.makedirs("log", exist_ok=True)
    
    # Reset flags at start of new generation
    reset_generation_flags()
    
    # Determine input source based on active tab
    is_url = (active_tab == "Website URL") and bool(url_input and url_input.strip())
    is_youtube = (active_tab == "YouTube Link") and bool(yt_url_input and yt_url_input.strip() and is_youtube_url(yt_url_input.strip()))
    is_audio_upload = (active_tab == "Upload Audio") and (audio_file_input is not None)
    is_pdf = (active_tab == "Upload PDF") and (pdf_file is not None)
    
    if not is_url and not is_youtube and not is_audio_upload and not is_pdf:
        yield "<p>Please provide input in the active tab.</p>", None, None, "", "", [], None
        return
        
    if is_youtube:
        current_source_hash = hashlib.md5(yt_url_input.strip().encode()).hexdigest()
    elif is_audio_upload:
        current_source_hash = hash_file(audio_file_input)
    elif is_url:
        current_source_hash = hashlib.md5(url_input.strip().encode()).hexdigest()
    else:
        current_source_hash = hash_file(pdf_file.name)
        
    vocab_list = []
    extracted_audio_path = None

    try:
        if is_youtube:
            progress(0, desc="Downloading YouTube audio (first 5 min)...")
            cookies_path = yt_cookies_file.name if yt_cookies_file else None
            audio_path, yt_title = extract_youtube_audio(yt_url_input.strip(), max_duration_sec=300, cookiefile=cookies_path)
            
            progress(0.1, desc=f"Transcribing audio with Cohere ASR...")
            content_text = transcribe_audio_with_asr(audio_path)
            images = []
            extracted_audio_path = audio_path
                
            if not content_text.strip():
                yield "<p>Could not transcribe any text from the YouTube video.</p>", current_source_hash, None, "", "", [], extracted_audio_path
                return
                
            # Log the transcription
            with open("log/debug_yt_transcription.txt", "w", encoding="utf-8") as f:
                f.write(f"Title: {yt_title}\n\n{content_text}")
        elif is_audio_upload:
            progress(0, desc="Transcribing uploaded audio with Cohere ASR...")
            content_text = transcribe_audio_with_asr(audio_file_input)
            images = []
            extracted_audio_path = audio_file_input
                
            if not content_text.strip():
                yield "<p>Could not transcribe any text from the uploaded audio.</p>", current_source_hash, None, "", "", [], extracted_audio_path
                return
                
            with open("log/debug_audio_transcription.txt", "w", encoding="utf-8") as f:
                f.write(content_text)
        elif is_url:
            progress(0, desc="Fetching Website...")
            content_text, images = extract_website_content(url_input.strip())
        else:
            progress(0, desc="Reading PDF...")
            content_text, images = extract_pdf_content(pdf_file.name)
            
        if not content_text.strip() and not images:
            yield "<p>No content found.</p>", current_source_hash, None, "", "", [], extracted_audio_path
            return
    except Exception as e:
        import traceback
        traceback.print_exc()
        yield f"<p>Error reading content: {e}</p>", None, None, "", "", [], None
        return

    vocab_list = []
    stream_text = ""
    for attempt in range(1, 4):
        if global_kill_threads[0]:
            print("[KILL] Kill flag detected, stopping extraction attempts.")
            break
        progress(0.2, desc=f"Extracting vocabulary (Attempt {attempt}/3)...")
        for stream_t, v_list in extract_vocabulary(content_text, images, translit_lang, translit_format, target_lang, max_text_char, repetition_penalty_val, auto_force_chars=auto_force_chars_val):
            stream_text = stream_t
            if v_list is not None:
                vocab_list = v_list
            yield "", current_source_hash, None, stream_text, content_text, images, extracted_audio_path
            
        if vocab_list:
            break

    # Reset kill flag after extraction so TTS can proceed
    global_kill_threads[0] = False

    # If generation was killed but we don't have vocab yet, try to salvage JSON from stream_text
    if not vocab_list and stream_text:
        print("[KILL] Attempting to salvage JSON from partial generation output...")
        try:
            import re
            json_matches = list(re.finditer(r'```(?:json)?\s*([\s\S]*?)```', stream_text))
            if json_matches:
                clean_text = json_matches[-1].group(1).strip()
            else:
                json_matches = list(re.finditer(r'(\[[\s\S]*\]|\{[\s\S]*\})', stream_text))
                clean_text = json_matches[-1].group(1).strip() if json_matches else ""
            
            if clean_text:
                try:
                    data = json.loads(clean_text)
                except:
                    import jiter
                    data = jiter.from_json(clean_text.encode("utf-8"), partial_mode=True)
                if not isinstance(data, list):
                    data = [data]
                if data and isinstance(data[0], dict) and 'korean' in data[0]:
                    vocab_list = data
                    print(f"[KILL] Salvaged {len(vocab_list)} vocab items from partial output!")
        except Exception as e:
            print(f"[KILL] Could not salvage JSON: {e}")

    if not vocab_list:
        yield "<p>Failed to extract or translate vocabulary after 3 attempts.</p>", current_source_hash, None, stream_text, content_text, images, extracted_audio_path
        return

    progress(0.6, desc="Generating TTS audio...")
    # Pre-generate TTS audio
    for i, item in enumerate(vocab_list):
        korean = item.get("korean", "")
        # Add dot
        if not korean.endswith("."):
            korean += "."
            
        try:
            wav, dur = tts.synthesize(
                korean, voice_style=voice_style, lang="ko",
                total_steps=12,
                speed=0.7,
            )
            
            # DEBUG: Save audio locally
            wav_1d = wav.squeeze()
            sf.write(f"log/debug_audio_{i}.wav", wav_1d, tts.sample_rate, format='WAV')
            
            audio_data_uri = numpy_to_base64_audio(wav, tts.sample_rate)
            item['audio_uri'] = audio_data_uri
        except Exception as e:
            print(f"TTS error for '{korean}': {e}")
            item['audio_uri'] = None

    cards_json = json.dumps(vocab_list).replace("</", "<\\/")

    iframe_html = f"""
    <!DOCTYPE html>
    <html>
    <head>
    <!-- Flaticon UIcons CDN -->
    <link rel='stylesheet' href='https://cdn-uicons.flaticon.com/uicons-regular-rounded/css/uicons-regular-rounded.css'>
    <style>
        body {{
            margin: 0;
            padding: 0;
            background: transparent;
        }}
        .flashcard-container {{
            perspective: 1000px;
            width: 100%;
            max-width: 500px;
            margin: 0 auto;
            font-family: 'Inter', sans-serif;
            padding-top: 20px;
        }}
        .flashcard {{
            width: 100%;
            min-height: 400px;
            display: grid;
            transition: transform 0.6s cubic-bezier(0.4, 0.2, 0.2, 1);
            transform-style: preserve-3d;
            cursor: pointer;
        }}
        .flashcard.is-flipped {{
            transform: rotateY(180deg);
        }}
        .card-face {{
            grid-area: 1 / 1;
            width: 100%;
            backface-visibility: hidden;
            display: flex;
            flex-direction: column;
            justify-content: center;
            align-items: center;
            border-radius: 20px;
            box-shadow: 0 10px 30px rgba(0,0,0,0.1);
            padding: 30px;
            box-sizing: border-box;
            background: rgba(15, 23, 42, 0.6);
            backdrop-filter: blur(15px);
            border: 1px solid rgba(255, 255, 255, 0.1);
            text-align: center;
            box-shadow: inset 0 0 0 1px rgba(255, 255, 255, 0.05), 0 10px 40px rgba(0, 0, 0, 0.5);
        }}
        .card-front {{
            background: linear-gradient(135deg, rgba(30, 41, 59, 0.9) 0%, rgba(15, 23, 42, 0.9) 100%);
            border-top: 2px solid rgba(139, 92, 246, 0.5);
        }}
        .card-back {{
            transform: rotateY(180deg);
            background: linear-gradient(135deg, rgba(30, 58, 138, 0.9) 0%, rgba(15, 23, 42, 0.9) 100%);
            border-top: 2px solid rgba(56, 189, 248, 0.5);
            color: #f8fafc;
        }}
        .korean-text {{
            font-size: 54px;
            font-weight: 700;
            color: #ffffff;
            text-shadow: 0 0 20px rgba(139, 92, 246, 0.6);
            margin-bottom: 20px;
        }}
        .english-text {{
            font-size: 34px;
            font-weight: 700;
            color: #ffffff;
            text-shadow: 0 0 15px rgba(56, 189, 248, 0.6);
            margin-bottom: 5px;
        }}
        .translit-text {{
            font-size: 20px;
            font-style: italic;
            color: #fca5a5;
            margin-bottom: 15px;
            letter-spacing: 1px;
        }}
        .explanation-text {{
            font-size: 16px;
            color: #cbd5e1;
            line-height: 1.6;
            background: rgba(0,0,0,0.2);
            padding: 15px;
            border-radius: 12px;
            border: 1px solid rgba(255,255,255,0.05);
        }}
        .nav-buttons {{
            display: flex;
            justify-content: space-between;
            margin-top: 30px;
            width: 100%;
            max-width: 500px;
            margin-left: auto;
            margin-right: auto;
        }}
        .nav-btn {{
            padding: 14px 28px;
            border: none;
            border-radius: 16px;
            background: rgba(139, 92, 246, 0.2);
            backdrop-filter: blur(5px);
            border: 1px solid rgba(139, 92, 246, 0.5);
            color: white;
            font-weight: 600;
            cursor: pointer;
            text-transform: uppercase;
            letter-spacing: 1px;
            transition: all 0.3s ease;
            box-shadow: 0 4px 15px rgba(0,0,0,0.2);
        }}
        .nav-btn:hover {{
            background: rgba(139, 92, 246, 0.5);
            transform: translateY(-2px);
            box-shadow: 0 6px 20px rgba(139, 92, 246, 0.4);
        }}
            transition: all 0.2s;
            box-shadow: 0 4px 12px rgba(124, 58, 237, 0.3);
            flex: 1;
            margin: 0 10px;
            display: flex;
            align-items: center;
            justify-content: center;
            gap: 8px;
        }}
        .nav-btn:hover {{
            background: #6d28d9;
            transform: translateY(-2px);
        }}
        .nav-btn:disabled {{
            background: #ccc;
            cursor: not-allowed;
            transform: none;
            box-shadow: none;
        }}
        .audio-btn {{
            margin-top: 15px;
            padding: 12px 24px;
            background: rgba(56, 189, 248, 0.2);
            color: #38bdf8;
            border: 1px solid rgba(56, 189, 248, 0.4);
            border-radius: 20px;
            font-size: 15px;
            font-weight: 600;
            cursor: pointer;
            transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1);
            display: flex;
            align-items: center;
            justify-content: center;
            gap: 8px;
            text-transform: uppercase;
            letter-spacing: 1px;
        }}
        .audio-btn:hover {{
            background: rgba(56, 189, 248, 0.4);
            transform: scale(1.05) translateY(-2px);
            box-shadow: 0 5px 15px rgba(56, 189, 248, 0.3);
        }}
        .progress {{
            text-align: center;
            margin-top: 15px;
            color: #666;
            font-size: 14px;
            font-weight: 600;
        }}
        .copy-icon {{
            cursor: pointer;
            color: #a0a0a0;
            font-size: 24px;
            transition: color 0.2s, transform 0.2s;
        }}
        .copy-icon:hover {{
            color: #7c3aed;
            transform: scale(1.1);
        }}
    </style>
    </head>
    <body>
    <div id="flashcard-app">
        <div class="flashcard-container">
            <div class="flashcard" id="card" onclick="flipCard()">
                <div class="card-face card-front">
                    <div style="display: flex; align-items: center; justify-content: center; gap: 15px; margin-bottom: 20px; width: 100%;">
                        <div class="korean-text" id="front-text" style="margin-bottom: 0;"><i class="fi fi-rr-spinner-third fa-spin"></i> Loading...</div>
                        <i class="fi fi-rr-copy copy-icon" id="copy-ko" onclick="copyText('front-text', event, this)" title="Copy text" style="display:none;"></i>
                    </div>
                    <button class="audio-btn" onclick="playAudio(event)" id="audio-btn" style="display:none;"><i class="fi fi-rr-play-circle"></i> Play Audio</button>
                    <p style="margin-top:20px; color:#999; font-size:13px; display:flex; align-items:center; gap:5px;"><i class="fi fi-rr-rotate-right"></i> Click card to flip 🎯</p>
                </div>
                <div class="card-face card-back">
                    <div style="display: flex; align-items: center; justify-content: center; gap: 15px; margin-bottom: 5px; width: 100%;">
                        <div class="english-text" id="back-en" style="margin-bottom: 0;"></div>
                        <i class="fi fi-rr-copy copy-icon" id="copy-en" onclick="copyText('back-en', event, this)" title="Copy text" style="display:none;"></i>
                    </div>
                    <div class="translit-text" id="back-translit"></div>
                    <div class="explanation-text"><i class="fi fi-rr-lightbulb-on" style="color:#f1c40f;"></i> <span id="back-exp"></span></div>
                </div>
            </div>
        </div>
        <div class="nav-buttons">
            <button class="nav-btn" id="prev-btn" onclick="prevCard()"><i class="fi fi-rr-angle-left"></i> Previous</button>
            <button class="nav-btn" id="next-btn" onclick="nextCard()">Next <i class="fi fi-rr-angle-right"></i></button>
        </div>
        <div class="progress" id="progress-text"></div>
    </div>

    <script>
        const cards = {cards_json};
        let currentIndex = 0;
        let audioPlayer = new Audio();

        function updateCard() {{
            if (!cards || cards.length === 0) {{
                document.getElementById('front-text').innerHTML = "No vocabulary found 😥";
                document.getElementById('prev-btn').disabled = true;
                document.getElementById('next-btn').disabled = true;
                if(document.getElementById('copy-ko')) document.getElementById('copy-ko').style.display = 'none';
                if(document.getElementById('copy-en')) document.getElementById('copy-en').style.display = 'none';
                return;
            }}
            const card = cards[currentIndex];
            document.getElementById('front-text').innerText = card.korean || "No word";
            document.getElementById('back-en').innerText = card.translation || card.english || "";
            document.getElementById('back-translit').innerText = card.transliteration ? `[${{card.transliteration}}]` : "";
            document.getElementById('back-exp').innerText = card.explanation || "";
            
            if(document.getElementById('copy-ko')) document.getElementById('copy-ko').style.display = 'block';
            if(document.getElementById('copy-en')) document.getElementById('copy-en').style.display = 'block';
            
            document.getElementById('prev-btn').disabled = currentIndex === 0;
            document.getElementById('next-btn').disabled = currentIndex === cards.length - 1;
            document.getElementById('progress-text').innerHTML = `📚 Card ${{currentIndex + 1}} of ${{cards.length}}`;
            
            const cardEl = document.getElementById('card');
            cardEl.classList.remove('is-flipped');
            
            if(card.audio_uri) {{
                audioPlayer.src = card.audio_uri;
                document.getElementById('audio-btn').style.display = 'flex';
            }} else {{
                document.getElementById('audio-btn').style.display = 'none';
            }}
        }}

        function flipCard() {{
            if (!cards || cards.length === 0) return;
            document.getElementById('card').classList.toggle('is-flipped');
        }}

        function playAudio(e) {{
            e.stopPropagation(); 
            audioPlayer.play().catch(err => console.log("Audio play error:", err));
        }}

        function nextCard() {{
            if (currentIndex < cards.length - 1) {{
                currentIndex++;
                updateCard();
            }}
        }}

        function prevCard() {{
            if (currentIndex > 0) {{
                currentIndex--;
                updateCard();
            }}
        }}

        function copyText(elementId, e, iconEl) {{
            e.stopPropagation();
            const textToCopy = document.getElementById(elementId).innerText;
            navigator.clipboard.writeText(textToCopy).then(() => {{
                const oldClass = iconEl.className;
                iconEl.className = "fi fi-rr-check copy-icon";
                iconEl.style.color = "#27ae60";
                setTimeout(() => {{
                    iconEl.className = oldClass;
                    iconEl.style.color = "";
                }}, 1500);
            }}).catch(err => console.log('Copy failed', err));
        }}

        window.onload = function() {{
            updateCard();
        }};
    </script>
    </body>
    </html>
    """
    
    import html
    safe_srcdoc = html.escape(iframe_html)
    
    # Return the iframe containing the whole SPA
    yield f'<iframe srcdoc="{safe_srcdoc}" style="width: 100%; height: 650px; border: none; overflow-y: auto;"></iframe>', current_source_hash, vocab_list, stream_text, content_text, images, extracted_audio_path

LANGUAGE_DATA = """Indo-European	Bengali, English, French, Portuguese, German, Romanian, Swedish, Danish, Bulgarian, Russian, Czech, Greek, Ukrainian, Spanish, Dutch, Slovak, Croatian, Polish, Lithuanian, Norwegian Bokmål, Norwegian Nynorsk, Persian, Slovenian, Gujarati, Latvian, Italian, Occitan, Nepali, Marathi, Belarusian, Serbian, Luxembourgish, Venetian, Assamese, Welsh, Silesian, Asturian, Chhattisgarhi, Awadhi, Maithili, Bhojpuri, Sindhi, Irish, Faroese, Hindi, Punjabi, Oriya, Tajik, Eastern Yiddish, Lombard, Ligurian, Sicilian, Friulian, Sardinian, Galician, Catalan, Icelandic, Tosk Albanian, Limburgish, Dari, Afrikaans, Macedonian, Sinhala, Urdu, Magahi, Bosnian, Armenian, Latgalian, Scottish Gaelic, Central Kurdish, Northern Kurdish, Southern Pashto, Sanskrit, Dhundari, Marwari, Ahirani, Bagheli, Bagri, Bundeli, Braj, Kumaoni, Kashmiri
Sino-Tibetan	Chinese (Simplified), Chinese (Traditional), Cantonese, Burmese, Standard Tibetan, Meitei
Afro-Asiatic	Arabic (Standard), Arabic (Najdi), Arabic (Levantine), Arabic (Egyptian), Arabic (Moroccan), Arabic (Mesopotamian), Arabic (Ta’izzi-Adeni), Arabic (Tunisian), Arabic (Gulf), Arabic (Algerian), Arabic (Sudanese), Arabic (Libyan), Hebrew, Maltese, Amharic, Tigrinya, Kabyle, Somali, West Central Oromo, Hausa
Austronesian	Indonesian, Malay, Tagalog, Cebuano, Javanese, Sundanese, Minangkabau, Balinese, Banjar, Pangasinan, Iloko, Waray (Philippines), Plateau Malagasy, Malagasy, Buginese, Maori, Samoan, Hawaiian, Fijian
Dravidian	Tamil, Telugu, Kannada, Malayalam
Turkic	Turkish, North Azerbaijani, Northern Uzbek, Kazakh, Bashkir, Tatar, Crimean Tatar, Kyrgyz, Turkmen, Uyghur
Tai-Kadai	Thai, Lao, Shan
Uralic	Finnish, Estonian, Hungarian, Meadow Mari
Austroasiatic	Vietnamese, Khmer
Niger–Congo	Yoruba, Ewe, Kinyarwanda, Lingala, Northern Sotho, Nyanja, Shona, Southern Sotho, Tswana, Xhosa, Zulu, Luganda, Swati, Tsonga, Tumbuka, Venda, Chokwe, Luba-Kasai, Rundi, Umbundu, Kikuyu, Kongo, Nigerian Fulfulde, Wolof, Fon, Kabiyè, Mossi, Akan, Twi, Bambara, Igbo"""
# Other	Japanese, Korean, Georgian, Basque, Haitian, Papiamento, Kabuverdianu, Tok Pisin, Swahili, Central Aymara, Tulu, Nagamese, Nigerian Pidgin, Mauritian Creole, Sango, Ayacucho Quechua, Halh Mongolian, Southwestern Dinka, Nuer, Guarani

LANGUAGE_CHOICES = []
for line in LANGUAGE_DATA.strip().split('\n'):
    family, langs = line.split('\t')
    for lang in langs.split(', '):
        LANGUAGE_CHOICES.append(f"{family} - {lang}")

import urllib.request

def get_example_pdf():
    url = "https://raw.githubusercontent.com/ShayekhBinIslam/file-host/main/cnp_korean_page7.pdf"
    file_path = "cnp_korean_page7.pdf"
    if not os.path.exists(file_path):
        try:
            urllib.request.urlretrieve(url, file_path)
        except Exception as e:
            print(f"Failed to download example PDF: {e}")
    return file_path if os.path.exists(file_path) else None

def get_example_audio():
    url = "https://raw.githubusercontent.com/ShayekhBinIslam/file-host/main/new_1min.wav"
    file_path = "new_1min.wav"
    if not os.path.exists(file_path):
        try:
            urllib.request.urlretrieve(url, file_path)
        except Exception as e:
            print(f"Failed to download example audio: {e}")
    return file_path if os.path.exists(file_path) else None

@spaces.GPU(duration=180)
def process_pdf_force(partial_text, pdf_file, url_input, translit_lang, translit_format, target_lang, max_text_char, repetition_penalty_val, last_source_state, last_korean_words_state):
    """Force JSON generation using the current partial stream_box text."""
    is_url = bool(url_input and url_input.strip())
    
    current_source_hash = ""
    if is_url:
        current_source_hash = "url:" + url_input.strip()
    elif pdf_file is not None:
        import hashlib
        with open(pdf_file.name, "rb") as f:
            current_source_hash = "pdf:" + hashlib.md5(f.read()).hexdigest()
            
    try:
        if is_url:
            progress(0, desc="Fetching Website...")
            content_text, images = extract_website_content(url_input.strip())
        else:
            progress(0, desc="Reading PDF...")
            content_text, images = extract_pdf_content(pdf_file.name)
            
        if not content_text.strip() and not images:
            yield "<p>No content found.</p>", current_source_hash, None, partial_text, "", [], None
            return
    except Exception as e:
        yield f"<p>Error reading content: {e}</p>", None, None, partial_text, "", [], None
        return

    vocab_list = []
    stream_text = partial_text
    
    progress(0.2, desc="Extracting vocabulary (Forced JSON)...")
    for stream_t, v_list in extract_vocabulary(content_text, images, translit_lang, translit_format, target_lang, max_text_char, repetition_penalty_val, partial_assistant_text=partial_text):
        stream_text = stream_t
        if v_list is not None:
            vocab_list = v_list
        yield "", current_source_hash, None, stream_text, content_text, images, None

    if not vocab_list:
        yield "<p>Failed to parse forced JSON.</p>", current_source_hash, None, stream_text, content_text, images, None
        return

    progress(0.6, desc="Generating TTS audio...")
    for i, item in enumerate(vocab_list):
        korean = item.get("korean", "")
        if korean and tts is not None:
            progress(0.6 + 0.3 * (i / len(vocab_list)), desc=f"Generating audio {i+1}/{len(vocab_list)}...")
            try:
                wav, dur = tts.synthesize(
                    korean, voice_style=voice_style, lang="ko",
                    total_steps=12,
                    speed=0.7,
                )
                import numpy as np
                import soundfile as sf
                
                audio_path = f"audio_{i}.wav"
                sf.write(audio_path, wav, 24000)
                item["audio_uri"] = numpy_to_base64_audio(wav, tts.sample_rate)
            except Exception as e:
                print(f"Failed to generate audio for {korean}: {e}")
                item["audio_uri"] = None

    progress(1.0, desc="Rendering flashcards...")
    
    from jinja2 import Environment, BaseLoader
    import json
    
    env = Environment(loader=BaseLoader())
    template = env.from_string(html_template)
    html_output = template.render(
        vocab_list=vocab_list,
        translit_lang=translit_lang,
        target_lang=target_lang
    )
    
    safe_srcdoc = html_output.replace('"', '&quot;')
    yield f'<iframe srcdoc="{safe_srcdoc}" style="width: 100%; height: 650px; border: none; overflow-y: auto;"></iframe>', current_source_hash, vocab_list, stream_text, content_text, images, None


def create_demo():
    example_pdf = get_example_pdf()
    example_audio = get_example_audio()
    
    custom_theme = gr.themes.Soft(
        primary_hue="violet",
        secondary_hue="indigo",
        neutral_hue="slate",
        font=[gr.themes.GoogleFont("Outfit"), "ui-sans-serif", "sans-serif"]
    )

    css = """
    /* Animated Background Styling */
    @keyframes gradientBG {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    
    body, .gradio-container {
        background: linear-gradient(-45deg, #0f2027, #203a43, #2c5364, #1a2a6c, #112240) !important;
        background-size: 400% 400% !important;
        animation: gradientBG 15s ease infinite !important;
        background-attachment: fixed !important;
        color: #f8fafc !important;
        font-family: 'Outfit', sans-serif !important;
    }
    
    /* Advanced Glassmorphism Panels */
    .gradio-container .form, 
    .gradio-container .panel, 
    .gradio-container .box {
        background: rgba(15, 23, 42, 0.4) !important;
        backdrop-filter: blur(20px) !important;
        -webkit-backdrop-filter: blur(20px) !important;
        border-radius: 24px !important;
        border: 1px solid rgba(255, 255, 255, 0.1) !important;
        box-shadow: 0 15px 35px rgba(0, 0, 0, 0.2), inset 0 0 0 1px rgba(255, 255, 255, 0.05) !important;
    }

    /* Headings */
    h1 {
        text-align: center;
        background: linear-gradient(to right, #c084fc, #60a5fa);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-weight: 800 !important;
        font-size: 3em !important;
        margin-bottom: 0.5em !important;
        text-shadow: 0px 4px 30px rgba(192, 132, 252, 0.3);
        letter-spacing: -1px;
    }
    
    /* Buttons */
    button.primary {
        background: linear-gradient(135deg, #8b5cf6, #3b82f6) !important;
        border: none !important;
        box-shadow: 0 0 20px rgba(139, 92, 246, 0.5) !important;
        transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
        font-weight: 700 !important;
        letter-spacing: 1px !important;
        text-transform: uppercase;
        border-radius: 16px !important;
        color: white !important;
    }
    button.primary:hover {
        transform: translateY(-3px) scale(1.02) !important;
        box-shadow: 0 0 30px rgba(139, 92, 246, 0.8) !important;
    }
    button.stop {
        background: linear-gradient(135deg, #ef4444, #b91c1c) !important;
        border: none !important;
        box-shadow: 0 0 20px rgba(239, 68, 68, 0.5) !important;
        transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
        font-weight: 700 !important;
        text-transform: uppercase;
        letter-spacing: 1px !important;
        border-radius: 16px !important;
        color: white !important;
    }
    button.stop:hover {
        transform: translateY(-3px) scale(1.02) !important;
        box-shadow: 0 0 30px rgba(239, 68, 68, 0.8) !important;
    }
    
    /* Inputs */
    textarea, input[type="text"], input[type="number"], select {
        background: rgba(0, 0, 0, 0.3) !important;
        border: 1px solid rgba(255,255,255,0.1) !important;
        border-radius: 12px !important;
        color: #f8fafc !important;
        transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    }
    textarea:focus, input[type="text"]:focus {
        border-color: #a855f7 !important;
        box-shadow: 0 0 15px rgba(168, 85, 247, 0.4), inset 0 0 0 1px rgba(168, 85, 247, 0.4) !important;
    }

    /* Hide Stream Box Duplicate Progress */
    #stream_box .progress-text, 
    #stream_box .progress-level, 
    #stream_box .progress, 
    #stream_box .progress-container { 
        display: none !important; 
    }
    """
    
    with gr.Blocks(title="LocalDuo", theme=custom_theme, css=css) as demo:
        gr.Markdown("# 🇰🇷✨ LocalDuo - Learn Korean from PDFs, Websites & YouTube")
        gr.Markdown("Enter a website URL 🌐, upload a PDF 📄, upload an audio file 🎵, or paste a YouTube link 🎬. The app uses a **Vision-Language Model (VLM)** 🧠, **ASR** 🎤, and **TTS** 🗣️ to generate vocabulary flashcards.")
        
        active_tab = gr.State("Website URL")
        
        with gr.Row():
            with gr.Column(scale=1):
                with gr.Tabs() as input_tabs:
                    with gr.Tab("Website URL", id="tab_url") as tab_url:
                        url_input = gr.Textbox(label="Enter a Website URL 🌐", 
                        placeholder=r"e.g. https://www.bbc.com/korean/articles/cn0p7rkvxdgo", 
                        value=r"https://www.bbc.com/korean/articles/cn0p7rkvxdgo")
                    
                    with gr.Tab("Upload PDF", id="tab_pdf") as tab_pdf:
                        pdf_input = gr.File(label="Upload Book PDF 📚", file_types=[".pdf"], value=example_pdf)
                    
                    with gr.Tab("Upload Audio", id="tab_audio") as tab_audio:
                        audio_file_input = gr.File(label="Upload Audio File 🎵", file_types=[".wav", ".mp3", ".m4a", ".ogg", ".flac", ".opus", ".webm"], value=example_audio)
                        gr.Markdown("*Upload a Korean audio file. It will be transcribed using Cohere ASR and vocabulary will be extracted from the transcript.*", elem_classes=["hint-text"])
                    
                    with gr.Tab("YouTube Link", id="tab_yt") as tab_yt:
                        yt_url_input = gr.Textbox(label="Enter a YouTube Link 🎬", 
                        placeholder=r"e.g. https://www.youtube.com/watch?v=...", 
                        value="https://www.youtube.com/watch?v=9Nj7l73PBWE",
                        info="Audio from the first 5 minutes will be transcribed using Cohere ASR")
                        yt_cookies_input = gr.File(label="YouTube Cookies (cookies.txt)", file_types=[".txt"], value=None,
                            type="filepath")
                        gr.Markdown("*Optional. Helps bypass YouTube bot detection. Install the [cookies.txt](https://addons.mozilla.org/firefox/addon/cookies-txt/) extension, go to youtube.com while logged in, click the extension → 'Current Site' to export.*", elem_classes=["hint-text"])
                
                # Track active tab
                tab_url.select(fn=lambda: "Website URL", inputs=None, outputs=active_tab)
                tab_pdf.select(fn=lambda: "Upload PDF", inputs=None, outputs=active_tab)
                tab_audio.select(fn=lambda: "Upload Audio", inputs=None, outputs=active_tab)
                tab_yt.select(fn=lambda: "YouTube Link", inputs=None, outputs=active_tab)
                
                gr.Markdown("### ⚙️ Customization Settings")
                max_text_char_input = gr.Slider(minimum=1000, maximum=30000, step=1000, value=1500, label="Max Input Text Length (Characters)")
                repetition_penalty_input = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.2, label="Repetition Penalty")
                auto_force_chars_input = gr.Slider(minimum=1_000, maximum=10_000, step=100, value=4_000, label="Auto-force JSON after (chars of thinking)")
                
                with gr.Accordion("🔧 Advanced", open=False):
                    translit_lang = gr.Dropdown(
                        label="Word Transliteration Language", 
                        choices=LANGUAGE_CHOICES, 
                        value="Indo-European - English"
                    )
                    translit_format = gr.Dropdown(label="Transliteration Format", choices=["dashed syllable", "regular word with space"], value="dashed syllable")
                    target_lang = gr.Dropdown(
                        label="Target Language (Full App)", 
                        choices=LANGUAGE_CHOICES, 
                        value="Indo-European - English"
                    )
                
                with gr.Row():
                    submit_btn = gr.Button("✨ Generate Flashcards ✨", variant="primary")
                    stop_thinking_btn = gr.Button("⚡ Stop thinking, Generate now", variant="secondary")
                    stop_btn = gr.Button("🛑 Stop Generation", variant="stop")
            
            with gr.Column(scale=2):
                output_html = gr.HTML(label="Flashcards will appear here")
                stream_box = gr.Textbox(label="Live Model Generation 🧠", lines=10, max_lines=20, interactive=False, autoscroll=True, elem_id="stream_box")
                
                with gr.Accordion("📄 Extracted Source Content", open=True):
                    extracted_text_box = gr.Textbox(label="Extracted Text", lines=10, max_lines=15, interactive=False)
                    extracted_images_gallery = gr.Gallery(label="Extracted Images", columns=4, height="auto", object_fit="contain")
                    extracted_audio_player = gr.Audio(label="Extracted Audio (YouTube / Uploaded)", type="filepath", interactive=False)
                
        last_source_state = gr.State(None)
        last_korean_words_state = gr.State(None)

        def reset_btn_text():
            return gr.update(value="⚡ Stop thinking, Generate now"), gr.update(value="🛑 Stop Generation")

        submit_btn.click(fn=reset_btn_text, inputs=None, outputs=[stop_thinking_btn, stop_btn], queue=False)

        generate_event = submit_btn.click(
            fn=process_pdf, 
            inputs=[pdf_input, url_input, audio_file_input, yt_url_input, yt_cookies_input, translit_lang, translit_format, target_lang, max_text_char_input, repetition_penalty_input, auto_force_chars_input, last_source_state, last_korean_words_state, active_tab], 
            outputs=[output_html, last_source_state, last_korean_words_state, stream_box, extracted_text_box, extracted_images_gallery, extracted_audio_player]
        )
        
        stop_thinking_btn.click(fn=set_stop_thinking, inputs=None, outputs=stop_thinking_btn, queue=False)
        
        stop_btn.click(fn=set_kill_threads, inputs=None, outputs=stop_btn, queue=False)
        
        # Force autoscroll using Custom JS
        stream_box.change(
            fn=None,
            js="""
            function() {
                const ta = document.querySelector('#stream_box textarea');
                if (ta) {
                    ta.scrollTop = ta.scrollHeight;
                }
            }
            """
        )
    return demo


if __name__ == "__main__":
    global model, processor, tts, voice_style, asr_model, asr_processor
    
    
    # model_id = "Qwen/Qwen3.5-9B"
    model_id = "Qwen/Qwen3.5-2B"
    
    print(f"Loading {model_id} model via Transformers...")
    
    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    
    # try:
    #     with open("chat_template.jinja", "r", encoding="utf-8") as f:
    #         processor.chat_template = f.read()
    # except Exception as e:
    #     print("Could not load custom chat template:", e)
    processor.chat_template = """
{%- set image_count = namespace(value=0) %}
{%- set video_count = namespace(value=0) %}
{%- macro render_content(content, do_vision_count, is_system_content=false) %}
    {%- if content is string %}
        {{- content }}
    {%- elif content is iterable and content is not mapping %}
        {%- for item in content %}
            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
                {%- if is_system_content %}
                    {{- raise_exception('System message cannot contain images.') }}
                {%- endif %}
                {%- if do_vision_count %}
                    {%- set image_count.value = image_count.value + 1 %}
                {%- endif %}
                {%- if add_vision_id %}
                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
                {%- endif %}
                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
            {%- elif 'video' in item or item.type == 'video' %}
                {%- if is_system_content %}
                    {{- raise_exception('System message cannot contain videos.') }}
                {%- endif %}
                {%- if do_vision_count %}
                    {%- set video_count.value = video_count.value + 1 %}
                {%- endif %}
                {%- if add_vision_id %}
                    {{- 'Video ' ~ video_count.value ~ ': ' }}
                {%- endif %}
                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
            {%- elif 'text' in item %}
                {{- item.text }}
            {%- else %}
                {{- raise_exception('Unexpected item type in content.') }}
            {%- endif %}
        {%- endfor %}
    {%- elif content is none or content is undefined %}
        {{- '' }}
    {%- else %}
        {{- raise_exception('Unexpected content type.') }}
    {%- endif %}
{%- endmacro %}
{%- if not messages %}
    {{- raise_exception('No messages provided.') }}
{%- endif %}
{%- if tools and tools is iterable and tools is not mapping %}
    {{- '<|im_start|>system\n' }}
    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
    {%- for tool in tools %}
        {{- "\n" }}
        {{- tool | tojson }}
    {%- endfor %}
    {{- "\n</tools>" }}
    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
    {%- if messages[0].role == 'system' %}
        {%- set content = render_content(messages[0].content, false, true)|trim %}
        {%- if content %}
            {{- '\n\n' + content }}
        {%- endif %}
    {%- endif %}
    {{- '<|im_end|>\n' }}
{%- else %}
    {%- if messages[0].role == 'system' %}
        {%- set content = render_content(messages[0].content, false, true)|trim %}
        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
    {%- endif %}
{%- endif %}
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
{%- for message in messages[::-1] %}
    {%- set index = (messages|length - 1) - loop.index0 %}
    {%- if ns.multi_step_tool and message.role == "user" %}
        {%- set content = render_content(message.content, false)|trim %}
        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
            {%- set ns.multi_step_tool = false %}
            {%- set ns.last_query_index = index %}
        {%- endif %}
    {%- endif %}
{%- endfor %}
{%- if ns.multi_step_tool %}
    {{- raise_exception('No user query found in messages.') }}
{%- endif %}
{%- for message in messages %}
    {%- set content = render_content(message.content, true)|trim %}
    {%- if message.role == "system" %}
        {%- if not loop.first %}
            {{- raise_exception('System message must be at the beginning.') }}
        {%- endif %}
    {%- elif message.role == "user" %}
        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
    {%- elif message.role == "assistant" %}
        {%- set reasoning_content = '' %}
        {%- if message.reasoning_content is string %}
            {%- set reasoning_content = message.reasoning_content %}
        {%- else %}
            {%- if '</think>' in content %}
                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
            {%- endif %}
        {%- endif %}
        {%- set reasoning_content = reasoning_content|trim %}
        {%- if loop.index0 > ns.last_query_index %}
            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
        {%- else %}
            {{- '<|im_start|>' + message.role + '\n' + content }}
        {%- endif %}
        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
            {%- for tool_call in message.tool_calls %}
                {%- if tool_call.function is defined %}
                    {%- set tool_call = tool_call.function %}
                {%- endif %}
                {%- if loop.first %}
                    {%- if content|trim %}
                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
                    {%- else %}
                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
                    {%- endif %}
                {%- else %}
                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
                {%- endif %}
                {%- if tool_call.arguments is defined %}
                    {%- for args_name, args_value in tool_call.arguments|items %}
                        {{- '<parameter=' + args_name + '>\n' }}
                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
                        {{- args_value }}
                        {{- '\n</parameter>\n' }}
                    {%- endfor %}
                {%- endif %}
                {{- '</function>\n</tool_call>' }}
            {%- endfor %}
        {%- endif %}
        {{- '<|im_end|>\n' }}
    {%- elif message.role == "tool" %}
        {%- if loop.previtem and loop.previtem.role != "tool" %}
            {{- '<|im_start|>user' }}
        {%- endif %}
        {{- '\n<tool_response>\n' }}
        {{- content }}
        {{- '\n</tool_response>' }}
        {%- if not loop.last and loop.nextitem.role != "tool" %}
            {{- '<|im_end|>\n' }}
        {%- elif loop.last %}
            {{- '<|im_end|>\n' }}
        {%- endif %}
    {%- else %}
        {{- raise_exception('Unexpected message role.') }}
    {%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
    {{- '<|im_start|>assistant\n' }}
    {%- if enable_thinking is defined and enable_thinking is false %}
        {{- '<think>\n\n</think>\n\n' }}
    {%- else %}
        {{- '<think>\n' }}
    {%- endif %}
{%- endif %}
""".strip()
        
    model = AutoModelForImageTextToText.from_pretrained(
        model_id, 
        torch_dtype=torch.bfloat16, 
        device_map="cpu", 
        trust_remote_code=True
    )

    print("Loading Cohere ASR model...")
    from transformers import CohereAsrForConditionalGeneration
    asr_processor = AutoProcessor.from_pretrained("CohereLabs/cohere-transcribe-03-2026")
    asr_model = CohereAsrForConditionalGeneration.from_pretrained(
        "CohereLabs/cohere-transcribe-03-2026",
        device_map="cpu",
    )

    print("Loading Supertonic TTS...")
    tts = TTS(model="supertonic-3")
    try:
        voice_style = tts.get_voice_style("F1")
    except Exception:
        voice_style = tts.get_voice_style(tts.voice_style_names[0])
        
    demo = create_demo()
    if IS_HF:
        demo.launch()
    else:
        demo.launch(server_name="0.0.0.0", server_port=7865)