#!/opt/homebrew/bin/python3.11
"""
ASRX Processor - AI-Driven Processor Contract Version 1.0

Compliant with AI-Driven Processor Contract v1.0
Effective Date: 2025-03-27

Features:
1. Standardized command-line interface
2. Redis progress reporting
3. Signal handling (SIGTERM, SIGINT)
4. Health check mode
5. Resource monitoring
6. Contract-compliant JSON output
7. Unified configuration
"""

import sys
import json
import os
import argparse
import signal
import time
import subprocess
import traceback
from datetime import datetime
from typing import Dict, Any

# Redis Publisher for progress reporting
try:
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from redis_publisher import RedisPublisher

    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False
    print(
        "WARNING: RedisPublisher not available, progress reporting disabled",
        file=sys.stderr,
    )

# Contract version
CONTRACT_VERSION = "1.0"
PROCESSOR_NAME = (
    "/Users/accusys/momentry_core_0.1/scripts/asrx_processor_contract_v1.py"
)
PROCESSOR_VERSION = "1.0.0"
MODEL_NAME = "pyannote"
MODEL_VERSION = "3.1"

# Unified configuration defaults
DEFAULT_TIMEOUT = 7200  # 2 hours for speaker diarization
DEFAULT_MODEL_SIZE = "base"
DEFAULT_DEVICE = "cpu"
DEFAULT_LANGUAGE = "auto"
DEFAULT_BATCH_SIZE = 16
DEFAULT_DIARIZATION = True
DEFAULT_MIN_SPEAKERS = 1
DEFAULT_MAX_SPEAKERS = 10


# Signal handling with timeout support
class SignalHandler:
    """Handle system signals for graceful shutdown"""

    def __init__(self):
        self.should_exit = False
        self.exit_code = 0
        signal.signal(signal.SIGTERM, self.handle_signal)
        signal.signal(signal.SIGINT, self.handle_signal)

    def handle_signal(self, signum, frame):
        """Handle termination signals"""
        print(f"\n收到信号 {signum}，正在优雅关闭...")
        self.should_exit = True
        self.exit_code = 128 + signum

    def should_stop(self):
        """Check if should stop processing"""
        return self.should_exit


