import gradio as gr
import json
import anthropic

client = anthropic.Anthropic()

LANGUAGES = [
    "Hindi", "Arabic", "French", "German", "Spanish", "Chinese",
    "Japanese", "Russian", "Urdu", "Turkish", "Korean", "Italian",
    "Portuguese", "English",
]

GAP_MS        = 200    # ms gap between words → new line
MAX_SPEED     = 2.5    # maximum TTS speed multiplier
MAX_CHARS_RATIO = MAX_SPEED  # at 2.5x speed, can fit 2.5x characters


def group_words_into_lines(words: list, gap_ms: float = GAP_MS) -> list:
    """
    Group word-level timestamps into subtitle lines.
    A new line starts whenever the gap between two consecutive words
    exceeds gap_ms milliseconds.

    Each word dict must have: {"start": float, "end": float, "word": str}
    Returns list of line dicts:
      {"start": float, "end": float, "text": str, "words": [...]}
    """
    if not words:
        return []

    lines = []
    current_words = [words[0]]

    for w in words[1:]:
        prev_end   = current_words[-1]["end"]
        gap_secs   = w["start"] - prev_end
        gap_ms_val = gap_secs * 1000

        if gap_ms_val > gap_ms:
            # flush current line
            lines.append({
                "start": current_words[0]["start"],
                "end":   current_words[-1]["end"],
                "text":  " ".join(cw["word"].strip() for cw in current_words),
                "words": current_words,
            })
            current_words = [w]
        else:
            current_words.append(w)

    # flush last line
    if current_words:
        lines.append({
            "start": current_words[0]["start"],
            "end":   current_words[-1]["end"],
            "text":  " ".join(cw["word"].strip() for cw in current_words),
            "words": current_words,
        })

    return lines


def fmt_time(s: float) -> str:
    m   = int(s) // 60
    sec = s - m * 60
    return f"{m:02d}:{sec:05.2f}"


def translate_line(text: str, src_lang: str, tgt_lang: str,
                   max_chars: int) -> str:
    """
    Translate one line via Claude, instructing it to stay under max_chars
    while NEVER cutting content — just use natural concise phrasing.
    """
    prompt = (
        f"Translate the following {src_lang} text to {tgt_lang}.\n"
        f"Rules:\n"
        f"1. Never cut or omit any meaning or content.\n"
        f"2. Try to keep the translation under {max_chars} characters "
        f"by using natural, concise phrasing in {tgt_lang}.\n"
        f"3. If it is impossible to stay under {max_chars} characters "
        f"without cutting content, go slightly over — completeness wins.\n"
        f"4. Output ONLY the translated text, nothing else.\n\n"
        f"Text: {text}"
    )

    message = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=1000,
        messages=[{"role": "user", "content": prompt}],
    )
    return message.content[0].text.strip()


