import os
import base64
import time
import json
import io
import requests
from PIL import Image
import gradio as gr

# --- Configuration ---
API_URL = os.environ.get("ARCISVLM_API_URL", "http://localhost:8000")
QUERY_ENDPOINT = f"{API_URL}/api/v1/query"

# --- Agent type defaults ---
AGENT_DEFAULTS = {
    "Caption": "Describe this surveillance scene in detail.",
    "Detect": "Detect and identify all objects and people in this image.",
    "Alert": "Identify any safety threats, intrusions, or anomalies that require immediate attention.",
    "VQA": "What is happening in this scene?",
    "Count": "Count the number of people and vehicles in this image.",
    "OCR": "Extract and transcribe all visible text, license plates, and signage.",
    "Reason": "Analyze this scene and reason about potential risks or notable behaviors.",
    "Track": "Describe the movement patterns and trajectories of subjects in this scene.",
}

AGENT_TYPES = list(AGENT_DEFAULTS.keys())

TASK_TYPE_MAP = {
    "Caption": "caption",
    "Detect": "detect",
    "Alert": "alert",
    "VQA": "vqa",
    "Count": "count",
    "OCR": "ocr",
    "Reason": "reason",
    "Track": "track",
}


# --- Helpers ---

def image_to_base64(pil_image: Image.Image) -> str:
    buffered = io.BytesIO()
    pil_image.save(buffered, format="JPEG", quality=90)
    return base64.b64encode(buffered.getvalue()).decode("utf-8")


def base64_to_pil(b64_str: str) -> Image.Image:
    img_bytes = base64.b64decode(b64_str)
    return Image.open(io.BytesIO(img_bytes))


def format_detections(detections: list) -> str:
    if not detections:
        return ""
    lines = [f"**Detections ({len(detections)} objects):**"]
    for i, det in enumerate(detections, 1):
        label = det.get("label", "unknown")
        conf = det.get("confidence", None)
        bbox = det.get("bbox", None)
        line = f"  {i}. {label}"
        if conf is not None:
            line += f" ({conf:.1%})"
        if bbox:
            line += f" — bbox: {bbox}"
        lines.append(line)
    return "\n".join(lines)


def format_counts(counts: dict) -> str:
    if not counts:
        return ""
    lines = ["**Object Counts:**"]
    for label, count in counts.items():
        lines.append(f"  • {label}: {count}")
    return "\n".join(lines)


def format_alert(alert: dict) -> str:
    if not alert:
        return ""
    severity = alert.get("severity", "unknown").upper()
    message = alert.get("message", "")
    triggered = alert.get("triggered", False)
    icon = "🔴" if triggered else "🟢"
    return f"**Alert [{severity}]** {icon}\n{message}"


# --- Core inference function ---

