import gradio as gr
import spaces
import requests
import json

BASE_URL = "https://api.fireworks.ai/inference/v1"
MODEL = "accounts/fireworks/models/minimax-m3"

STYLE_DEFINITIONS = {
    "formal": "Professional, objective, factual tone",
    "sarcastic": "Dry, ironic, lightly mocking",
    "humorous_tech": "Funny, with technology or programming references",
    "humorous_non_tech": "Funny, everyday humour with no technical jargon",
}
ALL_STYLES = list(STYLE_DEFINITIONS.keys())
MAX_CLIPS = 8


def get_video_description(video_source, api_key):
    resp = requests.post(
        f"{BASE_URL}/chat/completions",
        headers={"Authorization": f"Bearer {api_key}"},
        json={
            "model": MODEL,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": (
                            "Describe this video in detail: the setting, subjects, "
                            "actions, and mood. Be factual and specific. 3-5 sentences."
                        )},
                        {"type": "video_url", "video_url": {"url": video_source}},
                    ],
                }
            ],
            "max_tokens": 5000,
        },
        timeout=120,
    )
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"]


def normalize_key(key):
    """Convert various spellings to our standard keys"""
    key_lower = key.lower().strip()
    
    # Map common variations to correct keys
    key_mappings = {
        'sarcastic': 'sarcastic',
        'sarcasm': 'sarcastic',
        'sarcastik': 'sarcastic',
        'sarcastc': 'sarcastic',
        'humorous_tech': 'humorous_tech',
        'tech_humour': 'humorous_tech',
        'tech_humor': 'humorous_tech',
        'humorous_non_tech': 'humorous_non_tech',
        'nontech_humour': 'humorous_non_tech',
        'nontech_humor': 'humorous_non_tech',
        'formal': 'formal',
        'formel': 'formal',
    }
    
    # Try exact match first
    if key_lower in key_mappings:
        return key_mappings[key_lower]
    
    # Try fuzzy match for sarcastic variations
    if 'sarcas' in key_lower:
        return 'sarcastic'
    
    return key_lower


def get_styled_captions(description, styles, api_key, max_retries=2):
    style_list = ", ".join(styles)
    definitions_text = "\n".join(f"- {s}: {STYLE_DEFINITIONS[s]}" for s in styles)
    keys_example = ", ".join(f'"{s}": "..."' for s in styles)

    prompt = f"""Here is a factual description of a video:

"{description}"

Write a caption for this video in EACH of the following styles: {style_list}

Style definitions:
{definitions_text}

IMPORTANT: Respond with ONLY a valid JSON object. Try to use these EXACT key names:
{{{keys_example}}}

If you must use different key names, make them as close as possible to the requested ones.
Return ONLY the JSON, nothing else."""

    last_error = None
    for attempt in range(max_retries + 1):
        resp = requests.post(
            f"{BASE_URL}/chat/completions",
            headers={"Authorization": f"Bearer {api_key}"},
            json={
                "model": MODEL,
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": 5000,
            },
            timeout=120,
        )
        resp.raise_for_status()
        content = resp.json()["choices"][0]["message"]["content"]
        content = content.strip()
        if content.startswith("```"):
            content = content.split("```")[1]
            if content.startswith("json"):
                content = content[4:]
            content = content.strip()
        try:
            parsed = json.loads(content)
        except json.JSONDecodeError as e:
            last_error = f"invalid JSON: {e}"
            continue
        
        # Normalize keys and build result
        result = {}
        missing_styles = []
        
        for style in styles:
            # Try exact match first
            if style in parsed:
                result[style] = parsed[style]
            else:
                # Try to find a matching key
                found = False
                for key, value in parsed.items():
                    normalized = normalize_key(key)
                    if normalized == style:
                        result[style] = value
                        found = True
                        break
                
                if not found:
                    missing_styles.append(style)
        
        # If we have all styles, return them
        if not missing_styles:
            return result
        
        # If we're on the last attempt, try to be more lenient
        if attempt == max_retries:
            # Use whatever we found and fill missing with empty strings
            for style in styles:
                if style not in result:
                    result[style] = ""
            return result
        
        last_error = f"missing/misspelled keys: {missing_styles}"
        continue

    # Fallback: return whatever we have with empty strings for missing
    result = {}
    for style in styles:
        result[style] = ""
    return result


@spaces.GPU
def process_clip(video_url, api_key):
    if not api_key:
        return None, "Missing API key", "", "", "", ""
    if not video_url or not video_url.strip():
        return None, "", "", "", "", ""

    url = video_url.strip()
    try:
        description = get_video_description(url, api_key)
        captions = get_styled_captions(description, ALL_STYLES, api_key)
    except Exception as e:
        return url, f"Error: {e}", "", "", "", ""

    return (
        url,
        description,
        captions.get("formal", ""),
        captions.get("sarcastic", ""),
        captions.get("humorous_tech", ""),
        captions.get("humorous_non_tech", ""),
    )


