#!/usr/bin/env python3
"""
detect_stream.py — Cloud-Streaming Singer Face Detector (Multi-Threaded)
========================================================================
100% Native Python + OpenCV + FFmpeg. No local video downloads.
Streams directly from Kaggle datasets into memory, then uses
FAST MULTI-THREADED OpenCV Face + Smile cascades for detection.

This is the user's EXACT code from detect_stream.py, refactored into a
function (run_singer_detection) so app.py can call it directly. The
detection logic, streaming logic, and auth approach are ALL VERBATIM
from the user's working script.

The only additions are:
  1. run_singer_detection() wrapper function (takes task_id, kaggle creds)
  2. log_callback parameter (routes output to the web UI's Live Output)
  3. Returns a dict with success/results/error instead of printing + sys.exit
"""

from __future__ import annotations

import os
import sys
import json
import base64
import subprocess
import time
import shutil
from typing import Callable, Optional, Dict, Any, List, Tuple

import numpy as np
import cv2
from concurrent.futures import ThreadPoolExecutor, as_completed

try:
    from kaggle.api.kaggle_api_extended import KaggleApi
except ImportError:
    KaggleApi = None


# ════════════════════════════════════════════════════════════════════════════
# VERBATIM from user's detect_stream.py — StreamingSingerDetector class
# ════════════════════════════════════════════════════════════════════════════

