#!/usr/bin/env python3
# -*- coding: utf-8 -*-
r"""
HWS local AI service for ai_video_inspector_demo.

Step 1 behavior:
- Mock mode works without any AI package.
- YOLO mode works after installing ultralytics and pillow.

Recommended Windows environment:
    py -3 -m venv .venv
    .venv\Scripts\python.exe -m pip install --upgrade pip
    .venv\Scripts\python.exe local_ai_service.py --host 127.0.0.1 --port 8765 --model mock

Optional YOLO:
    .venv\Scripts\python.exe -m pip install ultralytics pillow
    .venv\Scripts\python.exe local_ai_service.py --model yolo11n.pt
"""

import sys

if sys.version_info[0] < 3:
    print("[ERROR] HWS local AI service requires Python 3.8+.")
    print("[ERROR] Your 'python' command is using Python {}.{}.".format(
        sys.version_info[0], sys.version_info[1]))
    print("[FIX] Install Python 3.10/3.11 x64, then run with: py -3 local_ai_service.py --model mock")
    sys.exit(1)

import argparse
import asyncio
import io
import json
import os
import struct
import threading
import time
from email.parser import BytesParser
from email.policy import default as email_policy
from http.server import BaseHTTPRequestHandler, HTTPServer

try:
    import websockets  # type: ignore
except Exception:
    websockets = None  # type: ignore

try:
    from ultralytics import YOLO  # type: ignore
except Exception:
    YOLO = None  # type: ignore

try:
    from PIL import Image  # type: ignore
except Exception:
    Image = None  # type: ignore

try:
    import numpy as np  # type: ignore
except Exception:
    np = None  # type: ignore

try:
    from faster_whisper import WhisperModel  # type: ignore
except Exception:
    WhisperModel = None  # type: ignore

_MODEL = None
_MODEL_NAME = "mock"
_CONF_THRESHOLD = 0.25
_CLASS_FILTER = set()
_EVENTS_ENABLED = False
_EVENT_STATE = {"counts": {}}
_AUDIO_STATE = {"chunks": {}, "active": {}, "sessions": {}, "streams": {}}

_ASR_ENABLED = False
_ASR_MODEL = None
_ASR_MODEL_NAME = "mock"
_ASR_DEVICE = "cpu"
_ASR_COMPUTE_TYPE = "int8"
_ASR_LANGUAGE = "auto"
_ASR_TASK = "transcribe"
_ASR_WINDOW_SECONDS = 3.0
_ASR_MAX_BUFFER_SECONDS = 12.0
_ASR_MIN_LEVEL = 0.006
_ASR_SAMPLE_RATE = 16000
_TRANSLATE_ENABLED = False
_TRANSLATE_DST_LANG = "en"
_TRANSLATOR_ENGINE = "mock"
_TRANSLATOR_MODEL_NAME = "Helsinki-NLP/opus-mt-en-zh"
_TRANSLATOR_PIPE = None
_TRANSLATOR_REQUIRED = False
_TRANSLATOR_LOCAL_ONLY = False
_ASR_LOCAL_ONLY = False
_AUDIO_DIAG = False


def _json_response(handler, status, payload):
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json; charset=utf-8")
    handler.send_header("Content-Length", str(len(data)))
    handler.end_headers()
    handler.wfile.write(data)


def _parse_multipart_form(content_type, body):
    """Parse multipart/form-data without the deprecated cgi module.

    Returns (fields, files), where fields[name] is text and files[name] is bytes.
    """
    raw = (
        b"Content-Type: " + content_type.encode("utf-8", "ignore") +
        b"\r\nMIME-Version: 1.0\r\n\r\n" + body
    )
    msg = BytesParser(policy=email_policy).parsebytes(raw)
    fields = {}
    files = {}
    if not msg.is_multipart():
        return fields, files

    for part in msg.iter_parts():
        if part.get_content_disposition() != "form-data":
            continue
        name = part.get_param("name", header="content-disposition")
        if not name:
            continue
        filename = part.get_filename()
        payload = part.get_payload(decode=True) or b""
        if filename:
            files[name] = payload
        else:
            charset = part.get_content_charset() or "utf-8"
            fields[name] = payload.decode(charset, "replace")
    return fields, files


def _load_yolo(model_name):
    global _MODEL, _MODEL_NAME
    _MODEL_NAME = model_name
    if not model_name or model_name.lower() in ("mock", "mock-box"):
        _MODEL = None
        _MODEL_NAME = model_name.lower() if model_name else "mock"
        print("[AI] {} mode enabled. No model package is required.".format(_MODEL_NAME))
        return
    if YOLO is None:
        print("[AI] ultralytics is not installed. Fallback to mock mode.")
        print("[AI] To enable YOLO: .venv\\Scripts\\python.exe -m pip install ultralytics pillow")
        _MODEL = None
        _MODEL_NAME = "mock"
        return
    print("[AI] loading YOLO model: {}".format(model_name))
    _MODEL = YOLO(model_name)
    print("[AI] YOLO model loaded.")


def _label_allowed(label):
    if not _CLASS_FILTER:
        return True
    return str(label).lower() in _CLASS_FILTER


def _run_yolo(image_bytes):
    if _MODEL is None:
        if _MODEL_NAME == "mock-box":
            label = "mock-object"
            if not _label_allowed(label):
                return []
            return [{
                "class_id": 0,
                "label": label,
                "confidence": 0.99,
                "x": 120.0,
                "y": 120.0,
                "w": 260.0,
                "h": 180.0,
            }]
        return []
    if Image is None:
        return []

    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    results = _MODEL.predict(image, verbose=False)
    detections = []

    if not results:
        return detections

    r0 = results[0]
    names = getattr(r0, "names", {}) or {}
    boxes = getattr(r0, "boxes", None)
    if boxes is None:
        return detections

    for box in boxes:
        try:
            xyxy = box.xyxy[0].tolist()
            cls_id = int(box.cls[0].item())
            conf = float(box.conf[0].item())
            label = str(names.get(cls_id, cls_id))
            if conf < _CONF_THRESHOLD or not _label_allowed(label):
                continue
            x1, y1, x2, y2 = xyxy
            detections.append({
                "class_id": cls_id,
                "label": label,
                "confidence": round(conf, 4),
                "x": round(float(x1), 2),
                "y": round(float(y1), 2),
                "w": round(float(x2 - x1), 2),
                "h": round(float(y2 - y1), 2),
            })
        except Exception:
            continue

    return detections