CUSTOM_CSS = """
:root {
    --bg: #0B0D10;
    --panel: #15181D;
    --panel-alt: #1B1F26;
    --border: #262B33;
    --text-main: #E8EAED;
    --text-sub: #8B93A1;
    --green: #3DDC84;
    --red: #FF5C5C;
    --yellow: #FFC93C;
    --cyan: #4FD1E8;
}

.gradio-container { 
    background: var(--bg) !important; 
    font-family: 'Segoe UI', system-ui, sans-serif !important;
    padding: 20px !important;
}

h1, h2, h3 { color: var(--text-main) !important; }
.markdown-body, p, span, label { color: var(--text-sub) !important; }

#topbar { 
    background: var(--panel); 
    border: 1px solid var(--border); 
    border-radius: 8px; 
    padding: 16px 20px !important;
    margin-bottom: 18px !important;
}

#hero-title { 
    font-size: 2.4rem !important;
    font-weight: 700 !important; 
    letter-spacing: -0.5px; 
    margin-bottom: 4px !important; 
    padding: 4px 0 !important;
}

#hero-title span.agent-highlight {
    color: #EE316B !important;
}

#hero-sub { 
    color: var(--text-sub) !important; 
    font-size: 1rem !important;
    margin-top: 0 !important;
    padding: 4px 0 !important;
}

.gr-panel, .block { 
    background: var(--panel) !important; 
    border: 1px solid var(--border) !important; 
    border-radius: 8px !important;
    padding: 16px !important;
}

/* Fix for URL input and Generate button - integrated design */
.url-row {
    display: flex !important;
    align-items: center !important;
    gap: 10px !important;
    background: var(--panel) !important;
    border: 1px solid var(--border) !important;
    border-radius: 8px !important;
    padding: 8px 12px !important;
    margin-bottom: 12px !important;
}

.url-row .gr-textbox {
    flex: 1 !important;
}

.url-row .gr-textbox label {
    display: none !important;
}

.url-row .gr-textbox input {
    background: #0F1216 !important;
    border: 1px solid var(--border) !important;
    border-radius: 6px !important;
    padding: 8px 12px !important;
    color: var(--text-main) !important;
    height: 38px !important;
}

.url-row .gr-button {
    height: 38px !important;
    min-width: 100px !important;
    white-space: nowrap !important;
    background: var(--yellow) !important;
    color: #0B0D10 !important;
    border: none !important;
    border-radius: 6px !important;
    font-weight: 700 !important;
    padding: 8px 20px !important;
}

.url-row .gr-button:hover {
    opacity: 0.85 !important;
}

/* Make style cards equal height with no gap */
.style-row {
    display: flex !important;
    gap: 0px !important;
    align-items: stretch !important;
    margin: 0 !important;
    padding: 0 !important;
}

.style-row .gr-column {
    display: flex !important;
    flex: 1 !important;
    padding: 0 !important;
    margin: 0 !important;
}

.style-row .gr-column:first-child .style-card {
    border-radius: 8px 0 0 8px !important;
    border-right: none !important;
}

.style-row .gr-column:last-child .style-card {
    border-radius: 0 8px 8px 0 !important;
}

.style-card {
    background: var(--panel-alt) !important; 
    border: 1px solid var(--border) !important; 
    padding: 12px 14px !important;
    height: 100% !important;
    min-height: 100px !important;
    display: flex !important;
    flex-direction: column !important;
    width: 100% !important;
    margin: 0 !important;
}

.style-card .markdown-body {
    flex: 1 !important;
    padding: 0 !important;
    overflow-wrap: break-word !important;
    word-wrap: break-word !important;
    margin-top: 0 !important;
}

.style-card .markdown-body p {
    margin: 0 !important;
    padding: 0 !important;
}

/* Style labels - no extra spacing */
.style-label-formal, 
.style-label-sarcastic,
.style-label-tech,
.style-label-nontech {
    color: var(--cyan) !important; 
    font-size: 0.7rem !important;
    font-weight: 700 !important; 
    text-transform: uppercase !important; 
    letter-spacing: 0.6px !important;
    padding: 0 0 6px 0 !important;
    margin: 0 0 6px 0 !important;
    flex-shrink: 0 !important;
    border-bottom: 1px solid var(--border) !important;
}

.style-label-sarcastic { color: var(--yellow) !important; }
.style-label-tech { color: var(--green) !important; }
.style-label-nontech { color: var(--red) !important; }

/* Description box styling */
.desc-box textarea {
    min-height: 70px !important;
    max-height: 100px !important;
    overflow-y: auto !important;
    white-space: pre-wrap !important;
    word-wrap: break-word !important;
    padding: 10px 12px !important;
    line-height: 1.5 !important;
    background: #0F1216 !important;
    border: 1px solid var(--border) !important;
    border-radius: 6px !important;
    color: var(--text-main) !important;
}

input, textarea { 
    background: #0F1216 !important; 
    color: var(--text-main) !important; 
    border: 1px solid var(--border) !important; 
    border-radius: 6px !important; 
    padding: 10px 12px !important;
}

label span { 
    color: var(--text-sub) !important; 
    font-size: 0.72rem !important; 
    text-transform: uppercase; 
    letter-spacing: 0.6px; 
}

button#add-btn { 
    background: transparent !important; 
    color: var(--cyan) !important; 
    border: 1px solid var(--cyan) !important; 
    font-weight: 600 !important;
    padding: 8px 16px !important;
}

button#remove-btn { 
    background: transparent !important; 
    color: var(--red) !important; 
    border: 1px solid var(--red) !important; 
    font-weight: 600 !important;
    padding: 8px 16px !important;
}

button#save-key-btn { 
    background: var(--yellow) !important; 
    color: #0B0D10 !important; 
    border: none !important; 
    font-weight: 700 !important;
    padding: 8px 16px !important;
}

/* Remove extra spacing */
.group {
    margin-bottom: 4px !important;
}

.gr-form {
    gap: 4px !important;
}

/* Clip header spacing */
.clip-header {
    margin-bottom: 4px !important;
    padding: 0 !important;
    color: var(--text-main) !important;
}
"""

