#!/opt/homebrew/bin/python3.11
"""
E4B per-CUT Multimodal Processor v1.11

For each CUT scene:
  1. Extract audio segment (ffmpeg)
  2. Build face grid composite from face detections
  3. Query OCR + YOLO data from DB
  4. Send multimodal prompt to E4B (vMLX, port 8000)
  5. Parse structured JSON response
  6. Output {file_uuid}.e4b_per_cut.json

Output:
{
  "file_uuid": "...",
  "scenes": [
    {
      "scene_index": 0,
      "start_time": 0.0,
      "end_time": 15.0,
      "corrected_text": "...",
      "summary": "...",
      "scene_type": "interview",
      "scene_description": "...",
      "face_qa": [{"trace_id": 1, "name": "Person A", "expression": "neutral"}],
      "key_objects": ["microphone", "table"],
      "original_asr": "..."
    }
  ]
}
"""

import json, os, sys, argparse, time, base64, io, subprocess, math, uuid
from typing import Optional
import requests
import psycopg2
import cv2
import numpy as np

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
OUTPUT_DIR = os.environ.get("MOMENTRY_OUTPUT_DIR", "/Users/accusys/momentry/output_dev")
LLM_URL = os.environ.get("E4B_URL",
    os.environ.get("MOMENTRY_LLM_URL", "http://127.0.0.1:8000/v1/chat/completions"))
LLM_MODEL = os.environ.get("E4B_MODEL", "gemma-4-E4B")
DB_URL = os.environ.get("DATABASE_URL", "postgresql://accusys@localhost:5432/momentry")
SCHEMA = os.environ.get("MOMENTRY_DB_SCHEMA", "dev")
FACE_JSON_DIR = OUTPUT_DIR


def extract_audio_segment(video_path: str, start_sec: float, end_sec: float) -> Optional[bytes]:
    """Extract audio segment as WAV bytes"""
    duration = end_sec - start_sec
    if duration <= 0:
        return None
    try:
        result = subprocess.run(
            ["ffmpeg", "-y", "-ss", str(start_sec), "-t", str(duration),
             "-i", video_path, "-vn", "-acodec", "pcm_s16le",
             "-ar", "16000", "-ac", "1", "-f", "wav", "pipe:1"],
            capture_output=True, timeout=min(duration + 30, 300),
        )
        if result.returncode == 0 and len(result.stdout) > 100:
            return result.stdout
    except Exception as e:
        print(f"  ⚠️ Audio extract error: {e}")
    return None


