#!/opt/homebrew/bin/python3.11
"""
Caption 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 tempfile
import time
import subprocess
import traceback
from datetime import datetime
from typing import Dict, Any, List

# 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/caption_processor_contract_v1.py"
)
PROCESSOR_VERSION = "1.0.0"
MODEL_NAME = "gpt-4-vision-preview"
MODEL_VERSION = "latest"

# Unified configuration defaults
DEFAULT_TIMEOUT = 1800  # 30 minutes for caption generation
DEFAULT_MAX_FRAMES = 30
DEFAULT_FRAME_INTERVAL = 2.0
DEFAULT_MODEL = "openai"  # openai, local, or none
DEFAULT_MODEL_NAME = "gpt-4-vision-preview"
DEFAULT_TEMPERATURE = 0.7
DEFAULT_MAX_TOKENS = 300


# 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: FFmpeg/FFprobe for frame extraction
    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 2: OpenAI API (optional)
    try:
        import openai

        checks.append(
            {
                "name": "openai",
                "status": "available",
                "version": openai.__version__,
            }
        )
    except ImportError:
        checks.append({"name": "openai", "status": "optional", "version": None})

    # Check 3: PIL/Pillow for image processing
    try:
        from PIL import Image

        checks.append(
            {
                "name": "pillow",
                "status": "available",
                "version": Image.__version__,
            }
        )
    except ImportError:
        checks.append({"name": "pillow", "status": "optional", "version": None})

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

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

    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)}


def extract_frames(
    video_path: str,
    max_frames: int = DEFAULT_MAX_FRAMES,
    frame_interval: float = DEFAULT_FRAME_INTERVAL,
) -> List[Dict[str, Any]]:
    """Extract frames from video at regular intervals"""

    frames = []
    temp_dir = tempfile.mkdtemp(prefix="caption_frames_")

    try:
        # Get video duration
        duration_result = subprocess.run(
            [
                "ffprobe",
                "-v",
                "quiet",
                "-show_entries",
                "format=duration",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                video_path,
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )

        if duration_result.returncode == 0:
            try:
                duration = float(duration_result.stdout.strip())
            except ValueError:
                duration = 60.0  # Default fallback
        else:
            duration = 60.0

        # Calculate actual number of frames to extract
        if frame_interval > 0:
            num_frames = min(max_frames, int(duration / frame_interval))
            if num_frames < 1:
                num_frames = 1
        else:
            num_frames = max_frames

        # Extract frames
        for i in range(num_frames):
            timestamp = (duration / num_frames) * i if num_frames > 1 else 0
            frame_filename = os.path.join(temp_dir, f"frame_{i:04d}.jpg")

            # Extract frame using ffmpeg
            cmd = [
                "ffmpeg",
                "-ss",
                str(timestamp),
                "-i",
                video_path,
                "-vframes",
                "1",
                "-q:v",
                "2",  # Quality factor (2 = high quality)
                "-y",  # Overwrite output file
                frame_filename,
            ]

            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=30,
            )

            if result.returncode == 0 and os.path.exists(frame_filename):
                frames.append(
                    {
                        "frame_id": i,
                        "timestamp": timestamp,
                        "file_path": frame_filename,
                        "file_size": os.path.getsize(frame_filename),
                    }
                )
            else:
                print(f"警告: 无法提取帧 {i} (时间戳: {timestamp})")

    except Exception as e:
        print(f"提取帧时出错: {e}")

    return frames


def generate_caption_for_frame(
    frame_path: str, model: str = DEFAULT_MODEL, **kwargs
) -> str:
    """Generate caption for a single frame"""

    if model == "openai":
        try:
            import openai
            from PIL import Image
            import base64

            # Read and encode image
            with open(frame_path, "rb") as image_file:
                base64_image = base64.b64encode(image_file.read()).decode("utf-8")

            # Prepare messages for GPT-4 Vision
            messages = [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Describe this image in detail. Include objects, actions, colors, and context.",
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            },
                        },
                    ],
                }
            ]

            # Call OpenAI API
            response = openai.chat.completions.create(
                model=kwargs.get("model_name", DEFAULT_MODEL_NAME),
                messages=messages,
                max_tokens=kwargs.get("max_tokens", DEFAULT_MAX_TOKENS),
                temperature=kwargs.get("temperature", DEFAULT_TEMPERATURE),
            )

            return response.choices[0].message.content

        except ImportError:
            return "OpenAI not available"
        except Exception as e:
            return f"Caption generation error: {str(e)}"

    elif model == "local":
        # Placeholder for local model implementation
        try:
            from PIL import Image

            image = Image.open(frame_path)
            width, height = image.size
            return f"Image size: {width}x{height} pixels. Local caption model not implemented."
        except ImportError:
            return "PIL not available"

    else:
        # Fallback: basic description
        try:
            from PIL import Image

            image = Image.open(frame_path)
            width, height = image.size
            return f"Image size: {width}x{height} pixels. No caption model specified."
        except ImportError:
            return "Basic image information not available"


# Main processing function
def process_caption(
    video_path: str,
    output_path: str,
    uuid: str = "",
    max_frames: int = DEFAULT_MAX_FRAMES,
    frame_interval: float = DEFAULT_FRAME_INTERVAL,
    model: str = DEFAULT_MODEL,
    model_name: str = DEFAULT_MODEL_NAME,
    temperature: float = DEFAULT_TEMPERATURE,
    max_tokens: int = DEFAULT_MAX_TOKENS,
    timeout: int = DEFAULT_TIMEOUT,
) -> Dict[str, Any]:
    """Process video for caption generation"""

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

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

    if publisher:
        publish("CAPTION_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": {
            "max_frames": max_frames,
            "frame_interval": frame_interval,
            "model": model,
            "model_name": model_name,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "timeout": timeout,
        },
        "success": False,
        "error": None,
        "frames": [],
        "captions": [],
        "processing_time": 0,
        "resource_usage": {},
    }

    start_time = time.time()
    temp_dir = None

    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
        if publisher:
            publish("CAPTION_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", {})

        # Extract frames
        if publisher:
            publish("CAPTION_EXTRACT_FRAMES", f"提取帧 (最多 {max_frames} 个)")

        frames = extract_frames(video_path, max_frames, frame_interval)

        if not frames:
            raise ValueError("无法从视频中提取帧")

        result["frames_extracted"] = len(frames)

        if publisher:
            publish("CAPTION_FRAMES_EXTRACTED", f"已提取 {len(frames)} 个帧")

        # Generate captions for each frame
        captions = []
        for i, frame in enumerate(frames):
            # Check timeout and signals periodically
            if timeout_manager.check_timeout():
                raise TimeoutError(f"超时 ({timeout} 秒)")
            if signal_handler.should_stop():
                raise KeyboardInterrupt("收到停止信号")

            if publisher:
                publish("CAPTION_GENERATING", f"生成字幕 {i + 1}/{len(frames)}")

            caption = generate_caption_for_frame(
                frame["file_path"],
                model=model,
                model_name=model_name,
                temperature=temperature,
                max_tokens=max_tokens,
            )

            captions.append(
                {
                    "frame_id": frame["frame_id"],
                    "timestamp": frame["timestamp"],
                    "caption": caption,
                    "frame_file": frame["file_path"],
                    "frame_size": frame["file_size"],
                }
            )

            # Clean up frame file
            try:
                os.remove(frame["file_path"])
            except:
                pass

        result["captions"] = captions
        result["caption_count"] = len(captions)
        result["success"] = True

        if publisher:
            publish("CAPTION_COMPLETE", f"完成: {len(captions)} 个字幕")

        # Clean up temp directory
        if temp_dir and os.path.exists(temp_dir):
            try:
                import shutil

                shutil.rmtree(temp_dir)
            except:
                pass

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

        # Clean up on error
        if temp_dir and os.path.exists(temp_dir):
            try:
                import shutil

                shutil.rmtree(temp_dir)
            except:
                pass

    # 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)
        if publisher:
            publish("CAPTION_SAVED", f"结果保存到: {output_path}")
    except Exception as e:
        result["error"] = f"保存结果失败: {str(e)}"
        if publisher:
            publish("CAPTION_SAVE_ERROR", f"保存失败: {str(e)}")

    return result


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(
        description=f"{PROCESSOR_NAME.upper()} Processor v{PROCESSOR_VERSION} - Video Caption Generation"
    )
    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(
        "--max-frames",
        help=f"Maximum frames to extract (default: {DEFAULT_MAX_FRAMES})",
        type=int,
        default=DEFAULT_MAX_FRAMES,
    )
    parser.add_argument(
        "--frame-interval",
        help=f"Seconds between frames (default: {DEFAULT_FRAME_INTERVAL})",
        type=float,
        default=DEFAULT_FRAME_INTERVAL,
    )
    parser.add_argument(
        "--model",
        help=f"Caption model to use (default: {DEFAULT_MODEL})",
        default=DEFAULT_MODEL,
        choices=["openai", "local", "none"],
    )
    parser.add_argument(
        "--model-name",
        help=f"Model name for OpenAI (default: {DEFAULT_MODEL_NAME})",
        default=DEFAULT_MODEL_NAME,
    )
    parser.add_argument(
        "--temperature",
        help=f"Temperature for generation (default: {DEFAULT_TEMPERATURE})",
        type=float,
        default=DEFAULT_TEMPERATURE,
    )
    parser.add_argument(
        "--max-tokens",
        help=f"Maximum tokens per caption (default: {DEFAULT_MAX_TOKENS})",
        type=int,
        default=DEFAULT_MAX_TOKENS,
    )
    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_caption(
        video_path=args.video_path,
        output_path=args.output_path,
        uuid=args.uuid,
        max_frames=args.max_frames,
        frame_interval=args.frame_interval,
        model=args.model,
        model_name=args.model_name,
        temperature=args.temperature,
        max_tokens=args.max_tokens,
        timeout=args.timeout,
    )

    # Print result summary
    if result.get("success", False):
        print(f"✅ {PROCESSOR_NAME.upper()} 处理成功")
        print(f"   帧数: {result.get('frames_extracted', 0)}")
        print(f"   字幕数: {result.get('caption_count', 0)}")
        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())
