server.py

"""gradio.Server backend for the custom Tiny Trigger dashboard.

Serves the built React frontend (``frontend/dist``) and exposes the Tiny Trigger engine as
``@app.api`` endpoints that keep Gradio's queuing / SSE / gradio_client
compatibility. The frontend talks to these via the ``@gradio/client`` JS library.

    poetry run python server.py        # serves the dashboard on :7860

The heavy ML imports stay lazy (inside ``tiny_trigger``); importing this module
does not require torch/ultralytics.
"""

from __future__ import annotations

import os
import logging
from pathlib import Path
from typing import Any

from gradio import FileData, Server
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse

from tiny_trigger import (
    DEFAULT_ANTHROPIC_MODEL,
    DEFAULT_OPENAI_MODEL,
    DEFAULT_REPLICATE_MODEL,
    compile_automation_with_anthropic,
    compile_automation_with_openai,
    compile_automation_with_replicate,
    evaluate_video_detections,
    load_automation_text,
    parse_class_prompt,
    process_video,
)
from tiny_trigger.automation import ActionSpec, AutomationDocument, AutomationRule
from tiny_trigger.automation import document_labels, rule_labels
from tiny_trigger.actions import dispatch_events
from tiny_trigger.store import (
    load_local_config,
    load_saved_automations,
    save_automations,
    append_events,
)
from tiny_trigger.video import render_automation_video

# ── paths ──────────────────────────────────────────────────────────────────
ROOT = Path(__file__).parent
DIST = ROOT / "frontend" / "dist"
ASSETS = DIST / "assets"
RENDERS = ROOT / ".local" / "renders"
RENDERS.mkdir(parents=True, exist_ok=True)