def _build_summary(mode, detections, image_size):
    if detections:
        top = detections[:5]
        parts = ["{} {:.2f}".format(d["label"], d["confidence"]) for d in top]
        return "YOLO检测到: " + ", ".join(parts)
    if _MODEL is not None:
        return "YOLO未检测到明显目标。"
    return "Mock AI正常收到视频帧，JPEG大小={} bytes。后续可接入YOLO模型。".format(image_size)


def _count_by_label(detections):
    counts = {}
    for d in detections:
        label = str(d.get("label", "object"))
        counts[label] = counts.get(label, 0) + 1
    return counts


def _build_events(slot, timestamp100ns, frame_index, detections):
    """Generate simple engineering events from per-class count changes.

    This is intentionally detection-based only. It does not assign stable track IDs
    and does not generate trajectories. Tracking is planned for the next demo step.
    """
    if not _EVENTS_ENABLED:
        return []

    prev = dict(_EVENT_STATE.get("counts", {}))
    curr = _count_by_label(detections)
    events = []

    all_labels = sorted(set(prev.keys()) | set(curr.keys()))
    for label in all_labels:
        old_count = int(prev.get(label, 0))
        new_count = int(curr.get(label, 0))
        if old_count == new_count:
            continue

        if new_count > old_count:
            event_type = "object_detected"
            level = "info"
            message = "{} appeared: {} -> {}".format(label, old_count, new_count)
        else:
            event_type = "object_lost"
            level = "warning"
            message = "{} disappeared: {} -> {}".format(label, old_count, new_count)

        events.append({
            "event_id": "{}_{}_{}".format(event_type, label, frame_index),
            "type": event_type,
            "level": level,
            "slot": slot,
            "timestamp100ns": timestamp100ns,
            "frameIndex": frame_index,
            "label": label,
            "prev_count": old_count,
            "current_count": new_count,
            "message": message,
        })

        events.append({
            "event_id": "class_count_changed_{}_{}".format(label, frame_index),
            "type": "class_count_changed",
            "level": "info",
            "slot": slot,
            "timestamp100ns": timestamp100ns,
            "frameIndex": frame_index,
            "label": label,
            "prev_count": old_count,
            "current_count": new_count,
            "message": "{} count changed: {} -> {}".format(label, old_count, new_count),
        })

    _EVENT_STATE["counts"] = curr
    return events




def _load_asr_model(enabled, model_name, device, compute_type, language, task, window_seconds, min_level, local_only=False):
    """Load faster-whisper for Step7 ASR.

    The audio wire protocol already exists from Step6. This function only chooses
    whether audio.pcm_chunk is handled by mock logic or by real ASR.
    """
    global _ASR_ENABLED, _ASR_MODEL, _ASR_MODEL_NAME, _ASR_DEVICE
    global _ASR_COMPUTE_TYPE, _ASR_LANGUAGE, _ASR_TASK
    global _ASR_WINDOW_SECONDS, _ASR_MIN_LEVEL, _ASR_LOCAL_ONLY

    _ASR_ENABLED = bool(enabled)
    _ASR_MODEL = None
    _ASR_MODEL_NAME = model_name or "base"
    _ASR_DEVICE = device or "cpu"
    _ASR_COMPUTE_TYPE = compute_type or "int8"
    _ASR_LANGUAGE = language or "auto"
    _ASR_TASK = task or "transcribe"
    _ASR_WINDOW_SECONDS = max(1.0, float(window_seconds or 3.0))
    _ASR_MIN_LEVEL = max(0.0, float(min_level or 0.006))
    _ASR_LOCAL_ONLY = bool(local_only)

    if not _ASR_ENABLED:
        print("[AI] audio ASR: mock mode. Use --asr to enable faster-whisper.")
        return

    if WhisperModel is None or np is None:
        print("[AI][WARN] faster-whisper/numpy not installed. Audio ASR falls back to mock.")
        print("[AI][FIX] Run: install_customer_demo_deps.bat")
        _ASR_ENABLED = False
        return

    if _ASR_LOCAL_ONLY:
        if not os.path.isdir(_ASR_MODEL_NAME):
            print("[AI][ERROR] ASR local model directory not found: {}".format(_ASR_MODEL_NAME))
            print("[AI][FIX] Run download_customer_demo_models.bat once on a networked PC, then start again.")
            raise SystemExit(2)

    print("[AI] loading faster-whisper ASR model: {} device={} compute_type={} task={} language={} local_only={}".format(
        _ASR_MODEL_NAME, _ASR_DEVICE, _ASR_COMPUTE_TYPE, _ASR_TASK, _ASR_LANGUAGE, _ASR_LOCAL_ONLY
    ))
    try:
        _ASR_MODEL = WhisperModel(_ASR_MODEL_NAME, device=_ASR_DEVICE, compute_type=_ASR_COMPUTE_TYPE, local_files_only=_ASR_LOCAL_ONLY)
    except TypeError:
        # Older faster-whisper does not expose local_files_only.  When a local
        # directory is passed, it still loads from disk and does not contact HF.
        _ASR_MODEL = WhisperModel(_ASR_MODEL_NAME, device=_ASR_DEVICE, compute_type=_ASR_COMPUTE_TYPE)
    print("[AI] faster-whisper ASR model loaded.")


