import asyncio
import os
import time

from zai import ZaiClient

class VideoGenerator:
    def __init__(self):
        self.client = ZaiClient()

    async def video_generate(
        self,
        prompt: str,
        quality: str = 'speed',
        with_audio: bool = True,
        max_wait_time: int = 300,
    ):
        """
        Asynchronous video generation method

        Args:
            prompt: Video generation prompt
            quality: Output mode, "quality" for quality priority, "speed" for speed priority
            with_audio: Whether to include audio
            size: Video resolution, supports up to 4K (e.g., "3840x2160")
            fps: Frame rate, can be 30 or 60
            max_wait_time: Maximum wait time (seconds)

        Returns:
            Video generation result
        """
        try:
            # Submit generation task
            print('Submitting video generation task...')
            response = self.client.videos.generations(
                model='cogvideox-2',
                prompt=prompt,
                quality=quality,
                with_audio=with_audio,
            )

            print(f'Task submitted successfully, Task ID: {response.id}')
            print(f'Initial response: {response}')

            # Asynchronously wait for task completion
            task_id = response.id
            start_time = time.time()

            while True:
                # Check for timeout
                if time.time() - start_time > max_wait_time:
                    raise TimeoutError(f'Video generation timeout, exceeded {max_wait_time} seconds')

                # Get task result
                print(f'Querying task status... (waited {int(time.time() - start_time)} seconds)')
                result = self.client.videos.retrieve_videos_result(id=task_id)

                print(f'Task status: {result}')

                # Check if task is completed
                if hasattr(result, 'task_status'):
                    if result.task_status == 'SUCCESS':
                        print('Video generation completed!')
                        return result
                    elif result.task_status == 'FAIL':
                        raise Exception(f'Video generation failed: {result}')
                    elif result.task_status in ['PROCESSING', 'SUBMITTED']:
                        print(f'Task in progress, status: {result.task_status}')
                    else:
                        print(f'Unknown status: {result.task_status}')
                else:
                    # If no task_status field, check other possible status fields
                    print(f'Checking result structure: {result}')

                # Wait 5 seconds before querying again
                await asyncio.sleep(5)

        except Exception as e:
            print(f'Error occurred during video generation: {e}')
            raise


# Usage example
async def main():
    generator = VideoGenerator()
    try:
        result = await generator.video_generate(
            prompt='A beautiful sunset beach scene with a beautiful girl',
            quality='quality',
            with_audio=True,
        )
        print('\n=== Final Result ===')
        print(result)

    except Exception as e:
        print(f'Generation failed: {e}')


if __name__ == '__main__':
    asyncio.run(main())