DEFAULT_MODEL = "yoloe-26s-seg.pt"
logging.basicConfig(
    level=os.environ.get("TINY_TRIGGER_LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
)
LOGGER = logging.getLogger(__name__)


# ── serialization helpers (JSON-shaped, for the custom frontend) ─────────────
def _detection_dict(d: Any) -> dict[str, Any]:
    return {
        "frame_index": d.frame_index,
        "timestamp_sec": round(d.timestamp_sec, 3),
        "label": d.label,
        "confidence": round(d.confidence, 4),
        "bbox_xyxy": [round(v, 1) for v in d.bbox_xyxy],
        "bbox_xyxy_norm": [round(v, 4) for v in d.bbox_xyxy_norm],
        "track_id": d.track_id,
    }


def _event_dict(e: Any) -> dict[str, Any]:
    return {
        "rule": e.rule,
        "action": e.action,
        "type": e.type,
        "frame_index": e.frame_index,
        "timestamp_sec": round(e.timestamp_sec, 3),
        "status": e.status,
        "url": e.url,
        "response_status": e.response_status,
        "error": e.error,
    }


def _media_url(path_str: str) -> str:
    """Map a rendered file path under RENDERS to a served /media URL."""
    return f"/media/{Path(path_str).name}"


def _file_path(video: Any) -> str | None:
    """FileData arrives as a dict over the wire; tolerate the object form too."""
    if video is None:
        return None
    if isinstance(video, dict):
        return video.get("path")
    return getattr(video, "path", None)


# ── the server ───────────────────────────────────────────────────────────────
app = Server(title="Tiny Trigger", description="Open-vocabulary video automations")


@app.api(name="detect_and_automate")
def detect_and_automate(
    video: FileData,
    classes: str,
    rules_text: str,
    confidence: float = 0.25,
    frame_stride: int = 5,
    sample_interval_sec: float | None = None,
    max_frames: int = 120,
    model_name: str = DEFAULT_MODEL,
    image_size: int = 0,
    device: str = "auto",
    max_detections: int = 0,
    enable_webhooks: bool = False,
    webhook_url: str = "",
) -> dict:
    """Detect once, evaluate rules, dispatch actions, render an overlay clip."""
    video_path = _file_path(video)
    if not video_path:
        raise ValueError("A video file is required.")

    rules = load_automation_text(rules_text)
    detection_classes = _merge_class_names(parse_class_prompt(classes), document_labels(rules))
    tracking_enabled = _document_uses_moving(rules)
    LOGGER.info(
        "Starting detect_and_automate video=%s classes=%s rules=%s sample_interval=%s max_frames=%s model=%s image_size=%s tracking=%s",
        Path(video_path).name,
        ", ".join(detection_classes),
        len(rules.rules),
        sample_interval_sec,
        max_frames,
        model_name or DEFAULT_MODEL,
        image_size or "default",
        tracking_enabled,
    )

    result = process_video(
        video_path=video_path,
        class_prompt=detection_classes,
        confidence=confidence,
        frame_stride=frame_stride,
        sample_interval_sec=sample_interval_sec,
        max_frames=max_frames,
        model_name=model_name or DEFAULT_MODEL,
        image_size=image_size or None,
        device=None if device in ("", "auto") else device,
        max_detections=max_detections or None,
        tracking_enabled=tracking_enabled,
        output_dir=str(RENDERS),
    )
    LOGGER.info(
        "Detection complete: sampled_frames=%s detections=%s annotated=%s",
        result.processed_frames,
        len(result.detections),
        result.output_video_path,
    )

    events, _last_fired = evaluate_video_detections(
        rules.rules,
        result.detections,
        frames=result.frames,
        # Uploaded videos use clip-relative timestamps, so cooldowns reset for
        # each run. A live camera mode can persist wall-clock cooldowns later.
        last_fired=None,
    )
    dispatched = dispatch_events(
        events, enable_webhooks=enable_webhooks, webhook_url=webhook_url or None
    )
    append_events(dispatched)
    LOGGER.info("Automation evaluation complete: events=%s dispatched=%s", len(events), len(dispatched))

    automation_path = render_automation_video(
        source_video_path=video_path,
        detections=result.detections,
        events=dispatched,
        frame_stride=result.frame_stride,
        max_frames=max_frames,
        output_dir=str(RENDERS),
    )
    LOGGER.info("detect_and_automate complete: output=%s", automation_path)

    fired = len(dispatched)
    return {
        "status": "fired" if fired else "no_match",
        "video_url": _media_url(automation_path),
        "annotated_url": _media_url(result.output_video_path),
        "classes": result.classes,
        "stats": {
            "detections": len(result.detections),
            "rules": len(rules.rules),
            "actions": fired,
            "processed_frames": result.processed_frames,
            "source_fps": round(result.source_fps, 2),
            "output_fps": round(result.output_fps, 2),
            "frame_stride": result.frame_stride,
            "sample_interval_sec": result.sample_interval_sec,
        },
        "detections": [_detection_dict(d) for d in result.detections],
        "events": [_event_dict(e) for e in dispatched],
    }


@app.api(name="compile_rules")
def compile_rules(
    instruction: str,
    classes: str = "",
    existing_rules_text: str = "",
    append: bool = True,
    provider: str = "anthropic",
    api_key: str = "",
    model: str = "",
    replicate_model: str = DEFAULT_REPLICATE_MODEL,
    replicate_reasoning_effort: str = "medium",
) -> dict:
    """Compile a natural-language request into validated automation rules."""
    class_names = parse_class_prompt(classes) if classes else []
    if existing_rules_text.strip():
        existing = load_automation_text(existing_rules_text)
        class_names = _merge_class_names(class_names, document_labels(existing))
    cfg = load_local_config()
    if provider == "replicate":
        api_token = api_key or os.environ.get("REPLICATE_API_TOKEN") or cfg.replicate_api_token
        if not api_token:
            raise ValueError("Paste a Replicate API token or set REPLICATE_API_TOKEN.")
        compiled = compile_automation_with_replicate(
            instruction=instruction,
            class_names=class_names,
            api_token=api_token,
            model=model or replicate_model or cfg.replicate_model or DEFAULT_REPLICATE_MODEL,
            reasoning_effort=(
                replicate_reasoning_effort or cfg.replicate_reasoning_effort or "medium"
            ),
        )
    elif provider == "openai":
        openai_key = api_key or os.environ.get("OPENAI_API_KEY") or cfg.openai_api_key
        if not openai_key:
            raise ValueError("Paste an OpenAI API key or set OPENAI_API_KEY.")
        compiled = compile_automation_with_openai(
            instruction=instruction,
            class_names=class_names,
            api_key=openai_key,
            model=model or cfg.openai_model or DEFAULT_OPENAI_MODEL,
        )
    elif provider in {"anthropic", "claude"}:
        anthropic_key = api_key or os.environ.get("ANTHROPIC_API_KEY") or cfg.anthropic_api_key
        if not anthropic_key:
            raise ValueError("Paste an Anthropic API key or set ANTHROPIC_API_KEY.")
        compiled = compile_automation_with_anthropic(
            instruction=instruction,
            class_names=class_names,
            api_key=anthropic_key,
            model=model or cfg.anthropic_model or DEFAULT_ANTHROPIC_MODEL,
        )
    else:
        raise ValueError("Provider must be replicate, openai, or anthropic.")
    document = compiled.document
    if append and existing_rules_text.strip():
        existing = load_automation_text(existing_rules_text)
        document = _merge_documents(existing, compiled.document)
    save_automations(document)
    return {
        "rules_text": document.model_dump_json(by_alias=True, indent=2),
        "raw_text": compiled.raw_text,
        "rule_count": len(document.rules),
    }


@app.api(name="validate_rules")
def validate_rules(rules_text: str) -> dict:
    """Validate JSON/YAML rules; return a structured summary or the error."""
    try:
        document = load_automation_text(rules_text)
    except Exception as exc:  # noqa: BLE001 — surface any validation error to the UI
        return {"ok": False, "error": str(exc), "rules": [], "document": None}
    return {
        "ok": True,
        "error": None,
        "document": document.model_dump(mode="json", by_alias=True),
        "rules": [
            {
                "name": r.name,
                "enabled": r.gate.enabled,
                "trigger": r.trigger.on,
                "labels": rule_labels(r),
                "conditions": len(r.when.all_conditions) + len(r.when.any_conditions),
                "actions": [_action_dict(a) for a in _rule_actions(r)],
            }
            for r in document.rules
        ],
    }


@app.api(name="save_rules")
def save_rules(rules_text: str) -> dict:
    document = load_automation_text(rules_text)
    save_automations(document)
    return {"ok": True, "rule_count": len(document.rules)}


@app.api(name="set_rule_enabled")
def set_rule_enabled(rules_text: str, rule_name: str, enabled: bool) -> dict:
    document = load_automation_text(rules_text)
    found = False
    for rule in document.rules:
        if rule.name == rule_name:
            rule.gate.enabled = enabled
            found = True
            break
    if not found:
        raise ValueError(f"Rule not found: {rule_name}")
    save_automations(document)
    return {
        "ok": True,
        "rules_text": document.model_dump_json(by_alias=True, indent=2),
        "rule_count": len(document.rules),
    }


@app.api(name="delete_rule")
def delete_rule(rules_text: str, rule_name: str) -> dict:
    document = load_automation_text(rules_text)
    remaining = [rule for rule in document.rules if rule.name != rule_name]
    if len(remaining) == len(document.rules):
        raise ValueError(f"Rule not found: {rule_name}")
    updated = AutomationDocument(rules=remaining)
    save_automations(updated)
    return {
        "ok": True,
        "rules_text": updated.model_dump_json(by_alias=True, indent=2),
        "rule_count": len(updated.rules),
    }


@app.api(name="load_rules")
def load_rules() -> dict:
    document = load_saved_automations()
    if document is None:
        return {"rules_text": None}
    return {"rules_text": document.model_dump_json(by_alias=True, indent=2)}


@app.api(name="get_config")
def get_config() -> dict:
    cfg = load_local_config()
    return cfg.model_dump(exclude={"replicate_api_token", "openai_api_key", "anthropic_api_key"})


def _merge_documents(existing: AutomationDocument, compiled: AutomationDocument) -> AutomationDocument:
    by_name = {rule.name: rule for rule in existing.rules}
    order = [rule.name for rule in existing.rules]
    for rule in compiled.rules:
        if rule.name not in by_name:
            order.append(rule.name)
        by_name[rule.name] = rule
    return AutomationDocument(rules=[by_name[name] for name in order])


def _rule_actions(rule: AutomationRule) -> list[ActionSpec]:
    if isinstance(rule.then, list):
        return rule.then
    return [*rule.then.enter, *rule.then.exit, *rule.then.while_actions]


def _action_dict(action: ActionSpec) -> dict[str, str]:
    return {"type": action.type, "name": action.name}


def _merge_class_names(*groups: list[str]) -> list[str]:
    seen: set[str] = set()
    merged: list[str] = []
    for group in groups:
        for label in group:
            normalized = " ".join(label.strip().split())
            key = normalized.lower()
            if normalized and key not in seen:
                seen.add(key)
                merged.append(normalized)
    return merged


def _document_uses_moving(document: AutomationDocument) -> bool:
    for rule in document.rules:
        for condition in [*rule.when.all_conditions, *rule.when.any_conditions]:
            if condition.moving is not None:
                return True
    return False


# ── static frontend + media (custom routes take priority over gradio's) ──────
@app.get("/", response_class=HTMLResponse)
def index() -> Any:
    index_html = DIST / "index.html"
    if not index_html.exists():
        return HTMLResponse(
            "<h1>Frontend not built</h1><p>Run <code>pnpm --dir frontend build</code>.</p>",
            status_code=503,
        )
    return FileResponse(index_html)


@app.get("/assets/{file_path:path}")
def assets(file_path: str) -> Any:
    target = (ASSETS / file_path).resolve()
    if not str(target).startswith(str(ASSETS.resolve())) or not target.is_file():
        return JSONResponse({"error": "not found"}, status_code=404)
    return FileResponse(target)


@app.get("/media/{file_path:path}")
def media(file_path: str) -> Any:
    target = (RENDERS / file_path).resolve()
    if not str(target).startswith(str(RENDERS.resolve())) or not target.is_file():
        return JSONResponse({"error": "not found"}, status_code=404)
    return FileResponse(target, media_type="video/mp4")


@app.get("/demo-video")
def demo_video() -> Any:
    path = ROOT / "tiny-trigger-demo.mp4"
    if not path.is_file():
        return JSONResponse({"error": "demo video not found"}, status_code=404)
    return FileResponse(path, media_type="video/mp4")


if __name__ == "__main__":
    app.launch(
        server_name=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
        show_error=True,
    )