def _pcm16_to_float32(pcm_bytes):
    if np is None or not pcm_bytes:
        return None
    samples = np.frombuffer(pcm_bytes, dtype=np.int16)
    if samples.size <= 0:
        return None
    return samples.astype(np.float32) / 32768.0


def _run_faster_whisper_asr(pcm_bytes, sample_rate):
    if not _ASR_ENABLED or _ASR_MODEL is None or np is None:
        return "", "", "", 0.0
    if sample_rate != _ASR_SAMPLE_RATE:
        # The C++ Audio AI Input Port is expected to send mono 16 kHz.
        # Keep service strict to avoid hidden timestamp/speed errors.
        return "", "", "", 0.0

    audio = _pcm16_to_float32(pcm_bytes)
    if audio is None:
        return "", "", "", 0.0

    language = None if _ASR_LANGUAGE.lower() in ("", "auto", "none") else _ASR_LANGUAGE
    task = _ASR_TASK if _ASR_TASK in ("transcribe", "translate") else "transcribe"

    segments, info = _ASR_MODEL.transcribe(
        audio,
        language=language,
        task=task,
        beam_size=1,
        vad_filter=True,
        condition_on_previous_text=False,
    )
    parts = []
    for seg in segments:
        t = getattr(seg, "text", "")
        if t:
            parts.append(t.strip())
    text = " ".join([x for x in parts if x]).strip()
    detected_lang = getattr(info, "language", "") or ""
    lang_prob = float(getattr(info, "language_probability", 0.0) or 0.0)

    translation = ""
    if task == "translate":
        # Whisper translate task returns English text directly.
        translation = text
    return text, translation, detected_lang, lang_prob



def _load_translator(engine, model_name, required=False, local_only=False):
    """Load optional Step9 text translator.

    The demo is designed so that video/ASR does not depend on translation packages.
    If --translator hf is requested and transformers/torch/model are available,
    English text is translated to Chinese locally. Otherwise we keep a clear mock
    translation fallback so the SDK subtitles and output-play path can still be tested.
    """
    global _TRANSLATOR_ENGINE, _TRANSLATOR_MODEL_NAME, _TRANSLATOR_PIPE, _TRANSLATOR_REQUIRED, _TRANSLATOR_LOCAL_ONLY
    _TRANSLATOR_REQUIRED = bool(required)
    _TRANSLATOR_LOCAL_ONLY = bool(local_only)
    _TRANSLATOR_ENGINE = (engine or "mock").lower()
    _TRANSLATOR_MODEL_NAME = model_name or "Helsinki-NLP/opus-mt-en-zh"
    _TRANSLATOR_PIPE = None

    if _TRANSLATOR_ENGINE in ("", "none", "off"):
        print("[AI] translator: disabled")
        return
    if _TRANSLATOR_ENGINE == "mock":
        if _TRANSLATOR_REQUIRED:
            print("[AI][ERROR] real translator is required but --translator mock was selected.")
            raise SystemExit(2)
        print("[AI] translator: mock mode")
        return
    if _TRANSLATOR_ENGINE != "hf":
        print("[AI][WARN] unknown translator engine '{}'; fallback to mock".format(_TRANSLATOR_ENGINE))
        _TRANSLATOR_ENGINE = "mock"
        return

    try:
        # Do not use transformers.pipeline("translation") here. Some recent or
        # minimal Transformers builds no longer expose the legacy generic
        # "translation" task name, which caused startup failure even though the
        # MarianMT model itself is usable.  Load the seq2seq model directly so the
        # demo is stable across Transformers versions.
        from transformers import AutoTokenizer, AutoModelForSeq2SeqLM  # type: ignore
        if _TRANSLATOR_LOCAL_ONLY and not os.path.isdir(_TRANSLATOR_MODEL_NAME):
            raise FileNotFoundError("translator local model directory not found: {}".format(_TRANSLATOR_MODEL_NAME))
        print("[AI] loading HF translator model directly: {} local_only={}".format(_TRANSLATOR_MODEL_NAME, _TRANSLATOR_LOCAL_ONLY))
        tokenizer = AutoTokenizer.from_pretrained(_TRANSLATOR_MODEL_NAME, local_files_only=_TRANSLATOR_LOCAL_ONLY)
        model = AutoModelForSeq2SeqLM.from_pretrained(_TRANSLATOR_MODEL_NAME, local_files_only=_TRANSLATOR_LOCAL_ONLY)
        model.eval()
        _TRANSLATOR_PIPE = {
            "kind": "seq2seq",
            "tokenizer": tokenizer,
            "model": model,
        }
        print("[AI] HF translator loaded.")
    except Exception as exc:
        print("[AI][ERROR] HF translator load failed: {}".format(exc))
        print("[AI][FIX] Run download_customer_demo_models.bat once, or pass a valid local --translator-model directory: {}".format(_TRANSLATOR_MODEL_NAME))
        _TRANSLATOR_PIPE = None
        if _TRANSLATOR_REQUIRED:
            raise SystemExit(2)
        _TRANSLATOR_ENGINE = "mock"



def _contains_cjk_text(s):
    if not s:
        return False
    for ch in s:
        o = ord(ch)
        if (0x4E00 <= o <= 0x9FFF) or (0x3400 <= o <= 0x4DBF) or (0xF900 <= o <= 0xFAFF):
            return True
    return False