with gr.Blocks(title="Video Caption Agent") as demo:
    with gr.Row(elem_id="topbar"):
        api_key = gr.Textbox(
            label="Fireworks API Key", type="password", placeholder="fw_...",
            scale=4, container=True,
        )

    gr.HTML('<h3 id="hero-title">VIDEO CAPTION <span class="agent-highlight">AGENT</span></h3>')
    gr.Markdown(
        "Watches any video and writes captions in four voices \u2014 formal, sarcastic, tech-humour, everyday humour.",
        elem_id="hero-sub",
    )

    visible_count = gr.State(1)
    url_boxes, gen_buttons, previews = [], [], []
    desc_boxes, formal_boxes, sarcastic_boxes, tech_boxes, nontech_boxes = [], [], [], [], []
    clip_groups = []

    for i in range(MAX_CLIPS):
        with gr.Group(visible=(i == 0)) as grp:
            # Clip header
            gr.HTML(f'<div class="clip-header"><strong>Clip {i + 1}</strong></div>')
            
            # Integrated URL + Generate button row
            with gr.Row(elem_classes="url-row"):
                url_box = gr.Textbox(
                    label="", 
                    placeholder="https://example.com/video.mp4",
                    scale=5,
                    container=False,
                )
                gen_btn = gr.Button("Generate", elem_id="gen-btn", scale=1)

            preview = gr.Video(label="Preview", height=280)
            
            # Description with better styling
            desc_box = gr.Textbox(
                label="Description", 
                lines=3, 
                interactive=False,
                elem_classes="desc-box"
            )

            # Row 1: Formal & Sarcastic - no gap
            with gr.Row(elem_classes="style-row"):
                with gr.Column():
                    with gr.Group(elem_classes="style-card"):
                        gr.HTML('<div class="style-label-formal">FORMAL</div>')
                        formal_box = gr.Markdown()
                with gr.Column():
                    with gr.Group(elem_classes="style-card"):
                        gr.HTML('<div class="style-label-sarcastic">SARCASTIC</div>')
                        sarcastic_box = gr.Markdown()
            
            # Row 2: Tech & Non-Tech - no gap
            with gr.Row(elem_classes="style-row"):
                with gr.Column():
                    with gr.Group(elem_classes="style-card"):
                        gr.HTML('<div class="style-label-tech">HUMOROUS — TECH</div>')
                        tech_box = gr.Markdown()
                with gr.Column():
                    with gr.Group(elem_classes="style-card"):
                        gr.HTML('<div class="style-label-nontech">HUMOROUS — NON-TECH</div>')
                        nontech_box = gr.Markdown()

        clip_groups.append(grp)
        url_boxes.append(url_box)
        gen_buttons.append(gen_btn)
        previews.append(preview)
        desc_boxes.append(desc_box)
        formal_boxes.append(formal_box)
        sarcastic_boxes.append(sarcastic_box)
        tech_boxes.append(tech_box)
        nontech_boxes.append(nontech_box)

        gen_btn.click(
            process_clip,
            inputs=[url_box, api_key],
            outputs=[preview, desc_box, formal_box, sarcastic_box, tech_box, nontech_box],
        )

    with gr.Row():
        add_btn = gr.Button("+ Add another video", elem_id="add-btn")
        remove_btn = gr.Button("\u2212 Remove last video", elem_id="remove-btn")

    def add_clip(count):
        count = min(count + 1, MAX_CLIPS)
        return [gr.update(visible=(i < count)) for i in range(MAX_CLIPS)] + [count]

    def remove_clip(count):
        count = max(count - 1, 1)
        return [gr.update(visible=(i < count)) for i in range(MAX_CLIPS)] + [count]

    add_btn.click(add_clip, inputs=[visible_count], outputs=clip_groups + [visible_count])
    remove_btn.click(remove_clip, inputs=[visible_count], outputs=clip_groups + [visible_count])

demo.launch(css=CUSTOM_CSS)