#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Video Generator Module
视频生成器模块

Main module for AI-powered video generation from images.
"""

import os
import glob
from typing import Dict, Any, List, Optional, Callable
import sys

from python_core.config import settings
from python_core.utils import setup_logger
# 导入同目录下的模块
from python_core.ai_video.cloud_storage import CloudStorage
from python_core.ai_video.api_client import APIClient
from python_core.utils.jsonrpc import create_response_handler, create_progress_reporter, JSONRPCError

logger = setup_logger(__name__)

class VideoGenerator:
    """AI video generator for converting images to videos."""
    
    def __init__(self,
                 api_key: str = None,
                 cos_config: Dict[str, str] = None):
        """
        Initialize video generator.

        Args:
            api_key: API key for video generation service
            cos_config: Cloud storage configuration
        """
        print("Initializing VideoGenerator...")
        sys.stdout.flush()

        print("Creating APIClient...")
        sys.stdout.flush()
        self.api_client = APIClient(api_key=api_key)
        print("APIClient created successfully")
        sys.stdout.flush()

        print("Creating CloudStorage...")
        sys.stdout.flush()
        if cos_config:
            self.cloud_storage = CloudStorage(**cos_config)
        else:
            self.cloud_storage = CloudStorage()
        print("CloudStorage created successfully")
        print("VideoGenerator initialization complete")
        sys.stdout.flush()
    
    def generate_video_from_image(self,
                                image_path: str,
                                prompt: str,
                                duration: str = '5',
                                model_type: str = 'lite',
                                save_path: str = None,
                                timeout: int = 180,
                                interval: int = 2,
                                progress_callback: Optional[Callable[[str], None]] = None,
                                request_id: str = None) -> Dict[str, Any]:
        """
        Generate video from a single image.
        
        Args:
            image_path: Path to input image
            prompt: Text prompt for video generation
            duration: Video duration ('5' or '10')
            model_type: Model type ('lite' or 'pro')
            save_path: Directory to save generated video
            timeout: Task timeout in seconds
            interval: Status check interval in seconds
            progress_callback: Optional progress callback function
            
        Returns:
            Dictionary with generation result
        """
        # Initialize JSON-RPC handlers
        rpc = create_response_handler(request_id)
        progress = create_progress_reporter()

        result = {'status': False, 'video_path': '', 'video_url': '', 'msg': ''}

        try:
            # Check if image file exists
            if not os.path.exists(image_path):
                # Try to find the file in common locations
                possible_paths = [
                    image_path,
                    os.path.join(os.getcwd(), image_path),
                    os.path.join(os.path.dirname(__file__), '..', '..', image_path),
                    os.path.join('uploads', image_path) if not os.path.isabs(image_path) else image_path
                ]

                found_path = None
                for path in possible_paths:
                    if os.path.exists(path):
                        found_path = path
                        break

                if found_path:
                    image_path = found_path
                else:
                    result['msg'] = f'Image file not found: {image_path}. Searched in: {possible_paths}'
                    logger.error(result['msg'])
                    return result
            
            # Step 1: Upload image to cloud storage
            progress.step("upload", "[1/4] 正在上传图片到云存储...")
            if progress_callback:
                progress_callback("[1/4] 正在上传图片到云存储...")
            
            upload_result = self.cloud_storage.upload_file(image_path)
            if not upload_result['status']:
                result['msg'] = f"Failed to upload image: {upload_result['msg']}"
                logger.error(result['msg'])
                progress.error(result['msg'])
                rpc.error(JSONRPCError.UPLOAD_FAILED, "Image upload failed", result['msg'])
                return result

            img_url = upload_result['data']
            progress.step("upload_success", "[1/4] 图片上传成功")
            if progress_callback:
                progress_callback("[1/4] 图片上传成功")
            
            # Step 2: Submit video generation task
            if progress_callback:
                progress_callback("[2/4] 正在提交视频生成任务...")
            
            task_result = self.api_client.submit_task(prompt, img_url, duration, model_type)
            if not task_result['status']:
                result['msg'] = f"Failed to submit task: {task_result['msg']}"
                logger.error(result['msg'])
                return result
            
            task_id = task_result['data']
            if progress_callback:
                progress_callback(f"[2/4] 任务提交成功，任务ID: {task_id}")
            
            # Step 3: Wait for task completion
            if progress_callback:
                progress_callback("[3/4] 正在等待视频生成完成...")
            
            completion_result = self.api_client.wait_for_completion(
                task_id, 
                timeout=timeout, 
                interval=interval,
                progress_callback=progress_callback
            )
            
            if not completion_result['status']:
                result['msg'] = f"Video generation failed: {completion_result['msg']}"
                logger.error(result['msg'])
                return result
            
            video_url = completion_result['data']
            result['video_url'] = video_url
            
            # Step 4: Download video if save_path is provided
            if save_path:
                if progress_callback:
                    progress_callback("[4/4] 正在下载视频到本地...")
                
                video_path = self.cloud_storage.download_file(video_url, save_path)
                if video_path and os.path.exists(video_path):
                    result['status'] = True
                    result['video_path'] = video_path
                    result['msg'] = '视频生成并下载成功'
                    if progress_callback:
                        progress_callback(f"[4/4] 视频下载成功: {os.path.basename(video_path)}")
                else:
                    result['msg'] = '视频下载失败'
                    if progress_callback:
                        progress_callback("[4/4] 视频下载失败")
                    logger.error(result['msg'])
            else:
                result['status'] = True
                result['video_path'] = video_url
                result['msg'] = '视频生成成功'
            
        except Exception as e:
            import traceback
            error_details = {
                'error_type': type(e).__name__,
                'error_message': str(e),
                'traceback': traceback.format_exc()
            }
            result['msg'] = f'处理过程中发生异常: {error_details["error_type"]}: {error_details["error_message"]}'
            result['error_details'] = error_details
            logger.error(f"Video generation failed: {result['msg']}")
            logger.error(f"Traceback: {error_details['traceback']}")
            progress.error(result['msg'])
            rpc.error(JSONRPCError.GENERATION_FAILED, "Video generation failed", error_details)

        # Send final result via JSON-RPC
        if result['status']:
            rpc.success(result)
        else:
            rpc.error(JSONRPCError.GENERATION_FAILED, result.get('msg', 'Unknown error'), result)

        return result
    
    def batch_generate_videos(self,
                            image_folder: str,
                            prompts: List[str],
                            output_folder: str,
                            duration: str = '5',
                            model_type: str = 'lite',
                            timeout: int = 300,
                            interval: int = 3,
                            progress_callback: Optional[Callable[[str], None]] = None,
                            request_id: str = None) -> Dict[str, Any]:
        """
        Batch generate videos from multiple images.
        
        Args:
            image_folder: Folder containing input images
            prompts: List of prompts (will cycle through them)
            output_folder: Folder to save generated videos
            duration: Video duration ('5' or '10')
            model_type: Model type ('lite' or 'pro')
            timeout: Task timeout in seconds
            interval: Status check interval in seconds
            progress_callback: Optional progress callback function
            
        Returns:
            Dictionary with batch processing result
        """
        # Initialize JSON-RPC handlers for batch processing
        rpc = create_response_handler(request_id)
        progress = create_progress_reporter()

        result = {'status': False, 'success_count': 0, 'failed_count': 0, 'results': [], 'msg': ''}

        try:
            if progress_callback:
                progress_callback(f"开始批量处理任务...")
                progress_callback(f"  图片文件夹: {image_folder}")
                progress_callback(f"  视频保存目录: {output_folder}")
                progress_callback(f"  视频时长: {duration}秒")
                progress_callback(f"  模型类型: {model_type}")
            
            # Find image files
            # Use a more robust approach to find image files
            image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp']
            image_files = []

            # Get all files in the directory
            try:
                all_files = os.listdir(image_folder)
                for file in all_files:
                    file_path = os.path.join(image_folder, file)
                    if os.path.isfile(file_path):
                        # Check if file extension matches (case-insensitive)
                        file_ext = os.path.splitext(file)[1].lower()
                        if file_ext in image_extensions:
                            image_files.append(file_path)
            except OSError as e:
                result['msg'] = f"Failed to read directory {image_folder}: {str(e)}"
                logger.error(result['msg'])
                return result

            if not image_files:
                result['msg'] = f"No image files found in {image_folder}"
                logger.error(result['msg'])
                return result

            if progress_callback:
                progress_callback(f'目录下找到 {len(image_files)} 张图片')

            # Debug: log found image files
            logger.info(f"Found {len(image_files)} image files in {image_folder}")
            for i, img_file in enumerate(image_files):
                logger.info(f"  {i+1}. {os.path.basename(img_file)}")
            
            if not prompts:
                result['msg'] = "No prompts provided"
                logger.error(result['msg'])
                return result
            
            if progress_callback:
                progress_callback(f'解析到 {len(prompts)} 个提示词，将循环使用')
            
            # Ensure output directory exists
            os.makedirs(output_folder, exist_ok=True)
            
            success_count = 0
            failed_count = 0
            
            # Process each image
            for i, image_path in enumerate(image_files):
                image_name = os.path.basename(image_path)
                # Use cycling prompts
                current_prompt = prompts[i % len(prompts)]
                
                if progress_callback:
                    progress_callback(f"[{i + 1}/{len(image_files)}] 处理图片: {image_name}")
                    progress_callback(f"  使用提示词[{i % len(prompts) + 1}]: {current_prompt[:50]}{'...' if len(current_prompt) > 50 else ''}")
                
                try:
                    # Generate video for this image
                    video_result = self.generate_video_from_image(
                        image_path=image_path,
                        prompt=current_prompt,
                        duration=duration,
                        model_type=model_type,
                        save_path=output_folder,
                        timeout=timeout,
                        interval=interval,
                        progress_callback=progress_callback
                    )
                    
                    if video_result and video_result.get('status'):
                        success_count += 1
                        video_path = video_result.get('video_path', '')
                        if progress_callback:
                            progress_callback(f"  [成功] 视频生成成功: {os.path.basename(video_path)}")
                    else:
                        failed_count += 1
                        error_msg = video_result.get('msg', '未知错误') if video_result else '处理失败'
                        if progress_callback:
                            progress_callback(f"  [失败] 视频生成失败: {error_msg}")
                    
                    result['results'].append({
                        'image_path': image_path,
                        'prompt': current_prompt,
                        'result': video_result
                    })
                    
                except Exception as e:
                    failed_count += 1
                    error_msg = str(e)
                    if progress_callback:
                        progress_callback(f"  [错误] 处理图片时出错: {error_msg}")
                    
                    result['results'].append({
                        'image_path': image_path,
                        'prompt': current_prompt,
                        'result': {'status': False, 'msg': error_msg}
                    })
            
            # Final statistics
            result['success_count'] = success_count
            result['failed_count'] = failed_count
            result['status'] = success_count > 0
            result['msg'] = f'处理完成！成功: {success_count}, 失败: {failed_count}'
            
            if progress_callback:
                progress_callback("=" * 50)
                progress_callback(f"处理完成！成功: {success_count}, 失败: {failed_count}")
                if success_count > 0:
                    progress_callback(f"生成的视频已保存到: {output_folder}")
            
        except Exception as e:
            result['msg'] = f'批量处理过程中发生错误: {str(e)}'
            logger.error(result['msg'])
            progress.error(result['msg'])
            rpc.error(JSONRPCError.GENERATION_FAILED, "Batch processing failed", str(e))

        # Send final batch result via JSON-RPC
        if result['status']:
            progress.complete(f"批量处理完成！成功: {result['success_count']}, 失败: {result['failed_count']}")
            rpc.success(result)
        else:
            rpc.error(JSONRPCError.GENERATION_FAILED, result.get('msg', 'Batch processing failed'), result)

        return result


def main():
    """Command line interface for AI video generation."""
    import argparse
    import json
    import os

    # Configure encoding for consistent output
    if os.name == 'nt':  # Windows
        # Set console code page to UTF-8 on Windows
        try:
            import subprocess
            subprocess.run(['chcp', '65001'], shell=True, capture_output=True)
        except:
            pass

    # Ensure stdout uses UTF-8 encoding
    if hasattr(sys.stdout, 'reconfigure'):
        try:
            sys.stdout.reconfigure(encoding='utf-8')
            sys.stderr.reconfigure(encoding='utf-8')
        except:
            pass

    parser = argparse.ArgumentParser(description="AI Video Generator")
    parser.add_argument("--action", required=True,
                       choices=["single", "batch"],
                       help="Action to perform")
    parser.add_argument("--image", help="Image file path (for single)")
    parser.add_argument("--folder", help="Image folder path (for batch)")
    parser.add_argument("--prompt", help="Generation prompt")
    parser.add_argument("--prompts", help="JSON string of prompts list (for batch)")
    parser.add_argument("--output", help="Output directory")
    parser.add_argument("--duration", default="5", choices=["5", "10"], help="Video duration")
    parser.add_argument("--model", default="lite", choices=["lite", "pro"], help="Model type")
    parser.add_argument("--timeout", type=int, default=180, help="Task timeout in seconds")

    args = parser.parse_args()

    try:
        generator = VideoGenerator()

        def progress_callback(message):
            print(message)

        if args.action == "single":
            if not args.image or not args.prompt or not args.output:
                raise ValueError("Single mode requires --image, --prompt, and --output")

            result = generator.generate_video_from_image(
                image_path=args.image,
                prompt=args.prompt,
                duration=args.duration,
                model_type=args.model,
                save_path=args.output,
                timeout=args.timeout,
                progress_callback=progress_callback,
                request_id="cli_request"
            )

        elif args.action == "batch":
            if not args.folder or not args.prompts or not args.output:
                raise ValueError("Batch mode requires --folder, --prompts, and --output")

            prompts = json.loads(args.prompts)

            result = generator.batch_generate_videos(
                image_folder=args.folder,
                prompts=prompts,
                output_folder=args.output,
                duration=args.duration,
                model_type=args.model,
                timeout=args.timeout,
                progress_callback=progress_callback,
                request_id="cli_batch_request"
            )

        # Use safe encoding for final output
        result_json = json.dumps(result, ensure_ascii=True, indent=2)
        if hasattr(sys.stdout, 'buffer'):
            sys.stdout.buffer.write(result_json.encode('utf-8'))
            sys.stdout.buffer.write(b'\n')
            sys.stdout.buffer.flush()
        else:
            print(result_json)
            sys.stdout.flush()

    except Exception as e:
        error_result = {
            "status": False,
            "error": str(e)
        }
        error_json = json.dumps(error_result, ensure_ascii=True, indent=2)
        if hasattr(sys.stdout, 'buffer'):
            sys.stdout.buffer.write(error_json.encode('utf-8'))
            sys.stdout.buffer.write(b'\n')
            sys.stdout.buffer.flush()
        else:
            print(error_json)
            sys.stdout.flush()
        sys.exit(1)


if __name__ == "__main__":
    main()