# Timeout manager
class TimeoutManager:
    """Manage processing timeouts"""

    def __init__(self, timeout_seconds: int):
        self.timeout_seconds = timeout_seconds
        self.start_time = time.time()
        self.timer = None

    def check_timeout(self) -> bool:
        """Check if timeout has been reached"""
        elapsed = time.time() - self.start_time
        return elapsed > self.timeout_seconds

    def get_remaining_time(self) -> float:
        """Get remaining time in seconds"""
        elapsed = time.time() - self.start_time
        return max(0, self.timeout_seconds - elapsed)

    def format_remaining_time(self) -> str:
        """Format remaining time as HH:MM:SS"""
        remaining = self.get_remaining_time()
        hours = int(remaining // 3600)
        minutes = int((remaining % 3600) // 60)
        seconds = int(remaining % 60)
        return f"{hours:02d}:{minutes:02d}:{seconds:02d}"


# Health check functions
def check_environment() -> Dict[str, Any]:
    """Check environment and dependencies"""
    checks = []

    # Check 1: whisperx for speaker diarization
    try:
        import whisperx

        checks.append(
            {
                "name": "whisperx",
                "status": "available",
                "version": getattr(whisperx, "__version__", "unknown"),
            }
        )
    except ImportError:
        checks.append({"name": "whisperx", "status": "missing", "version": None})

    # Check 2: FFmpeg/FFprobe
    try:
        ffprobe_result = subprocess.run(
            ["ffprobe", "-version"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if ffprobe_result.returncode == 0:
            version_line = ffprobe_result.stdout.split("\n")[0]
            checks.append(
                {"name": "ffprobe", "status": "available", "version": version_line}
            )
        else:
            checks.append({"name": "ffprobe", "status": "error", "version": None})
    except (subprocess.TimeoutExpired, FileNotFoundError):
        checks.append({"name": "ffprobe", "status": "missing", "version": None})

    # Check 3: Redis (optional)
    checks.append(
        {
            "name": "redis",
            "status": "available" if REDIS_AVAILABLE else "optional",
            "version": None,
        }
    )

    # Check 4: Python version
    checks.append(
        {
            "name": "python",
            "status": "available",
            "version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
        }
    )

    # Check 5: CUDA/GPU availability (optional)
    try:
        import torch

        cuda_available = torch.cuda.is_available()
        checks.append(
            {
                "name": "cuda",
                "status": "available" if cuda_available else "optional",
                "version": torch.version.cuda if cuda_available else None,
            }
        )
    except ImportError:
        checks.append({"name": "cuda", "status": "optional", "version": None})

    return {
        "timestamp": datetime.now().isoformat(),
        "processor_name": PROCESSOR_NAME,
        "processor_version": PROCESSOR_VERSION,
        "contract_version": CONTRACT_VERSION,
        "model_name": MODEL_NAME,
        "model_version": MODEL_VERSION,
        "checks": checks,
    }


def check_video_file(video_path: str) -> Dict[str, Any]:
    """Check video file properties"""
    try:
        result = subprocess.run(
            [
                "ffprobe",
                "-v",
                "error",
                "-select_streams",
                "v:0",
                "-show_entries",
                "stream=codec_name,width,height,duration,r_frame_rate",
                "-show_entries",
                "format=duration,size",
                "-of",
                "json",
                video_path,
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )

        if result.returncode != 0:
            return {
                "valid": False,
                "error": result.stderr[:200] if result.stderr else "Unknown error",
            }

        info = json.loads(result.stdout)

        video_info = {}
        if "streams" in info and len(info["streams"]) > 0:
            stream = info["streams"][0]
            video_info = {
                "codec": stream.get("codec_name", "unknown"),
                "width": int(stream.get("width", 0)),
                "height": int(stream.get("height", 0)),
                "duration": float(stream.get("duration", 0)),
                "frame_rate": stream.get("r_frame_rate", "0/0"),
            }

        format_info = {}
        if "format" in info:
            format_info = {
                "format_duration": float(info["format"].get("duration", 0)),
                "file_size": int(info["format"].get("size", 0)),
            }

        return {
            "valid": True,
            "video_info": video_info,
            "format_info": format_info,
            "exists": os.path.exists(video_path),
            "file_size": os.path.getsize(video_path)
            if os.path.exists(video_path)
            else 0,
        }

    except Exception as e:
        return {"valid": False, "error": str(e)}


# Main processing function
def process_asrx(
    video_path: str,
    output_path: str,
    uuid: str = "",
    model_size: str = DEFAULT_MODEL_SIZE,
    device: str = DEFAULT_DEVICE,
    language: str = DEFAULT_LANGUAGE,
    batch_size: int = DEFAULT_BATCH_SIZE,
    diarization: bool = DEFAULT_DIARIZATION,
    min_speakers: int = DEFAULT_MIN_SPEAKERS,
    max_speakers: int = DEFAULT_MAX_SPEAKERS,
    timeout: int = DEFAULT_TIMEOUT,
) -> Dict[str, Any]:
    """Process video for speaker diarization using whisperx"""

    # Initialize
    signal_handler = SignalHandler()
    timeout_manager = TimeoutManager(timeout)
    publisher = RedisPublisher(uuid) if REDIS_AVAILABLE and uuid else None

    def publish(stage: str, message: str, data: Dict = None):
        if publisher:
            publisher.info(PROCESSOR_NAME, stage, message, data)

    publish("ASRX_START", f"开始处理: {os.path.basename(video_path)}")

    result = {
        "processor_name": PROCESSOR_NAME,
        "processor_version": PROCESSOR_VERSION,
        "contract_version": CONTRACT_VERSION,
        "model_name": MODEL_NAME,
        "model_version": MODEL_VERSION,
        "video_path": video_path,
        "output_path": output_path,
        "uuid": uuid,
        "timestamp": datetime.now().isoformat(),
        "parameters": {
            "model_size": model_size,
            "device": device,
            "language": language,
            "batch_size": batch_size,
            "diarization": diarization,
            "min_speakers": min_speakers,
            "max_speakers": max_speakers,
            "timeout": timeout,
        },
        "success": False,
        "error": None,
        "segments": [],
        "speakers": [],
        "processing_time": 0,
        "resource_usage": {},
    }

    start_time = time.time()

    try:
        # Check timeout
        if timeout_manager.check_timeout():
            raise TimeoutError(f"超时 ({timeout} 秒)")

        # Check if should exit
        if signal_handler.should_stop():
            raise KeyboardInterrupt("收到停止信号")

        # Check video file
        publish("ASRX_CHECK_VIDEO", "检查视频文件")
        video_check = check_video_file(video_path)
        if not video_check.get("valid", False):
            raise ValueError(f"无效的视频文件: {video_check.get('error', '未知错误')}")

        result["video_info"] = video_check.get("video_info", {})
        result["format_info"] = video_check.get("format_info", {})

        # Import whisperx
        publish("ASRX_LOAD_MODEL", f"加载模型: {model_size}")
        try:
            import whisperx
        except ImportError as e:
            raise ImportError(f"whisperx 未安装: {e}")

        # Load model
        publish("ASRX_LOADING", f"加载 whisperx 模型 ({model_size}, {device})")
        model = whisperx.load_model(
            model_size,
            device=device,
            compute_type="int8" if device == "cpu" else "float16",
        )

        # Transcribe
        publish("ASRX_TRANSCRIBING", "转录音频")
        transcript = model.transcribe(
            video_path,
            language=language if language != "auto" else None,
            batch_size=batch_size,
        )

        # Align timestamps
        publish("ASRX_ALIGNING", "对齐时间戳")
        model_a, metadata = whisperx.load_align_model(
            language_code=transcript["language"]
        )
        transcript = whisperx.align(
            transcript["segments"],
            model_a,
            metadata,
            video_path,
            device,
            return_char_alignments=False,
        )

        # Speaker diarization
        if diarization:
            publish("ASRX_DIARIZATION", "说话人分离")
            diarize_model = whisperx.DiarizationPipeline(
                use_auth_token=None, device=device
            )

            # Add min/max speakers
            diarize_segments = diarize_model(
                video_path,
                min_speakers=min_speakers,
                max_speakers=max_speakers,
            )

            transcript = whisperx.assign_word_speakers(diarize_segments, transcript)

            # Extract speaker information
            speakers = {}
            for segment in transcript["segments"]:
                if "speaker" in segment:
                    speaker_id = segment["speaker"]
                    if speaker_id not in speakers:
                        speakers[speaker_id] = {
                            "id": speaker_id,
                            "segment_count": 0,
                            "total_words": 0,
                            "total_duration": 0.0,
                        }

                    speakers[speaker_id]["segment_count"] += 1
                    speakers[speaker_id]["total_words"] += len(
                        segment.get("text", "").split()
                    )
                    speakers[speaker_id]["total_duration"] += segment.get(
                        "end", 0
                    ) - segment.get("start", 0)

            result["speakers"] = list(speakers.values())

        # Format segments
        segments = []
        for segment in transcript.get("segments", []):
            segments.append(
                {
                    "start": segment.get("start", 0.0),
                    "end": segment.get("end", 0.0),
                    "text": segment.get("text", ""),
                    "speaker": segment.get("speaker", None),
                    "words": segment.get("words", []),
                    "confidence": segment.get("confidence", 0.0),
                }
            )

        result["segments"] = segments
        result["language"] = transcript.get("language", "unknown")
        result["success"] = True

        publish("ASRX_COMPLETE", f"完成: {len(segments)} 个片段")

    except TimeoutError as e:
        result["error"] = f"处理超时: {e}"
        publish("ASRX_TIMEOUT", f"超时: {e}")
    except KeyboardInterrupt:
        result["error"] = "处理被用户中断"
        publish("ASRX_INTERRUPTED", "处理被中断")
    except ImportError as e:
        result["error"] = f"依赖缺失: {e}"
        publish("ASRX_MISSING_DEPS", f"缺少依赖: {e}")
    except Exception as e:
        result["error"] = f"处理错误: {str(e)}"
        publish("ASRX_ERROR", f"错误: {str(e)}")
        traceback.print_exc()

    # Calculate processing time
    processing_time = time.time() - start_time
    result["processing_time"] = processing_time

    # Add resource usage
    try:
        import psutil

        process = psutil.Process()
        memory_info = process.memory_info()
        result["resource_usage"] = {
            "cpu_percent": process.cpu_percent(),
            "memory_mb": memory_info.rss / (1024 * 1024),
            "user_time": process.cpu_times().user,
            "system_time": process.cpu_times().system,
        }
    except ImportError:
        result["resource_usage"] = {"error": "psutil not available"}

    # Save result
    try:
        with open(output_path, "w") as f:
            json.dump(result, f, indent=2, ensure_ascii=False)
        publish("ASRX_SAVED", f"结果保存到: {output_path}")
    except Exception as e:
        result["error"] = f"保存结果失败: {str(e)}"
        publish("ASRX_SAVE_ERROR", f"保存失败: {str(e)}")

    return result


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(
        description=f"{PROCESSOR_NAME.upper()} Processor v{PROCESSOR_VERSION} - Speaker Diarization"
    )
    parser.add_argument("video_path", help="Path to input video file")
    parser.add_argument("output_path", help="Path to output JSON file")
    parser.add_argument("--uuid", help="UUID for progress tracking", default="")
    parser.add_argument(
        "--model-size",
        help=f"Model size (default: {DEFAULT_MODEL_SIZE})",
        default=DEFAULT_MODEL_SIZE,
        choices=["tiny", "base", "small", "medium", "large-v3"],
    )
    parser.add_argument(
        "--device",
        help=f"Device to use (default: {DEFAULT_DEVICE})",
        default=DEFAULT_DEVICE,
        choices=["cpu", "cuda"],
    )
    parser.add_argument(
        "--language",
        help=f"Language code or 'auto' (default: {DEFAULT_LANGUAGE})",
        default=DEFAULT_LANGUAGE,
    )
    parser.add_argument(
        "--batch-size",
        help=f"Batch size for processing (default: {DEFAULT_BATCH_SIZE})",
        type=int,
        default=DEFAULT_BATCH_SIZE,
    )
    parser.add_argument(
        "--no-diarization",
        help="Disable speaker diarization",
        action="store_true",
    )
    parser.add_argument(
        "--min-speakers",
        help=f"Minimum number of speakers (default: {DEFAULT_MIN_SPEAKERS})",
        type=int,
        default=DEFAULT_MIN_SPEAKERS,
    )
    parser.add_argument(
        "--max-speakers",
        help=f"Maximum number of speakers (default: {DEFAULT_MAX_SPEAKERS})",
        type=int,
        default=DEFAULT_MAX_SPEAKERS,
    )
    parser.add_argument(
        "--timeout",
        help=f"Timeout in seconds (default: {DEFAULT_TIMEOUT})",
        type=int,
        default=DEFAULT_TIMEOUT,
    )
    parser.add_argument(
        "--health-check",
        help="Run health check and exit",
        action="store_true",
    )
    parser.add_argument(
        "--check-video",
        help="Check video file and exit",
        action="store_true",
    )

    args = parser.parse_args()

    # Health check mode
    if args.health_check:
        health = check_environment()
        print(json.dumps(health, indent=2, ensure_ascii=False))
        return (
            0
            if all(c["status"] in ["available", "optional"] for c in health["checks"])
            else 1
        )

    # Video check mode
    if args.check_video:
        video_check = check_video_file(args.video_path)
        print(json.dumps(video_check, indent=2, ensure_ascii=False))
        return 0 if video_check.get("valid", False) else 1

    # Normal processing mode
    result = process_asrx(
        video_path=args.video_path,
        output_path=args.output_path,
        uuid=args.uuid,
        model_size=args.model_size,
        device=args.device,
        language=args.language,
        batch_size=args.batch_size,
        diarization=not args.no_diarization,
        min_speakers=args.min_speakers,
        max_speakers=args.max_speakers,
        timeout=args.timeout,
    )

    # Print result summary
    if result.get("success", False):
        print(f"✅ {PROCESSOR_NAME.upper()} 处理成功")
        print(f"   片段数: {len(result.get('segments', []))}")
        print(f"   说话人数: {len(result.get('speakers', []))}")
        print(f"   处理时间: {result.get('processing_time', 0):.1f} 秒")
        print(f"   输出文件: {args.output_path}")
        return 0
    else:
        print(f"❌ {PROCESSOR_NAME.upper()} 处理失败")
        print(f"   错误: {result.get('error', '未知错误')}")
        return 1


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