def build_face_grid(face_json_path: str, start_sec: float, end_sec: float,
                    max_frames: int = 4, grid_size: int = 200) -> Optional[str]:
    """Build a composite face grid image from face detections in time range"""
    if not os.path.exists(face_json_path):
        return None
    try:
        with open(face_json_path) as f:
            face_data = json.load(f)
    except Exception:
        return None

    frames_in_range = []
    for fr in face_data.get("frames", []):
        ts = fr.get("timestamp", 0)
        if start_sec <= ts <= end_sec and fr.get("faces"):
            frames_in_range.append(fr)

    if not frames_in_range:
        return None

    # Sample up to max_frames evenly
    step = max(1, len(frames_in_range) // max_frames)
    sampled = frames_in_range[::step][:max_frames]

    # Get video path
    meta = face_data.get("metadata", {})
    video_path = meta.get("video_path", "")

    cells = []
    for fr in sampled:
        ts = fr.get("timestamp", 0)
        face = fr["faces"][0]
        x, y, w, h = face.get("x", 0), face.get("y", 0), \
                     face.get("width", 50), face.get("height", 50)

        # Extract face crop from video
        cmd = [
            "ffmpeg", "-y", "-ss", str(ts), "-i", video_path,
            "-vf", f"crop={w}:{h}:{x}:{y},scale={grid_size}:{grid_size}",
            "-vframes", "1", "-f", "image2pipe", "-",
        ]
        try:
            r = subprocess.run(cmd, capture_output=True, timeout=30)
            if r.returncode == 0 and len(r.stdout) > 100:
                cells.append(r.stdout)
        except Exception:
            continue

    if not cells:
        return None

    # Composite into a grid
    n = len(cells)
    cols = min(2, n)
    rows = math.ceil(n / cols)
    cell_imgs = []
    for raw in cells:
        arr = np.frombuffer(raw, dtype=np.uint8)
        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
        if img is not None:
            img = cv2.resize(img, (grid_size, grid_size))
            cell_imgs.append(img)
        else:
            cell_imgs.append(np.zeros((grid_size, grid_size, 3), dtype=np.uint8))

    canvas = np.zeros((rows * grid_size, cols * grid_size, 3), dtype=np.uint8)
    for i, img in enumerate(cell_imgs):
        r, c = i // cols, i % cols
        canvas[r * grid_size:(r + 1) * grid_size, c * grid_size:(c + 1) * grid_size] = img

    _, buf = cv2.imencode(".jpg", canvas, [cv2.IMWRITE_JPEG_QUALITY, 85])
    return base64.b64encode(buf).decode("utf-8")


def query_ocr_text(file_uuid: str, start_time: float, end_time: float) -> str:
    """Query OCR text from pre_chunks for this time range"""
    try:
        conn = psycopg2.connect(DB_URL)
        cur = conn.cursor()
        cur.execute(f"""
            SELECT text_content, data FROM {SCHEMA}.pre_chunks
            WHERE file_uuid = %s AND processor_type = 'ocr'
              AND start_time BETWEEN %s AND %s
            ORDER BY start_time
        """, (file_uuid, start_time, end_time))
        texts = []
        for text_content, data in cur.fetchall():
            if text_content:
                texts.append(text_content)
            elif data and isinstance(data, dict) and "texts" in data:
                for t in data["texts"]:
                    if isinstance(t, dict) and t.get("text"):
                        texts.append(t["text"])
        cur.close()
        conn.close()
        return " | ".join(texts) if texts else ""
    except Exception as e:
        print(f"  ⚠️ OCR query failed: {e}")
        return ""


def query_yolo_objects(file_uuid: str, start_time: float, end_time: float) -> list:
    """Query YOLO objects from pre_chunks for this time range"""
    try:
        conn = psycopg2.connect(DB_URL)
        cur = conn.cursor()
        cur.execute(f"""
            SELECT DISTINCT obj->>'class_name', COUNT(*) as cnt FROM {SCHEMA}.pre_chunks,
              jsonb_array_elements(data->'objects') AS obj
            WHERE file_uuid = %s AND processor_type = 'yolo'
              AND start_time BETWEEN %s AND %s
            GROUP BY obj->>'class_name'
            ORDER BY cnt DESC
            LIMIT 20
        """, (file_uuid, start_time, end_time))
        objects = [{"label": row[0], "count": row[1]} for row in cur.fetchall() if row[0]]
        cur.close()
        conn.close()
        return objects
    except Exception as e:
        print(f"  ⚠️ YOLO query failed: {e}")
        return []


def query_asr_text(file_uuid: str, start_time: float, end_time: float) -> str:
    """Get original ASR text for this time range"""
    try:
        conn = psycopg2.connect(DB_URL)
        cur = conn.cursor()
        cur.execute(f"""
            SELECT text_content FROM {SCHEMA}.chunk
            WHERE file_uuid = %s AND chunk_type = 'sentence'
              AND start_time >= %s AND end_time <= %s
            ORDER BY start_time
        """, (file_uuid, start_time, end_time))
        texts = [row[0] for row in cur.fetchall() if row[0]]
        cur.close()
        conn.close()
        return " ".join(texts) if texts else ""
    except Exception as e:
        print(f"  ⚠️ ASR query failed: {e}")
        return ""


def call_e4b_multimodal(prompt_text: str,
                        audio_bytes: Optional[bytes] = None,
                        face_grid_b64: Optional[str] = None,
                        max_retries: int = 3) -> Optional[dict]:
    """Send multimodal request to E4B and parse structured response"""
    content_parts = [{"type": "text", "text": prompt_text}]

    if audio_bytes:
        audio_b64 = base64.b64encode(audio_bytes).decode("utf-8")
        content_parts.append({
            "type": "audio_url",
            "audio_url": {"url": f"data:audio/wav;base64,{audio_b64}"},
        })

    if face_grid_b64:
        content_parts.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{face_grid_b64}"},
        })

    messages = [
        {"role": "system", "content": (
            "You are a precise video scene analyst. Given audio, face images, "
            "OCR text, and detected objects for a single scene:\n"
            "1. Transcribe and CORRECT the spoken audio accurately (fix ASR errors)\n"
            "2. Briefly summarize the scene (who, what, where, tone)\n"
            "3. Classify scene type: interview|presentation|monologue|dialogue|action|other\n"
            "4. Describe the visual scene\n"
            "5. For each visible face, identify the person, expression, and gaze direction\n"
            "6. List key objects visible\n\n"
            "ALWAYS respond with valid JSON ONLY (no markdown wrapping):\n"
            '{"corrected_text": "...", "summary": "...", "scene_type": "...", '
            '"scene_description": "...", "face_qa": [{"trace_id": 0, "name": "...", '
            '"expression": "..."}], "key_objects": ["..."]}'
        )},
        {"role": "user", "content": content_parts},
    ]

    payload = {
        "model": LLM_MODEL,
        "messages": messages,
        "max_tokens": 2048,
        "temperature": 0.2,
    }

    for attempt in range(max_retries):
        try:
            resp = requests.post(LLM_URL, json=payload, timeout=180)
            if resp.status_code == 200:
                content = resp.json()["choices"][0]["message"]["content"].strip()
                # Strip markdown code fences
                if content.startswith("```"):
                    content = content.split("\n", 1)[-1]
                    if "```" in content:
                        content = content.rsplit("```", 1)[0]
                return json.loads(content)
            else:
                print(f"  ⚠️ E4B HTTP {resp.status_code}: {resp.text[:200]}")
        except json.JSONDecodeError as e:
            print(f"  ⚠️ JSON parse error (attempt {attempt+1}): {e}")
        except Exception as e:
            print(f"  ⚠️ E4B call error (attempt {attempt+1}): {e}")
        time.sleep(2)

    return None