class StreamingSingerDetector:
    def __init__(self, stream_url: str, b64_auth: str, output_dir: str = "singer_snapshots",
                 log_callback: Optional[Callable[[str], None]] = None):
        self.stream_url = stream_url
        self.b64_auth = b64_auth
        self.output_dir = os.path.abspath(output_dir)
        self.temp_dir = os.path.join(self.output_dir, "temp_hq")
        self.log = log_callback or (lambda msg: print(msg))

        os.makedirs(self.output_dir, exist_ok=True)
        os.makedirs(self.temp_dir, exist_ok=True)

        self.headers_str = f"Authorization: Basic {self.b64_auth}\r\nUser-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64)\r\n"
        self.log("  [detect] Streaming initialized (Thread-safe mode activated)\n")

    # --------------------------------------------------------- FFmpeg utilities
    def probe_stream(self) -> tuple:
        self.log("[INFO] Probing remote stream properties...")
        cmd = [
            "ffprobe", "-v", "error",
            "-show_entries", "format=duration:stream=width,height",
            "-print_format", "json",
            "-headers", self.headers_str,
            self.stream_url
        ]
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
            if result.returncode == 0:
                data = json.loads(result.stdout)
                stream = data.get("streams", [{}])[0]
                fmt = data.get("format", {})
                w = int(stream.get("width", 1920))
                h = int(stream.get("height", 1080))
                duration = float(fmt.get("duration", 3600.0))
                return w, h, duration
        except Exception as e:
            self.log(f"[WARNING] Remote probe failed ({e}).")
        self.log("[WARNING] Falling back to default 1920x1080, 1 hour duration.")
        return 1920, 1080, 3600.0

    def fetch_hq_snapshot(self, ts: float, out_path: str, quality: int = 1):
        """Fetches a single high-quality frame directly from the remote stream."""
        cmd = [
            "ffmpeg", "-y",
            "-headers", self.headers_str,
            "-ss", f"{ts:.4f}",
            "-i", self.stream_url,
            "-vframes", "1",
            "-q:v", str(quality),
            out_path
        ]
        subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    # -------------------------------------------------------- Face & Mouth Detection
    def detect_faces(self, img: np.ndarray, face_cascade, smile_cascade) -> list:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        gray = cv2.equalizeHist(gray)

        faces_rects = face_cascade.detectMultiScale(
            gray, scaleFactor=1.15, minNeighbors=5, minSize=(60, 60)
        )

        faces = []
        for (x, y, w, h) in faces_rects:
            mouth_score = self._get_mouth_score(img, x, y, w, h, smile_cascade)
            faces.append({
                "bbox": (x, y, x + w, y + h),
                "conf": 0.9,
                "cx": x + w // 2,
                "cy": y + h // 2,
                "area": w * h,
                "mouth": mouth_score,
            })
        return faces

    def _get_mouth_score(self, img: np.ndarray, x: int, y: int, w: int, h: int, smile_cascade) -> float:
        mouth_y_start = y + int(h * 0.6)
        mouth_roi = img[mouth_y_start: y + h, x: x + w]

        if mouth_roi.size == 0:
            return 0.5

        gray_mouth = cv2.cvtColor(mouth_roi, cv2.COLOR_BGR2GRAY)
        gray_mouth = cv2.equalizeHist(gray_mouth)

        smiles = smile_cascade.detectMultiScale(
            gray_mouth, scaleFactor=1.7, minNeighbors=12, minSize=(int(w * 0.3), int(h * 0.08))
        )
        return 0.85 if len(smiles) > 0 else 0.25

    # ---------------------------------------------- Face similarity & Clustering
    @staticmethod
    def _face_similarity(a: np.ndarray, b: np.ndarray) -> float:
        if a.size == 0 or b.size == 0: return 0.0
        sz = (64, 64)
        a = cv2.resize(a, sz)
        b = cv2.resize(b, sz)
        ha = cv2.calcHist([cv2.cvtColor(a, cv2.COLOR_BGR2HSV)], [0, 1], None, [30, 40], [0, 180, 0, 256])
        hb = cv2.calcHist([cv2.cvtColor(b, cv2.COLOR_BGR2HSV)], [0, 1], None, [30, 40], [0, 180, 0, 256])
        cv2.normalize(ha, ha); cv2.normalize(hb, hb)
        return cv2.compareHist(ha, hb, cv2.HISTCMP_CORREL)

    def _cluster_faces(self, all_detections: list) -> list:
        SIM_THRESHOLD = 0.38
        clusters = []
        for det in all_detections:
            best_c, best_s = None, SIM_THRESHOLD
            for c in clusters:
                s = self._face_similarity(det["crop"], c["rep_img"])
                if s > best_s:
                    best_s, best_c = s, c
            if best_c is not None:
                best_c["members"].append(det)
            else:
                clusters.append({"rep_img": det["crop"], "members": [det]})
        return clusters

    def _score_cluster(self, cluster: dict, video_area: int) -> float:
        mems = cluster["members"]
        n = len(mems)
        freq = min(n / 8.0, 1.0)
        mouth = np.mean([m["mouth"] for m in mems])
        center = np.mean([m["center_s"] for m in mems])
        size = np.mean([m["size_s"] for m in mems])
        score = 0.30 * freq + 0.35 * mouth + 0.20 * center + 0.15 * size
        cluster["_score"] = score
        cluster["_n"] = n
        cluster["_mouth"] = mouth
        return score

    # ------------------------------------------------ Auto Crop Logic
    def _auto_crop_half_body(self, img: np.ndarray, bbox: tuple) -> np.ndarray:
        x1, y1, x2, y2 = bbox
        face_h, face_w = y2 - y1, x2 - x1
        img_h, img_w = img.shape[:2]

        new_y1 = max(0, y1 - int(face_h * 0.5))
        new_y2 = min(img_h, y2 + int(face_h * 3.5))
        new_x1 = max(0, x1 - int(face_w * 2.5))
        new_x2 = min(img_w, x2 + int(face_w * 2.5))

        return img[new_y1:new_y2, new_x1:new_x2]

    # ====================================================== PUBLIC API
    def run(self, num_samples: int = 45, num_snapshots: int = 3) -> List[Dict]:
        self.log("=" * 62)
        self.log("   [detect] CLOUD-STREAMING EXTRACTOR (Fast Threaded)")
        self.log("=" * 62)

        w, h, duration = self.probe_stream()
        video_area = w * h
        frame_bytes = w * h * 3
        self.log(f"  Duration : {duration:.1f} s   Resolution : {w}x{h}\n")

        # ── 1. Stream frames into memory ──
        self.log(f"[1/4] Streaming {num_samples} frames from Kaggle cloud into RAM ...")
        t_start = time.time()

        start_t = duration * 0.08
        end_t = duration * 0.95
        window = end_t - start_t
        extraction_fps = num_samples / window

        ffmpeg_cmd = [
            "ffmpeg",
            "-headers", self.headers_str,
            "-ss", str(start_t),
            "-t", str(window),
            "-i", self.stream_url,
            "-vf", f"fps={extraction_fps}",
            "-f", "image2pipe",
            "-pix_fmt", "bgr24",
            "-vcodec", "rawvideo", "pipe:1"
        ]

        p_ffmpeg = subprocess.Popen(ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)

        frames = []
        frame_idx = 0
        while True:
            raw_frame = p_ffmpeg.stdout.read(frame_bytes)
            if not raw_frame or len(raw_frame) < frame_bytes:
                break

            img = np.frombuffer(raw_frame, dtype=np.uint8).reshape((h, w, 3))
            ts = start_t + (frame_idx / extraction_fps)
            frames.append((img.copy(), ts))
            frame_idx += 1
            self.log(f"\r  -> Buffered Frame {frame_idx}/{num_samples}")

        p_ffmpeg.terminate()
        t_extract = time.time() - t_start
        self.log(f"\n  {len(frames)} frames loaded in {t_extract:.1f}s\n")

        if not frames:
            self.log("[ERROR] Stream failed to return frames. Aborting.")
            return []

        # ── 2. Detect faces + mouth openness (THREADED) ──
        self.log("[2/4] Detecting faces & checking for singing (threaded) ...")
        t_detect_start = time.time()
        all_dets = []
        MAX_WORKERS = min(8, os.cpu_count() or 4)

        def _detect_one(img, ts):
            local_face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
            local_smile_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_smile.xml')

            fh, fw = img.shape[:2]
            fc = (fw // 2, fh // 2)
            max_dist = np.hypot(fw, fh) / 2
            dets = []

            for face in self.detect_faces(img, local_face_cascade, local_smile_cascade):
                cx, cy = face["cx"], face["cy"]
                dist = np.hypot(cx - fc[0], cy - fc[1])
                dets.append({
                    "ts": ts, "bbox": face["bbox"], "conf": face["conf"],
                    "crop": img[face["bbox"][1]:face["bbox"][3], face["bbox"][0]:face["bbox"][2]].copy(),
                    "mouth": face["mouth"], "center_s": 1.0 - dist / max_dist, "size_s": face["area"] / video_area,
                })
            return dets

        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
            futures = [pool.submit(_detect_one, img, ts) for img, ts in frames]
            for future in as_completed(futures):
                all_dets.extend(future.result())

        all_dets.sort(key=lambda x: x["ts"])

        t_detect = time.time() - t_detect_start
        self.log(f"  Detection done in {t_detect:.1f}s ({MAX_WORKERS} threads)")
        self.log(f"  Total detections: {len(all_dets)}\n")

        if not all_dets:
            self.log("  No faces found — aborting.")
            return []

        # ── 3. Cluster & identify singer ──
        self.log("[3/4] Clustering faces to find the singer ...")
        clusters = self._cluster_faces(all_dets)
        for c in clusters:
            self._score_cluster(c, video_area)
        clusters.sort(key=lambda c: c["_score"], reverse=True)

        self.log("\n  Top clusters:")
        for i, c in enumerate(clusters[:6]):
            tag = " ★ SINGER" if i == 0 else ""
            self.log(f"    #{i+1}  score={c['_score']:.3f}  faces={c['_n']:2d}  mouth={c['_mouth']:.2f}{tag}")
        singer = clusters[0]
        self.log(f"\n  -> Singer cluster: {singer['_n']} detections\n")

        # ── 4. Pick best snapshots & extract HQ from stream (THREADED) ──
        self.log(f"[4/4] Fetching HQ frames from cloud and auto-cropping (threaded) ...\n")
        scored = sorted(singer["members"], key=lambda m: 0.50 * m["mouth"] + 0.25 * m["conf"] + 0.25 * m["center_s"], reverse=True)

        picked = []
        MIN_GAP = max(4.0, duration / (num_snapshots * 3))
        for m in scored:
            if len(picked) >= num_snapshots: break
            if all(abs(m["ts"] - p["ts"]) >= MIN_GAP for p in picked): picked.append(m)
        for m in scored:
            if len(picked) >= num_snapshots: break
            if m not in picked: picked.append(m)

        def _extract_hq_snapshot(idx, det):
            ts = det["ts"]
            bbox = det["bbox"]

            raw_path = os.path.join(self.temp_dir, f"hq_{idx}.png")
            self.fetch_hq_snapshot(ts, raw_path, quality=1)

            hq = cv2.imread(raw_path)
            if hq is None: return None

            x1, y1, x2, y2 = bbox

            # 1. Annotated
            ann = hq.copy()
            cv2.rectangle(ann, (x1, y1), (x2, y2), (0, 220, 80), 3)
            lbl = "SINGER"
            (tw, th), _ = cv2.getTextSize(lbl, cv2.FONT_HERSHEY_SIMPLEX, 1.0, 2)
            cv2.rectangle(ann, (x1, y1 - th - 14), (x1 + tw + 8, y1), (0, 220, 80), -1)
            cv2.putText(ann, lbl, (x1 + 4, y1 - 6), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2)
            info_txt = f"t = {ts:.1f}s  |  singing = {det['mouth']:.2f}"
            cv2.putText(ann, info_txt, (12, ann.shape[0] - 16), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 255), 2)
            ann_path = os.path.join(self.output_dir, f"singer_{idx+1}_annotated.png")
            cv2.imwrite(ann_path, ann, [cv2.IMWRITE_PNG_COMPRESSION, 0])

            # 2. Clean
            clean_path = os.path.join(self.output_dir, f"singer_{idx+1}_clean.png")
            cv2.imwrite(clean_path, hq, [cv2.IMWRITE_PNG_COMPRESSION, 0])

            # 3. Half-body crop
            cropped_body = self._auto_crop_half_body(hq, bbox)
            crop_path = os.path.join(self.output_dir, f"singer_{idx+1}_halfbody.png")
            cv2.imwrite(crop_path, cropped_body, [cv2.IMWRITE_PNG_COMPRESSION, 0])

            c_h, c_w = cropped_body.shape[:2]
            return {
                "annotated": ann_path, "clean": clean_path, "crop": crop_path,
                "_info": f"t={ts:.1f}s  singing={det['mouth']:.2f}  crop={c_w}x{c_h}px",
                "_idx": idx,
            }

        results = []
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
            futures = {pool.submit(_extract_hq_snapshot, idx, det): idx for idx, det in enumerate(picked[:num_snapshots])}
            for future in as_completed(futures):
                r = future.result()
                if r: results.append(r)

        results.sort(key=lambda x: x["_idx"])
        for i, r in enumerate(results):
            self.log(f"  Snapshot {i+1}:  {r['_info']}")
            self.log(f"     -> {os.path.basename(r['annotated'])}")
            self.log(f"     -> {os.path.basename(r['clean'])}")
            self.log(f"     -> {os.path.basename(r['crop'])}")

        shutil.rmtree(self.temp_dir, ignore_errors=True)
        total_t = time.time() - t_start
        self.log(f"\n  Total time: {total_t:.1f}s")
        self.log("=" * 62)
        self.log(f"  DONE — {len(results)} sets saved to: {self.output_dir}")
        self.log("=" * 62)

        return results