def process(words_json: str, src_lang: str, tgt_lang: str,
            max_speed: float, gap_ms: float) -> tuple:
    """
    Main processing function.
    Returns (display_text, json_output).
    """
    if not words_json.strip():
        return "❌ Please paste word-level JSON.", ""

    try:
        data = json.loads(words_json)
    except json.JSONDecodeError as e:
        return f"❌ Invalid JSON: {e}", ""

    # Accept both {"words": [...]} and bare [...]
    if isinstance(data, dict):
        words = data.get("words", [])
    elif isinstance(data, list):
        words = data
    else:
        return "❌ JSON must be a list of words or {\"words\": [...]}", ""

    if not words:
        return "❌ No words found in JSON.", ""

    # Validate first word has required fields
    required = {"start", "end", "word"}
    if not required.issubset(words[0].keys()):
        return f"❌ Each word must have: {required}", ""

    lines = group_words_into_lines(words, gap_ms_override := gap_ms)

    results   = []
    display   = []
    warnings  = []

    for i, line in enumerate(lines):
        src_text  = line["text"]
        src_chars = len(src_text)
        max_chars = int(src_chars * max_speed)

        # Translate
        tgt_text  = translate_line(src_text, src_lang, tgt_lang, max_chars)
        tgt_chars = len(tgt_text)

        # Speed factor needed to fit translation in the same time slot
        # (proportional to character count ratio)
        ratio          = tgt_chars / max(src_chars, 1)
        required_speed = round(ratio, 3)
        over_limit     = required_speed > max_speed

        slot_start = line["start"]
        slot_end   = line["end"]

        result = {
            "line":           i + 1,
            "start":          round(slot_start, 3),
            "end":            round(slot_end, 3),
            "original":       src_text,
            "translated":     tgt_text,
            "original_chars": src_chars,
            "translated_chars": tgt_chars,
            "char_ratio":     round(ratio, 3),
            "required_speed": required_speed,
            "max_speed":      max_speed,
            "over_limit":     over_limit,
        }
        results.append(result)

        # Human-readable display line
        flag = "⚠️ " if over_limit else "✅ "
        display.append(
            f"{flag}[{fmt_time(slot_start)} → {fmt_time(slot_end)}]\n"
            f"  {src_lang}: {src_text}\n"
            f"  {tgt_lang}: {tgt_text}\n"
            f"  chars: {src_chars} → {tgt_chars}  "
            f"speed: {required_speed}x"
            + (" ⚠️ OVER {:.1f}x LIMIT".format(max_speed) if over_limit else "")
        )

        if over_limit:
            warnings.append(
                f"Line {i+1}: needs {required_speed}x but limit is {max_speed}x"
            )

    summary = (
        f"✅ {len(lines)} lines processed. "
        f"{'⚠️ ' + str(len(warnings)) + ' lines over speed limit.' if warnings else 'All within speed limit.'}"
    )

    display_text = summary + "\n\n" + "\n\n".join(display)
    json_out     = json.dumps(results, ensure_ascii=False, indent=2)

    return display_text, json_out


with gr.Blocks(title="Dubbing Line Builder") as demo:
    gr.Markdown("""
    # 🎬 Dubbing Line Builder
    Paste word-level timestamps → groups into lines (200ms gap rule) →
    translates → calculates required TTS speed per line.

    **Input JSON format** (from your Whisper API):
    ```json
    [{"start": 0.5, "end": 0.9, "word": "Hello"},
     {"start": 1.0, "end": 1.4, "word": "world"}]
    ```
    or `{"words": [...]}` shape also accepted.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            words_input = gr.Textbox(
                label="📋 Word-level JSON",
                lines=15,
                placeholder='[{"start": 0.5, "end": 0.9, "word": "Hello"}, ...]'
            )
            with gr.Row():
                src_lang = gr.Dropdown(
                    label="📢 Source language",
                    choices=LANGUAGES, value="English"
                )
                tgt_lang = gr.Dropdown(
                    label="🎯 Target language",
                    choices=LANGUAGES, value="Hindi"
                )
            with gr.Row():
                max_speed_input = gr.Slider(
                    label="⚡ Max TTS speed (x)",
                    minimum=1.0, maximum=3.0, step=0.1, value=2.5,
                    info="Max chars = original × this value"
                )
                gap_input = gr.Slider(
                    label="⏱ Gap rule (ms)",
                    minimum=50, maximum=1000, step=50, value=200,
                    info="New line if silence > this many ms"
                )
            btn = gr.Button("🚀 Build Lines & Translate", variant="primary")

        with gr.Column(scale=1):
            display_out = gr.Textbox(label="📝 Result", lines=25)
            json_out    = gr.Textbox(label="📦 JSON output (for your app)", lines=15)

    btn.click(
        fn=process,
        inputs=[words_input, src_lang, tgt_lang, max_speed_input, gap_input],
        outputs=[display_out, json_out],
    )

    gr.api(process, api_name="build_lines")

if __name__ == "__main__":
    demo.launch()