def load_scenes(file_uuid: str) -> list:
    """Load CUT scenes from JSON"""
    cut_path = os.path.join(OUTPUT_DIR, f"{file_uuid}.cut.json")
    if not os.path.exists(cut_path):
        print(f"[E4B] ❌ cut.json not found: {cut_path}")
        return []
    with open(cut_path) as f:
        data = json.load(f)
    return data.get("scenes", [])


def main():
    parser = argparse.ArgumentParser(description="E4B per-CUT Multimodal Processor v1.11")
    parser.add_argument("--file-uuid", required=True)
    parser.add_argument("--video-path", help="Video file path (auto-detect if not set)")
    parser.add_argument("--max-scenes", type=int, default=999)
    parser.add_argument("--no-audio", action="store_true", help="Skip audio extraction")
    parser.add_argument("--no-faces", action="store_true", help="Skip face grid")
    args = parser.parse_args()

    file_uuid = args.file_uuid
    print(f"[E4B] Processing {file_uuid}")

    # Find video path
    video_path = args.video_path
    if not video_path:
        try:
            conn = psycopg2.connect(DB_URL)
            cur = conn.cursor()
            schema_prefix = f"{SCHEMA}." if SCHEMA else ""
            cur.execute(
                f"SELECT file_path FROM {schema_prefix}videos WHERE file_uuid = %s",
                (file_uuid,),
            )
            row = cur.fetchone()
            cur.close()
            conn.close()
            if row and row[0] and os.path.exists(row[0]):
                video_path = row[0]
        except Exception as e:
            print(f"  ⚠️ DB lookup for video path failed: {e}")
        if not video_path:
            candidates = [
                f"/Users/accusys/momentry/var/sftpgo/data/demo/{file_uuid}.mp4",
                f"/Users/accusys/momentry/output/{file_uuid}.mp4",
                f"/Users/accusys/momentry/output_dev/{file_uuid}.mp4",
            ]
            for c in candidates:
                if os.path.exists(c):
                    video_path = c
                    break
    if not video_path or not os.path.exists(video_path):
        print(f"[E4B] ❌ Video not found: {video_path or args.video_path}")
        sys.exit(1)

    # Load scenes
    scenes = load_scenes(file_uuid)
    if not scenes:
        print("[E4B] ❌ No CUT scenes found")
        sys.exit(1)
    scenes = scenes[:args.max_scenes]
    print(f"[E4B] {len(scenes)} scenes to process")

    # Face JSON path
    face_json_path = os.path.join(FACE_JSON_DIR, f"{file_uuid}.face.json")

    results = []
    t0 = time.time()

    for i, scene in enumerate(scenes):
        st = scene.get("start_time", 0)
        et = scene.get("end_time", st + 30)
        print(f"[E4B] Scene {i}: {st:.1f}s - {et:.1f}s")

        # Gather context
        audio_bytes = None if args.no_audio else extract_audio_segment(video_path, st, et)
        face_grid = None if args.no_faces else build_face_grid(face_json_path, st, et)
        ocr_text = query_ocr_text(file_uuid, st, et)
        yolo_objs = query_yolo_objects(file_uuid, st, et)
        original_asr = query_asr_text(file_uuid, st, et)

        # Build prompt with context
        context_parts = []
        if original_asr:
            context_parts.append(f"Original transcript: \"{original_asr}\"")
        if ocr_text:
            context_parts.append(f"On-screen text: \"{ocr_text}\"")
        if yolo_objs:
            labels = [o["label"] for o in yolo_objs[:10]]
            context_parts.append(f"Detected objects: {', '.join(labels)}")
        context_parts.append(
            "Correct the transcript if it has errors. "
            "Describe the scene and identify people visible."
        )
        prompt_text = "\n".join(context_parts)

        # Call E4B
        e4b_resp = call_e4b_multimodal(prompt_text, audio_bytes, face_grid)

        scene_result = {
            "scene_index": i,
            "start_time": st,
            "end_time": et,
            "original_asr": original_asr or "",
            "ocr_text": ocr_text or "",
            "yolo_objects": yolo_objs,
            "corrected_text": "",
            "summary": "",
            "scene_type": "unknown",
            "scene_description": "",
            "face_qa": [],
            "key_objects": [],
        }

        if e4b_resp:
            scene_result["corrected_text"] = e4b_resp.get("corrected_text", "")
            scene_result["summary"] = e4b_resp.get("summary", "")
            scene_result["scene_type"] = e4b_resp.get("scene_type", "unknown")
            scene_result["scene_description"] = e4b_resp.get("scene_description", "")
            scene_result["face_qa"] = e4b_resp.get("face_qa", [])
            scene_result["key_objects"] = e4b_resp.get("key_objects", [])
            print(f"  ✅ Type: {scene_result['scene_type']}, faces: {len(scene_result['face_qa'])}")
        else:
            print(f"  ⚠️ E4B returned no result, using original ASR")
            scene_result["corrected_text"] = original_asr

        results.append(scene_result)

        elapsed = time.time() - t0
        print(f"  → {elapsed:.0f}s elapsed")

    # Save output
    output = {
        "file_uuid": file_uuid,
        "video_path": video_path,
        "total_scenes": len(results),
        "scenes": results,
    }
    out_path = os.path.join(OUTPUT_DIR, f"{file_uuid}.e4b_per_cut.json")
    with open(out_path, "w") as f:
        json.dump(output, f, indent=2, ensure_ascii=False)

    elapsed = time.time() - t0
    print(f"[E4B] ✅ Done: {len(results)} scenes, {elapsed:.0f}s → {out_path}")


if __name__ == "__main__":
    main()