def _translate_text(text, src_lang, dst_lang):
    """Translate recognized text for Step9 ASR demo.

    Primary engineering target: English speech -> Chinese subtitle.
    - --translator hf uses a local HuggingFace translation model, defaulting to
      Helsinki-NLP/opus-mt-en-zh.
    - --translator mock keeps the demo usable without heavy translation packages.
    """
    if not text:
        return ""
    dst = (dst_lang or _TRANSLATE_DST_LANG or "zh").lower()
    src = (src_lang or "auto").lower()

    if _ASR_TASK == "translate" and dst.startswith("en"):
        return text
    if not _TRANSLATE_ENABLED:
        return ""

    if _TRANSLATOR_ENGINE == "hf" and _TRANSLATOR_PIPE is not None:
        try:
            # Direct seq2seq inference for Helsinki-NLP/opus-mt-en-zh.
            if isinstance(_TRANSLATOR_PIPE, dict) and _TRANSLATOR_PIPE.get("kind") == "seq2seq":
                tokenizer = _TRANSLATOR_PIPE.get("tokenizer")
                model = _TRANSLATOR_PIPE.get("model")
                if tokenizer is None or model is None:
                    return "[翻译模型错误] translator object invalid"
                inputs = tokenizer([text], return_tensors="pt", padding=True, truncation=True, max_length=256)
                output_ids = model.generate(**inputs, max_new_tokens=256, num_beams=4)
                decoded = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
                t = (decoded[0] if decoded else "").strip()
            else:
                # Compatibility fallback for older builds where pipeline exists.
                r = _TRANSLATOR_PIPE(text, max_length=256)
                t = ""
                if isinstance(r, list) and r:
                    t = str(r[0].get("translation_text", "") or "").strip()
            if t:
                # In EN->CN demo, never accept an English-only result as the
                # final translation. This prevents Out Play from reading the
                # original English when the wrong model or task is selected.
                if dst.startswith("zh") and not _contains_cjk_text(t):
                    return "[翻译未输出中文，请检查模型/脚本] " + t
                return t
        except Exception as exc:
            if dst.startswith("zh"):
                return "[翻译模型错误] {}".format(exc)
            return "[translation error] {}".format(exc)

    # No real translator is available. Do NOT append the original English text here.
    # Otherwise the output player may read English after a Chinese warning prefix.
    # The C++ demo will not speak this diagnostic text in ASR EN->CN mode.
    if dst.startswith("zh"):
        if _TRANSLATOR_REQUIRED:
            return ""
        return "[翻译不可用：请先安装并加载英文转中文翻译模型]"
    return "[translation unavailable: translator not loaded]"

def _pcm16_rms_level(pcm_bytes):
    if not pcm_bytes:
        return 0.0
    sample_count = len(pcm_bytes) // 2
    if sample_count <= 0:
        return 0.0
    # Keep this lightweight. Inspect at most about 1 second of 16 kHz samples.
    max_samples = min(sample_count, 16000)
    total = 0.0
    for i in range(max_samples):
        v = int.from_bytes(pcm_bytes[i * 2:i * 2 + 2], byteorder="little", signed=True)
        total += float(v) * float(v)
    rms = (total / max_samples) ** 0.5
    return rms / 32768.0