def run_inference(image, agent_type, question):
    if image is None:
        return (
            "Please upload an image to analyze.",
            "—",
            None,
        )

    if not question or not question.strip():
        question = AGENT_DEFAULTS.get(agent_type, "Describe this image.")

    task_type = TASK_TYPE_MAP.get(agent_type, "vqa")

    # Convert PIL image to base64
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    image_b64 = image_to_base64(image)

    payload = {
        "question": question.strip(),
        "task_type": task_type,
        "image_base64": image_b64,
    }

    start = time.time()
    try:
        response = requests.post(
            QUERY_ENDPOINT,
            json=payload,
            timeout=60,
            headers={"Content-Type": "application/json"},
        )
        response.raise_for_status()
        data = response.json()
    except requests.exceptions.ConnectionError:
        elapsed = (time.time() - start) * 1000
        return (
            f"Connection error: Could not reach API at {API_URL}.\n"
            "Please ensure the ArcisVLM backend is running.",
            f"{elapsed:.0f} ms (failed)",
            None,
        )
    except requests.exceptions.Timeout:
        elapsed = (time.time() - start) * 1000
        return (
            "Request timed out after 60 seconds.",
            f"{elapsed:.0f} ms (timeout)",
            None,
        )
    except requests.exceptions.HTTPError as e:
        elapsed = (time.time() - start) * 1000
        try:
            detail = response.json().get("detail", str(e))
        except Exception:
            detail = str(e)
        return (
            f"API error ({response.status_code}): {detail}",
            f"{elapsed:.0f} ms (error)",
            None,
        )
    except Exception as e:
        elapsed = (time.time() - start) * 1000
        return (
            f"Unexpected error: {str(e)}",
            f"{elapsed:.0f} ms (error)",
            None,
        )

    elapsed = (time.time() - start) * 1000

    # Parse response
    answer = data.get("answer", "(no answer returned)")
    api_time = data.get("processing_time_ms", None)
    annotated_b64 = data.get("annotated_frame_base64", None)
    detections = data.get("detections", [])
    alert = data.get("alert", None)
    counts = data.get("counts", {})

    # Build full answer text
    parts = [answer]

    det_text = format_detections(detections)
    if det_text:
        parts.append("\n" + det_text)

    count_text = format_counts(counts)
    if count_text:
        parts.append("\n" + count_text)

    alert_text = format_alert(alert)
    if alert_text:
        parts.append("\n" + alert_text)

    full_answer = "\n".join(parts)

    # Timing display
    if api_time is not None:
        time_display = f"{api_time:.0f} ms (API) / {elapsed:.0f} ms (total)"
    else:
        time_display = f"{elapsed:.0f} ms (round-trip)"

    # Annotated image
    annotated_image = None
    if annotated_b64:
        try:
            annotated_image = base64_to_pil(annotated_b64)
        except Exception:
            annotated_image = None

    return full_answer, time_display, annotated_image


def update_question(agent_type):
    return AGENT_DEFAULTS.get(agent_type, "")


# --- Gradio UI ---

THEME = gr.themes.Default(
    primary_hue="violet",
    secondary_hue="slate",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "sans-serif"],
).set(
    body_background_fill="#0f1117",
    body_background_fill_dark="#0f1117",
    block_background_fill="#1a1d27",
    block_background_fill_dark="#1a1d27",
    block_border_color="#2d3148",
    block_border_color_dark="#2d3148",
    block_label_background_fill="#1a1d27",
    block_label_background_fill_dark="#1a1d27",
    input_background_fill="#12141e",
    input_background_fill_dark="#12141e",
    button_primary_background_fill="*primary_500",
    button_primary_background_fill_hover="*primary_400",
    button_primary_text_color="white",
    block_title_text_color="#e2e8f0",
    block_title_text_color_dark="#e2e8f0",
    body_text_color="#cbd5e1",
    body_text_color_dark="#cbd5e1",
)

CSS = """
#title-block {
    text-align: center;
    padding: 24px 0 8px 0;
}
#title-block h1 {
    font-size: 2rem;
    font-weight: 700;
    background: linear-gradient(135deg, #818cf8 0%, #a78bfa 50%, #38bdf8 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin-bottom: 6px;
}
#title-block p {
    color: #94a3b8;
    font-size: 0.95rem;
    line-height: 1.5;
}
.badge-row {
    display: flex;
    justify-content: center;
    gap: 10px;
    flex-wrap: wrap;
    margin-top: 10px;
    margin-bottom: 4px;
}
.badge {
    background: #1e2235;
    border: 1px solid #374151;
    border-radius: 20px;
    padding: 3px 12px;
    font-size: 0.78rem;
    color: #94a3b8;
}
#run-btn {
    font-size: 1rem;
    font-weight: 600;
    letter-spacing: 0.03em;
    height: 48px;
}
#answer-output textarea {
    font-size: 0.92rem;
    line-height: 1.6;
}
#timing-output input {
    font-size: 0.82rem;
    color: #64748b;
}
footer {
    display: none !important;
}
"""