# ════════════════════════════════════════════════════════════════════════════
# Main entry point — called from app.py
# ════════════════════════════════════════════════════════════════════════════

def run_singer_detection(
    task_id: str,
    kaggle_username: str,
    kaggle_key: str,
    log_callback: Optional[Callable[[str], None]] = None,
    num_samples: int = 45,
    num_snapshots: int = 3,
    output_dir: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Run singer detection on the video that step 1 uploaded to Kaggle.

    Uses the user's EXACT StreamingSingerDetector class from detect_stream.py.
    Streams the video from Kaggle, samples frames, detects faces + singing,
    clusters to find the singer, and outputs:
      - singer_1_halfbody.png, singer_2_halfbody.png, etc.
      - singer_1_annotated.png, singer_1_clean.png, etc.

    Args:
        task_id: The task ID (e.g. "20260619_142543_0b331b"). Underscores
                 are auto-converted to hyphens for the Kaggle dataset slug.
        kaggle_username: Kaggle username (also the dataset owner).
        kaggle_key: Kaggle API key.
        log_callback: Function called with progress messages.
        num_samples: Number of frames to sample. Default 45.
        num_snapshots: Number of final snapshots to export. Default 3.
        output_dir: Directory for output images. Defaults to task dir.

    Returns:
        dict with keys: success, results, output_dir, error
    """
    def log(msg):
        if log_callback:
            log_callback(msg)
        else:
            print(msg)

    # ── Pre-flight checks ─────────────────────────────────────────────────
    if KaggleApi is None:
        msg = "Kaggle Python package not installed. Run: pip install kaggle"
        log(f"[Singer] ERROR: {msg}")
        return {"success": False, "results": [], "output_dir": None, "error": msg}

    if not kaggle_username or not kaggle_key:
        msg = "KAGGLE_USERNAME / KAGGLE_KEY are not set."
        log(f"[Singer] ERROR: {msg}")
        return {"success": False, "results": [], "output_dir": None, "error": msg}

    # ── Set env vars for KaggleApi.authenticate() ─────────────────────────
    os.environ["KAGGLE_USERNAME"] = kaggle_username
    os.environ["KAGGLE_KEY"] = kaggle_key

    # ── Authenticate with Kaggle ──────────────────────────────────────────
    log("[Singer] Initializing Kaggle API client...")
    try:
        api = KaggleApi()
        api.authenticate()
    except Exception as e:
        msg = f"Kaggle authentication failed: {e}"
        log(f"[Singer] ERROR: {msg}")
        return {"success": False, "results": [], "output_dir": None, "error": msg}

    username = api.config_values.get('username') or kaggle_username
    key = api.config_values.get('key') or kaggle_key

    auth_bytes = f"{username}:{key}".encode('utf-8')
    b64_auth = base64.b64encode(auth_bytes).decode('utf-8')

    # ── Derive dataset slug and filename from task_id ─────────────────────
    # Same underscore→hyphen conversion as detect_edges_video_stream.py
    dataset_name = task_id.replace("_", "-")
    filename = f"output_{task_id}.mkv"
    stream_url = f"https://www.kaggle.com/api/v1/datasets/download/{username}/{dataset_name}/{filename}"

    log(f"[Singer] Task ID:      {task_id}")
    log(f"[Singer] Dataset slug: {username}/{dataset_name}")
    log(f"[Singer] Video file:   {filename}")
    log(f"[Singer] Stream URL:   {stream_url}")

    # ── Set up output directory ───────────────────────────────────────────
    if not output_dir:
        output_dir = os.path.join(os.getcwd(), "singer_snapshots")
    os.makedirs(output_dir, exist_ok=True)

    # ═══════════════════════════════════════════════════════════════════════
    # RETRY LOOP — same as detect_edges_video_stream.py
    # Kaggle returns 403 when the dataset is still being processed.
    # ═══════════════════════════════════════════════════════════════════════
    import time as _time

    MAX_ATTEMPTS = 6
    RETRY_DELAY = 15

    results = []

    for attempt in range(1, MAX_ATTEMPTS + 1):
        log(f"[Singer] === Attempt {attempt}/{MAX_ATTEMPTS} ===")

        detector = StreamingSingerDetector(
            stream_url=stream_url,
            b64_auth=b64_auth,
            output_dir=output_dir,
            log_callback=log,
        )

        results = detector.run(num_samples=num_samples, num_snapshots=num_snapshots)

        if results:
            log(f"[Singer] SUCCESS: {len(results)} singer snapshots generated.")
            break
        else:
            log(f"[Singer] WARNING: No results on attempt {attempt}.")
            if attempt < MAX_ATTEMPTS:
                log(f"[Singer] Waiting {RETRY_DELAY}s before retry... (Kaggle may still be processing)")
                _time.sleep(RETRY_DELAY)
            else:
                msg = "Singer detection failed after all retries."
                log(f"[Singer] ERROR: {msg}")
                return {"success": False, "results": [], "output_dir": output_dir, "error": msg}

    return {
        "success": True,
        "results": results,
        "output_dir": output_dir,
        "error": None,
    }


# ── CLI wrapper (kept for manual testing) ────────────────────────────────────

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Streaming Singer Face Detector")
    parser.add_argument("-d", "--dataset", required=True, help="Kaggle dataset (e.g. owner/dataset)")
    parser.add_argument("-f", "--filename", required=True, help="Target video filename (e.g. video.mkv)")
    parser.add_argument("-o", "--output", default="singer_snapshots", help="Output directory")
    parser.add_argument("-n", "--samples", type=int, default=45, help="Number of frames to sample")
    parser.add_argument("-s", "--snapshots", type=int, default=3, help="Number of final snapshots to export")
    args = parser.parse_args()

    if "/" not in args.dataset:
        print("[ERROR] Dataset must be in the form <username>/<dataset_name>")
        sys.exit(1)
    kaggle_username, task_id_raw = args.dataset.split("/", 1)
    task_id = task_id_raw.replace("-", "_")

    kaggle_key = os.environ.get("KAGGLE_KEY", "")
    if not kaggle_key:
        print("[ERROR] KAGGLE_KEY env var not set")
        sys.exit(1)

    result = run_singer_detection(
        task_id=task_id,
        kaggle_username=kaggle_username,
        kaggle_key=kaggle_key,
        log_callback=lambda msg: print(msg),
        num_samples=args.samples,
        num_snapshots=args.snapshots,
        output_dir=args.output,
    )

    if not result["success"]:
        sys.exit(1)


if __name__ == "__main__":
    main()