def _process_audio_chunk(header, pcm_bytes):
    """Handle Step6/Step7/Step10 audio stream.

    Step10 keeps the service strict about the wire format and adds continuity checks:
    - payload must be PCM16 mono 16 kHz
    - chunkIndex should be continuous
    - timestamp should be continuous within a small tolerance
    The ASR worker path is still simple and synchronous, but now problems are visible in
    both the Python console and the C++ demo log/raw JSON.
    """
    audio_source = int(header.get("audioSource", 0))
    chunk_index = int(header.get("chunkIndex", 0))
    timestamp100ns = int(header.get("timestamp100ns", 0))
    sample_rate = int(header.get("sampleRate", 16000))
    channels = int(header.get("channels", 1))
    bits_per_sample = int(header.get("bitsPerSample", 16))
    sample_count = int(header.get("sampleCount", 0))
    src_lang_req = str(header.get("srcLang", "auto"))
    dst_lang = str(header.get("dstLang", "en"))
    duration100ns = int(sample_count * 10000000 / sample_rate) if sample_rate > 0 else 0
    rms = _pcm16_rms_level(pcm_bytes)

    key = str(audio_source)
    old_chunks = int(_AUDIO_STATE["chunks"].get(key, 0))
    _AUDIO_STATE["chunks"][key] = old_chunks + 1

    expected_bytes = sample_count * 2 if sample_count > 0 else 0
    warnings = []
    if sample_rate != _ASR_SAMPLE_RATE:
        warnings.append("sample_rate={} expected={}".format(sample_rate, _ASR_SAMPLE_RATE))
    if channels != 1:
        warnings.append("channels={} expected=1".format(channels))
    if bits_per_sample != 16:
        warnings.append("bits={} expected=16".format(bits_per_sample))
    if expected_bytes > 0 and len(pcm_bytes) != expected_bytes:
        warnings.append("payload_bytes={} expected={}".format(len(pcm_bytes), expected_bytes))

    streams = _AUDIO_STATE.setdefault("streams", {})
    stream = streams.get(key)
    if stream is None:
        stream = {
            "next_chunk": chunk_index,
            "next_ts": timestamp100ns,
            "missing_chunks": 0,
            "timestamp_gaps": 0,
            "last_recv_wall_ms": 0.0,
            "received_chunks": 0,
            "total_payload_bytes": 0,
            "max_recv_gap_ms": 0.0,
            "max_asr_ms": 0.0,
            "max_translate_ms": 0.0,
        }
        streams[key] = stream

    now_wall_ms = time.time() * 1000.0
    last_recv_wall_ms = float(stream.get("last_recv_wall_ms", 0.0) or 0.0)
    recv_gap_ms = 0.0
    if last_recv_wall_ms > 0.0:
        recv_gap_ms = now_wall_ms - last_recv_wall_ms
        if recv_gap_ms > float(stream.get("max_recv_gap_ms", 0.0) or 0.0):
            stream["max_recv_gap_ms"] = recv_gap_ms
    stream["last_recv_wall_ms"] = now_wall_ms
    stream["received_chunks"] = int(stream.get("received_chunks", 0)) + 1
    stream["total_payload_bytes"] = int(stream.get("total_payload_bytes", 0)) + len(pcm_bytes)

    expected_chunk = int(stream.get("next_chunk", chunk_index))
    if chunk_index != expected_chunk:
        delta = chunk_index - expected_chunk
        if delta > 0:
            stream["missing_chunks"] = int(stream.get("missing_chunks", 0)) + delta
            warnings.append("missing_chunks={} expected_chunk={} got={}".format(delta, expected_chunk, chunk_index))
        else:
            warnings.append("out_of_order_chunk expected={} got={}".format(expected_chunk, chunk_index))
    stream["next_chunk"] = chunk_index + 1

    expected_ts = int(stream.get("next_ts", timestamp100ns))
    if timestamp100ns and expected_ts:
        gap = timestamp100ns - expected_ts
        if abs(gap) > 1200000:  # 120 ms tolerance
            stream["timestamp_gaps"] = int(stream.get("timestamp_gaps", 0)) + 1
            warnings.append("timestamp_gap_ms={:.1f}".format(gap / 10000.0))
    stream["next_ts"] = timestamp100ns + duration100ns

    expected_wall_gap_ms = (duration100ns / 10000.0) if duration100ns > 0 else 0.0
    if recv_gap_ms > 0.0 and expected_wall_gap_ms > 0.0 and recv_gap_ms > (expected_wall_gap_ms * 3.0 + 200.0):
        warnings.append("recv_wall_gap_ms={:.1f} expected_ms={:.1f}".format(recv_gap_ms, expected_wall_gap_ms))

    current_asr_ms = 0.0
    current_translate_ms = 0.0

    if warnings:
        print("[AI][AUDIO][WARN] chunk={} {}".format(chunk_index, "; ".join(warnings)))
    # Customer demo: do not print every normal audio chunk.
    # Keep only warnings by default. Use --audio-diag for periodic RX summaries.
    if warnings or (_AUDIO_DIAG and (chunk_index % 25) == 0):
        print("[AI][AUDIO][RX] chunk={} bytes={} expected={} level={:.4f} recvGapMs={:.1f} totalRx={} missingTotal={} tsGapTotal={} warnings={}".format(
            chunk_index, len(pcm_bytes), expected_bytes, rms, recv_gap_ms,
            int(stream.get("received_chunks", 0)),
            int(stream.get("missing_chunks", 0)),
            int(stream.get("timestamp_gaps", 0)),
            ";".join(warnings) if warnings else "none"))

    def _decorate(payload):
        payload["payloadBytes"] = len(pcm_bytes)
        payload["expectedBytes"] = expected_bytes
        payload["streamWarnings"] = warnings
        payload["missingChunksTotal"] = int(stream.get("missing_chunks", 0))
        payload["timestampGapTotal"] = int(stream.get("timestamp_gaps", 0))
        payload["audioDiag"] = {
            "receivedChunks": int(stream.get("received_chunks", 0)),
            "totalPayloadBytes": int(stream.get("total_payload_bytes", 0)),
            "recvGapMs": round(float(recv_gap_ms), 2),
            "maxRecvGapMs": round(float(stream.get("max_recv_gap_ms", 0.0) or 0.0), 2),
            "missingChunksTotal": int(stream.get("missing_chunks", 0)),
            "timestampGapTotal": int(stream.get("timestamp_gaps", 0)),
            "bufferedSeconds": round(float(stream.get("buffered_seconds", 0.0) or 0.0), 3),
            "asrMs": round(float(current_asr_ms), 2),
            "translateMs": round(float(current_translate_ms), 2),
            "maxAsrMs": round(float(stream.get("max_asr_ms", 0.0) or 0.0), 2),
            "maxTranslateMs": round(float(stream.get("max_translate_ms", 0.0) or 0.0), 2),
        }
        if warnings:
            old = payload.get("summary", "") or ""
            payload["summary"] = "Audio stream warning: {}{}{}".format(
                "; ".join(warnings), " | " if old else "", old)
        return payload

    if not _ASR_ENABLED or _ASR_MODEL is None:
        active = rms > _ASR_MIN_LEVEL
        if active:
            text = "[mock] audio stream active, chunk={}, level={:.3f}".format(chunk_index, rms)
            summary = "Audio mock transcript updated."
        else:
            text = "[mock] listening... chunk={}, level={:.3f}".format(chunk_index, rms)
            summary = "Audio mock stream received silence/low level."

        is_final = (chunk_index % 6) == 5
        result_type = "audio.transcript_final" if is_final else "audio.transcript_partial"
        translation = ""
        if is_final:
            translation = _translate_text(text, src_lang_req, dst_lang)
        return _decorate({
            "type": result_type,
            "status": "ok",
            "audioSource": audio_source,
            "chunkIndex": chunk_index,
            "timestamp100ns": timestamp100ns,
            "startTimestamp100ns": timestamp100ns,
            "endTimestamp100ns": timestamp100ns + duration100ns,
            "sampleRate": sample_rate,
            "channels": channels,
            "bitsPerSample": bits_per_sample,
            "sampleCount": sample_count,
            "srcLang": src_lang_req,
            "dstLang": dst_lang,
            "text": text,
            "translation": translation,
            "isFinal": is_final,
            "level": round(rms, 4),
            "summary": summary,
            "transport": "websocket",
            "engine": "mock",
        })

    session = _AUDIO_STATE["sessions"].get(key)
    if session is None:
        session = {
            "pcm": bytearray(),
            "start_ts": timestamp100ns,
            "last_ts": timestamp100ns,
            "active_seen": False,
            "last_text": "",
            "last_lang": "",
        }
        _AUDIO_STATE["sessions"][key] = session

    if not session["pcm"]:
        session["start_ts"] = timestamp100ns
    session["last_ts"] = timestamp100ns + duration100ns
    session["pcm"].extend(pcm_bytes)
    if rms > _ASR_MIN_LEVEL:
        session["active_seen"] = True

    max_bytes = int(_ASR_MAX_BUFFER_SECONDS * _ASR_SAMPLE_RATE * 2)
    if len(session["pcm"]) > max_bytes:
        session["pcm"] = bytearray(session["pcm"][-max_bytes:])
        session["start_ts"] = max(timestamp100ns - int(_ASR_MAX_BUFFER_SECONDS * 10000000), 0)

    buffered_seconds = len(session["pcm"]) / float(_ASR_SAMPLE_RATE * 2)
    stream["buffered_seconds"] = buffered_seconds
    should_transcribe = buffered_seconds >= _ASR_WINDOW_SECONDS

    if not should_transcribe:
        return _decorate({
            "type": "audio.transcript_partial",
            "status": "ok",
            "audioSource": audio_source,
            "chunkIndex": chunk_index,
            "timestamp100ns": timestamp100ns,
            "startTimestamp100ns": session["start_ts"],
            "endTimestamp100ns": timestamp100ns + duration100ns,
            "sampleRate": sample_rate,
            "channels": channels,
            "bitsPerSample": bits_per_sample,
            "sampleCount": sample_count,
            "srcLang": src_lang_req,
            "dstLang": dst_lang,
            "text": session.get("last_text", "") or "[ASR] listening... level={:.3f}".format(rms),
            "translation": "",
            "isFinal": False,
            "level": round(rms, 4),
            "summary": "ASR buffering {:.1f}/{:.1f}s".format(buffered_seconds, _ASR_WINDOW_SECONDS),
            "transport": "websocket",
            "engine": "faster-whisper",
        })

    pcm_for_asr = bytes(session["pcm"])
    start_ts = int(session["start_ts"])
    end_ts = int(session["last_ts"])

    text = ""
    translation = ""
    detected_lang = ""
    lang_prob = 0.0
    error = ""
    asr_ms = 0.0
    translate_ms = 0.0
    try:
        if session.get("active_seen", False) or _pcm16_rms_level(pcm_for_asr) > _ASR_MIN_LEVEL:
            asr_t0 = time.time()
            text, translation, detected_lang, lang_prob = _run_faster_whisper_asr(pcm_for_asr, sample_rate)
            asr_ms = (time.time() - asr_t0) * 1000.0
            if asr_ms > float(stream.get("max_asr_ms", 0.0) or 0.0):
                stream["max_asr_ms"] = asr_ms
            if not translation:
                tr_t0 = time.time()
                translation = _translate_text(text, detected_lang or src_lang_req, dst_lang)
                translate_ms = (time.time() - tr_t0) * 1000.0
                if translate_ms > float(stream.get("max_translate_ms", 0.0) or 0.0):
                    stream["max_translate_ms"] = translate_ms
    except Exception as exc:
        error = str(exc)

    current_asr_ms = asr_ms
    current_translate_ms = translate_ms

    if _AUDIO_DIAG or asr_ms > 800.0 or translate_ms > 800.0:
        print("[AI][AUDIO][ASR] window={:.2f}s chunk={} asrMs={:.1f} translateMs={:.1f} textLen={} transLen={} missingTotal={} tsGapTotal={}".format(
            (len(pcm_for_asr) / float(_ASR_SAMPLE_RATE * 2)) if _ASR_SAMPLE_RATE else 0.0,
            chunk_index, asr_ms, translate_ms, len(text or ""), len(translation or ""),
            int(stream.get("missing_chunks", 0)), int(stream.get("timestamp_gaps", 0))))

    session["pcm"] = bytearray()
    session["start_ts"] = timestamp100ns + duration100ns
    session["active_seen"] = False
    if text:
        session["last_text"] = text
    if detected_lang:
        session["last_lang"] = detected_lang

    if error:
        return _decorate({
            "type": "audio.transcript_final",
            "status": "error",
            "audioSource": audio_source,
            "chunkIndex": chunk_index,
            "timestamp100ns": timestamp100ns,
            "startTimestamp100ns": start_ts,
            "endTimestamp100ns": end_ts,
            "sampleRate": sample_rate,
            "channels": channels,
            "bitsPerSample": bits_per_sample,
            "sampleCount": sample_count,
            "srcLang": detected_lang or src_lang_req,
            "dstLang": dst_lang,
            "text": "",
            "translation": "",
            "isFinal": True,
            "level": round(rms, 4),
            "summary": "ASR error: {}".format(error),
            "transport": "websocket",
            "engine": "faster-whisper",
        })

    if not text:
        return _decorate({
            "type": "audio.transcript_partial",
            "status": "ok",
            "audioSource": audio_source,
            "chunkIndex": chunk_index,
            "timestamp100ns": timestamp100ns,
            "startTimestamp100ns": start_ts,
            "endTimestamp100ns": end_ts,
            "sampleRate": sample_rate,
            "channels": channels,
            "bitsPerSample": bits_per_sample,
            "sampleCount": sample_count,
            "srcLang": detected_lang or src_lang_req,
            "dstLang": dst_lang,
            "text": "[ASR] no speech detected",
            "translation": "",
            "isFinal": False,
            "level": round(rms, 4),
            "summary": "ASR window processed, no speech detected.",
            "transport": "websocket",
            "engine": "faster-whisper",
        })

    return _decorate({
        "type": "audio.transcript_final",
        "status": "ok",
        "audioSource": audio_source,
        "chunkIndex": chunk_index,
        "timestamp100ns": timestamp100ns,
        "startTimestamp100ns": start_ts,
        "endTimestamp100ns": end_ts,
        "sampleRate": sample_rate,
        "channels": channels,
        "bitsPerSample": bits_per_sample,
        "sampleCount": sample_count,
        "srcLang": detected_lang or src_lang_req,
        "dstLang": dst_lang,
        "text": text,
        "translation": translation,
        "isFinal": True,
        "level": round(rms, 4),
        "languageProbability": round(lang_prob, 4),
        "summary": "ASR final: {}".format(text[:120]),
        "transport": "websocket",
        "engine": "faster-whisper",
        "asrModel": _ASR_MODEL_NAME,
    })