DESCRIPTION_HTML = """
<div id="title-block">
  <h1>ArcisVLM — Intelligent Video Analytics</h1>
  <p>
    Powered by <strong>Gemma 4 E2B</strong> multimodal backbone with 8 specialized surveillance agents.<br>
    Upload a camera frame and select an agent to analyze the scene in real time.
  </p>
  <div class="badge-row">
    <span class="badge">Gemma 4 E2B</span>
    <span class="badge">VL-JEPA Encoder</span>
    <span class="badge">MoE Router</span>
    <span class="badge">8 Agents</span>
    <span class="badge">REST API</span>
  </div>
</div>
"""

AGENT_INFO = {
    "Caption":  "Generates a natural-language description of the scene.",
    "Detect":   "Localizes and labels objects and persons with bounding boxes.",
    "Alert":    "Flags threats, intrusions, or policy violations with severity.",
    "VQA":      "Answers free-form questions about the image.",
    "Count":    "Counts objects or people by category.",
    "OCR":      "Reads text, license plates, and signage.",
    "Reason":   "Multi-step reasoning about behaviors and risks.",
    "Track":    "Describes subject movement and trajectory patterns.",
}

def get_agent_info(agent_type):
    return AGENT_INFO.get(agent_type, "")


with gr.Blocks(theme=THEME, css=CSS, title="ArcisVLM") as demo:

    gr.HTML(DESCRIPTION_HTML)

    with gr.Row():
        # Left column: inputs
        with gr.Column(scale=1, min_width=340):
            image_input = gr.Image(
                label="Camera Frame",
                type="pil",
                sources=["upload", "clipboard"],
                height=300,
            )

            agent_dropdown = gr.Dropdown(
                choices=AGENT_TYPES,
                value="Caption",
                label="Agent Type",
                interactive=True,
            )

            agent_info_box = gr.Markdown(
                value=get_agent_info("Caption"),
                elem_id="agent-info",
            )

            question_input = gr.Textbox(
                label="Question / Prompt",
                value=AGENT_DEFAULTS["Caption"],
                lines=2,
                placeholder="Enter your question or use the default for the selected agent.",
            )

            run_btn = gr.Button(
                "Run Analysis",
                variant="primary",
                elem_id="run-btn",
            )

        # Right column: outputs
        with gr.Column(scale=1, min_width=340):
            answer_output = gr.Textbox(
                label="Agent Response",
                lines=10,
                interactive=False,
                elem_id="answer-output",
                # # show_copy_button=True,
            )

            timing_output = gr.Textbox(
                label="Processing Time",
                interactive=False,
                elem_id="timing-output",
            )

            annotated_output = gr.Image(
                label="Annotated Frame",
                type="pil",
                interactive=False,
                height=280,
                visible=True,
            )

    # Example images section
    gr.Markdown("### Example Prompts by Agent")
    with gr.Row():
        for agent, desc in AGENT_INFO.items():
            with gr.Column(min_width=120, scale=1):
                gr.Markdown(f"**{agent}**\n\n_{desc}_")

    # --- Event handlers ---

    def on_agent_change(agent_type):
        question = AGENT_DEFAULTS.get(agent_type, "")
        info = get_agent_info(agent_type)
        return question, info

    agent_dropdown.change(
        fn=on_agent_change,
        inputs=[agent_dropdown],
        outputs=[question_input, agent_info_box],
    )

    run_btn.click(
        fn=run_inference,
        inputs=[image_input, agent_dropdown, question_input],
        outputs=[answer_output, timing_output, annotated_output],
    )

    question_input.submit(
        fn=run_inference,
        inputs=[image_input, agent_dropdown, question_input],
        outputs=[answer_output, timing_output, annotated_output],
    )

    gr.Markdown(
        "<center style='color:#475569;font-size:0.8rem;margin-top:16px;'>"
        "ArcisVLM · Adiance Technologies · "
        f"API: <code>{API_URL}</code>"
        "</center>"
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
    )