def _process_analyze_frame(slot, timestamp100ns, frame_index, mode, prompt, image_bytes):
    if not image_bytes:
        return {
            "slot": slot,
            "timestamp100ns": timestamp100ns,
            "frameIndex": frame_index,
            "status": "error",
            "summary": "missing image",
            "detections": [],
            "events": [],
        }

    detections = _run_yolo(image_bytes)
    summary = _build_summary(mode, detections, len(image_bytes))
    events = _build_events(slot, timestamp100ns, frame_index, detections)
    if events:
        summary = summary + " | events={}".format(len(events))

    overlays = [{
        "type": "text",
        "x": 20,
        "y": 20,
        "text": "AI: " + summary,
        "duration_ms": 1000,
    }]
    for d in detections[:20]:
        overlays.append({
            "type": "box",
            "x": d["x"],
            "y": d["y"],
            "w": d["w"],
            "h": d["h"],
            "label": "{} {:.2f}".format(d["label"], d["confidence"]),
            "duration_ms": 1000,
        })

    return {
        "slot": slot,
        "timestamp100ns": timestamp100ns,
        "frameIndex": frame_index,
        "status": "ok",
        "mode": mode,
        "model": _MODEL_NAME,
        "transport": "websocket_or_http",
        "summary": summary,
        "detections": detections,
        "events": events,
        "overlays": overlays,
        "prompt_echo": prompt[:120],
    }


async def _ws_handler(websocket, path=None):
    # WebSocket path is checked softly because the Python websockets API changed
    # handler signatures across versions.
    client = getattr(websocket, "remote_address", None)
    print("[WS] client connected: {}".format(client))
    try:
        async for message in websocket:
            try:
                if isinstance(message, str):
                    req = json.loads(message)
                    if req.get("type") == "ping":
                        await websocket.send(json.dumps({"type": "pong", "time": time.time()}, ensure_ascii=False))
                    else:
                        await websocket.send(json.dumps({"status": "ok", "type": "ack", "time": time.time()}, ensure_ascii=False))
                    continue

                if not isinstance(message, (bytes, bytearray)) or len(message) < 4:
                    await websocket.send(json.dumps({"status": "error", "summary": "invalid binary message"}, ensure_ascii=False))
                    continue

                header_size = struct.unpack_from("<I", message, 0)[0]
                if header_size <= 0 or header_size > 1024 * 1024 or 4 + header_size > len(message):
                    await websocket.send(json.dumps({"status": "error", "summary": "invalid header size"}, ensure_ascii=False))
                    continue

                header_bytes = message[4:4 + header_size]
                image_bytes = bytes(message[4 + header_size:])
                header = json.loads(header_bytes.decode("utf-8", "replace"))

                msg_type = header.get("type", "video.jpeg_frame")
                if msg_type == "video.jpeg_frame":
                    slot = int(header.get("slot", 0))
                    timestamp100ns = int(header.get("timestamp100ns", 0))
                    frame_index = int(header.get("frameIndex", 0))
                    mode = str(header.get("mode", "yolo"))
                    prompt = str(header.get("prompt", ""))

                    payload = _process_analyze_frame(slot, timestamp100ns, frame_index, mode, prompt, image_bytes)
                    payload["type"] = "video.detections"
                    payload["transport"] = "websocket"
                    await websocket.send(json.dumps(payload, ensure_ascii=False))
                    continue

                if msg_type == "audio.pcm_chunk":
                    payload = _process_audio_chunk(header, image_bytes)
                    await websocket.send(json.dumps(payload, ensure_ascii=False))
                    continue

                await websocket.send(json.dumps({"status": "error", "summary": "unsupported ws message type: {}".format(msg_type)}, ensure_ascii=False))
            except Exception as exc:
                await websocket.send(json.dumps({
                    "status": "error",
                    "summary": "AI WebSocket exception: {}".format(exc),
                }, ensure_ascii=False))
    finally:
        print("[WS] client disconnected: {}".format(client))


def _run_http_server(host, port):
    httpd = HTTPServer((host, port), AiHandler)
    print("[AI] HTTP service listening on http://{}:{}".format(host, port))
    print("[AI] health: GET /v1/health")
    print("[AI] analyze: POST /v1/analyze_frame")
    httpd.serve_forever()


async def _run_ws_server(host, ws_port):
    if websockets is None:
        raise RuntimeError("websockets is not installed. Run install_ws_deps.bat first.")
    print("[AI] WebSocket service listening on ws://{}:{}/v1/ai/stream".format(host, ws_port))
    async with websockets.serve(_ws_handler, host, ws_port, max_size=32 * 1024 * 1024):
        await asyncio.Future()


class AiHandler(BaseHTTPRequestHandler):
    server_version = "HWSLocalAI/0.1"

    def do_GET(self):  # noqa: N802
        if self.path.startswith("/v1/health"):
            _json_response(self, 200, {
                "status": "ok",
                "service": "hws-local-ai-service",
                "model": _MODEL_NAME,
                "python": sys.version.split()[0],
                "time": time.time(),
            })
            return
        _json_response(self, 404, {"status": "error", "summary": "not found"})

    def do_POST(self):  # noqa: N802
        if not self.path.startswith("/v1/analyze_frame"):
            _json_response(self, 404, {"status": "error", "summary": "not found"})
            return

        try:
            content_type = self.headers.get("Content-Type", "")
            content_length = int(self.headers.get("Content-Length", "0"))
            body = self.rfile.read(content_length)

            fields, files = _parse_multipart_form(content_type, body)

            slot = int(fields.get("slot", "0"))
            timestamp100ns = int(fields.get("timestamp100ns", "0"))
            frame_index = int(fields.get("frameIndex", "0"))
            mode = fields.get("mode", "yolo")
            prompt = fields.get("prompt", "")

            image_bytes = files.get("image", b"")
            if not image_bytes:
                _json_response(self, 400, {"status": "error", "summary": "missing image"})
                return
            payload = _process_analyze_frame(slot, timestamp100ns, frame_index, mode, prompt, image_bytes)
            payload["transport"] = "http"
            _json_response(self, 200, payload)
        except Exception as exc:
            _json_response(self, 500, {
                "status": "error",
                "summary": "AI service exception: {}".format(exc),
            })

    def log_message(self, fmt, *args):
        sys.stdout.write("[HTTP] " + (fmt % args) + "\n")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8765, help="HTTP port for health and fallback HTTP analyze_frame.")
    parser.add_argument("--ws", action="store_true", help="Enable WebSocket streaming transport.")
    parser.add_argument("--ws-port", type=int, default=8766, help="WebSocket port for /v1/ai/stream.")
    parser.add_argument("--model", default="mock", help="mock, mock-box, or yolo11n.pt/yolov8n.pt/etc.")
    parser.add_argument("--events", action="store_true", help="Enable detection event generation.")
    parser.add_argument("--conf", type=float, default=0.25, help="Detection confidence threshold.")
    parser.add_argument("--classes", default="", help="Optional comma-separated class filter, e.g. person,car.")
    parser.add_argument("--asr", action="store_true", help="Enable faster-whisper ASR for audio.pcm_chunk.")
    parser.add_argument("--asr-model", default="base", help="faster-whisper model name/path, e.g. tiny, base, small, medium, large-v3, or models\\faster-whisper-base.")
    parser.add_argument("--asr-local-only", action="store_true", help="Load faster-whisper only from a local model directory. Do not download during service startup.")
    parser.add_argument("--asr-device", default="cpu", help="ASR device: cpu or cuda.")
    parser.add_argument("--asr-compute-type", default="int8", help="ASR compute type, e.g. int8/cpu, float16/cuda.")
    parser.add_argument("--asr-language", default="auto", help="ASR language, e.g. auto, zh, en, ja.")
    parser.add_argument("--asr-task", default="transcribe", choices=["transcribe", "translate"], help="Whisper task. translate outputs English text.")
    parser.add_argument("--asr-window-seconds", type=float, default=3.0, help="Audio seconds accumulated per ASR segment.")
    parser.add_argument("--asr-min-level", type=float, default=0.006, help="RMS threshold for speech activity hint.")
    parser.add_argument("--translate", action="store_true", help="Enable Step9 translation subtitle output. For EN->CN use --translator hf --dst-lang zh.")
    parser.add_argument("--dst-lang", default="en", help="Target subtitle language, e.g. en, zh, ja. Whisper translate supports English output directly.")
    parser.add_argument("--translator", default="mock", choices=["mock", "hf", "none"], help="Step9 text translator engine. Use hf for local HuggingFace English->Chinese translation.")
    parser.add_argument("--translator-model", default="Helsinki-NLP/opus-mt-en-zh", help="HuggingFace translation model/repo or local directory for --translator hf.")
    parser.add_argument("--translator-local-only", action="store_true", help="Load translator only from local model directory. Do not download during service startup.")
    parser.add_argument("--translator-required", action="store_true", help="Fail service startup instead of falling back to mock when the requested translator cannot load.")
    parser.add_argument("--audio-diag", action="store_true", help="Print verbose audio transport/ASR diagnostics for loss and latency analysis.")
    args = parser.parse_args()

    global _CONF_THRESHOLD, _CLASS_FILTER, _EVENTS_ENABLED
    global _TRANSLATE_ENABLED, _TRANSLATE_DST_LANG, _AUDIO_DIAG
    _CONF_THRESHOLD = float(args.conf)
    _EVENTS_ENABLED = bool(args.events)
    _TRANSLATE_ENABLED = bool(args.translate)
    _TRANSLATE_DST_LANG = str(args.dst_lang or "en")
    _AUDIO_DIAG = bool(args.audio_diag)
    _CLASS_FILTER = set([x.strip().lower() for x in args.classes.split(',') if x.strip()])

    _load_yolo(args.model)
    #_load_asr_model(args.asr, args.asr_model, args.asr_device, args.asr_compute_type,
    #                args.asr_language, args.asr_task, args.asr_window_seconds, args.asr_min_level, args.asr_local_only)
    #_load_translator(args.translator, args.translator_model, args.translator_required, args.translator_local_only)
    #print("[AI] events: {}  conf_threshold: {}  class_filter: {}".format(_EVENTS_ENABLED, _CONF_THRESHOLD, sorted(_CLASS_FILTER)))
    #print("[AI] translation: {} dst_lang={} translator={} translatorLocalOnly={}".format(_TRANSLATE_ENABLED, _TRANSLATE_DST_LANG, args.translator, args.translator_local_only))
    #print("[AI] audio_diag: {}".format(_AUDIO_DIAG))

    if args.ws:
        http_thread = threading.Thread(target=_run_http_server, args=(args.host, args.port), daemon=True)
        http_thread.start()
        try:
            asyncio.run(_run_ws_server(args.host, args.ws_port))
        except KeyboardInterrupt:
            print("\n[AI] stopped.")
        return 0

    httpd = HTTPServer((args.host, args.port), AiHandler)
    print("[AI] service listening on http://{}:{}".format(args.host, args.port))
    print("[AI] health: GET /v1/health")
    print("[AI] analyze: POST /v1/analyze_frame")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\n[AI] stopped.")
    return 0


if __name__ == "__main__":
    sys.exit(main())
