# complete_workflow.py

import os
import json
import random
import subprocess
import requests
import base64
import argparse
from PIL import Image
from gtts import gTTS
from moviepy.editor import ImageClip, AudioFileClip
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

class Config:
    ROOT_DIR = '/root/mm/app'
    ANSWERS_PATH = os.path.join(ROOT_DIR, 'answers.json')
    PROMPT_PATH = os.path.join(ROOT_DIR, 'prompt.json')
    IMAGE_DIR = os.path.join(ROOT_DIR, 'images')
    IMAGE_PATH = os.path.join(IMAGE_DIR, 'output_image.png')
    AUDIO_DIR = os.path.join(ROOT_DIR, 'audio')
    NARRATIVE_PATH = os.path.join(AUDIO_DIR, 'narrative.txt')
    AUDIO_OUTPUT_PATH = os.path.join(AUDIO_DIR, 'output_audio.mp3')
    VIDEO_OUTPUT_PATH = os.path.join(AUDIO_DIR, 'output_video.mp4')
    BASE_LOGO_PATH = os.path.join(IMAGE_DIR, 'logo.png')
    CASCADING_LOGO_PATH = os.path.join(IMAGE_DIR, 'cascading_logo.gif')

    STYLES_LIST = [
        "Minimalist Style", "Flat Design", "Vintage/Retro Style",
        "Bold & Vibrant", "Photographic Style", "Illustrative Style",
        "Typography-Focused", "Monochrome", "Surreal/Abstract",
        "Collage", "Grunge", "3D Design", "Futuristic/Techno",
        "Pop Art", "Geometric Patterns", "Cinematic/Film Noir",
        "Hand-Drawn Doodles", "Pastel Aesthetic", "Urban/Street Style"
    ]

    @classmethod
    def ensure_directories(cls):
        for directory in [cls.IMAGE_DIR, cls.AUDIO_DIR]:
            os.makedirs(directory, exist_ok=True)

class LogoGenerator:
    def __init__(self):
        self.base_logo_path = Config.BASE_LOGO_PATH
        self.output_path = Config.CASCADING_LOGO_PATH

    def create_cascading_logo(self):
        """Create a cascading color-changing logo."""
        try:
            if not os.path.exists(self.base_logo_path):
                raise FileNotFoundError(f"Base logo not found at {self.base_logo_path}")

            base_logo = Image.open(self.base_logo_path)
            frames = []

            colors = [
                (255, 0, 0, 128),    # Red with alpha
                (0, 255, 0, 128),    # Green with alpha
                (0, 0, 255, 128),    # Blue with alpha
                (255, 255, 0, 128),  # Yellow with alpha
                (255, 0, 255, 128),  # Magenta with alpha
                (0, 255, 255, 128)   # Cyan with alpha
            ]

            for color in colors:
                frame = base_logo.copy()
                frame = frame.convert("RGBA")
                overlay = Image.new('RGBA', frame.size, color=color)
                frame = Image.blend(frame, overlay, alpha=0.5)
                frames.append(frame)

            os.makedirs(os.path.dirname(self.output_path), exist_ok=True)

            frames[0].save(
                self.output_path,
                save_all=True,
                append_images=frames[1:],
                duration=200,
                loop=0,
                optimize=True
            )
            print(f"Cascading logo saved to {self.output_path}")
            return True
        except Exception as e:
            print(f"Error creating cascading logo: {e}")
            return False

class WorkflowManager:
    def __init__(self, api_key, test_mode=False):
        Config.ensure_directories()
        self.api_key = api_key
        self.test_mode = test_mode

    def ask_questions(self):
        if self.test_mode:
            answers = {
                "main_subject": "A cat holding a sign that says hello world",
                "environment": "At home",
                "mood": "Happy",
                "style": random.choice(Config.STYLES_LIST)
            }
            print("Test mode: Using auto-filled answers")
        else:
            answers = {
                "main_subject": input("What is the main subject of your image? "),
                "environment": input("Where is the subject? (e.g., at home, at work, etc.) "),
                "mood": input("What is the mood of your image? "),
                "style": self.get_style_from_user()
            }
        
        with open(Config.ANSWERS_PATH, 'w') as f:
            json.dump(answers, f, indent=4)
        print(f"Answers saved to {Config.ANSWERS_PATH}")
        return answers

    def get_style_from_user(self):
        print("\nAvailable Styles:")
        for i, style in enumerate(Config.STYLES_LIST, 1):
            print(f"{i}. {style}")
        
        while True:
            try:
                choice = int(input("\nChoose a style number: "))
                if 1 <= choice <= len(Config.STYLES_LIST):
                    return Config.STYLES_LIST[choice - 1]
            except ValueError:
                pass
            print("Invalid choice. Please try again.")

    def generate_prompt(self):
        with open(Config.ANSWERS_PATH, 'r') as f:
            answers = json.load(f)
        
        prompt_data = {
            "prompt": f"{answers['main_subject']} in {answers['environment']} with a {answers['mood']} mood.",
            "guidance_scale": round(random.uniform(2.5, 5.0), 2),
            "num_inference_steps": 50,
            "height": 1024,
            "width": 1024,
            "style": answers['style']
        }
        
        with open(Config.PROMPT_PATH, 'w') as f:
            json.dump(prompt_data, f, indent=4)
        print(f"Prompt saved to {Config.PROMPT_PATH}")
        return prompt_data

    def generate_image(self):
        print("Generating image...")
        with open(Config.PROMPT_PATH, 'r') as f:
            prompt_data = json.load(f)
        
        headers = {"Authorization": f"Bearer {self.api_key}"}
        response = requests.post(
            "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0",
            headers=headers,
            json=prompt_data
        )
        
        if response.status_code == 200:
            with open(Config.IMAGE_PATH, 'wb') as f:
                f.write(response.content)
            print(f"Image saved to {Config.IMAGE_PATH}")
            return True
        print(f"Failed to generate image: {response.text}")
        return False

    def analyze_image_and_generate_narrative(self):
        print("Analyzing image and generating narrative...")
        try:
            with open(Config.IMAGE_PATH, 'rb') as image_file:
                image_data = base64.b64encode(image_file.read()).decode('utf-8')
            
            headers = {"Authorization": f"Bearer {self.api_key}"}
            response = requests.post(
                "https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-large",
                headers=headers,
                json={
                    "image": image_data,
                    "parameters": {
                        "max_length": 200,
                        "num_return_sequences": 1
                    }
                }
            )
            
            if response.status_code == 200:
                caption = response.json()[0]['generated_text']
                narrative = f"""In this fascinating image, {caption} The scene perfectly captures 
                the essence of the moment, inviting viewers to explore every detail and imagine 
                the story behind it. As we look closer, we can sense the {self.get_answers()['mood']} 
                mood that permeates the entire composition, creating an atmosphere that's both 
                engaging and memorable."""
                
                with open(Config.NARRATIVE_PATH, 'w') as f:
                    f.write(narrative)
                print(f"Narrative saved to {Config.NARRATIVE_PATH}")
                return narrative
            else:
                raise Exception(f"Failed to generate narrative: {response.text}")
        except Exception as e:
            print(f"Error during narrative generation: {e}")
            return None

    def get_answers(self):
        with open(Config.ANSWERS_PATH, 'r') as f:
            return json.load(f)

    def generate_audio_from_narrative(self, narrative):
        print("Generating audio from narrative...")
        try:
            tts = gTTS(text=narrative, lang='en')
            tts.save(Config.AUDIO_OUTPUT_PATH)
            
            # Create video with audio
            image_clip = ImageClip(Config.IMAGE_PATH)
            audio_clip = AudioFileClip(Config.AUDIO_OUTPUT_PATH)
            video_clip = image_clip.set_duration(audio_clip.duration).set_audio(audio_clip)
            video_clip.write_videofile(Config.VIDEO_OUTPUT_PATH, fps=24)
            print(f"Video saved to {Config.VIDEO_OUTPUT_PATH}")
            return True
        except Exception as e:
            print(f"Error during audio/video generation: {e}")
            return False

    def generate_logo(self):
        """Generate the cascading logo."""
        print("Generating cascading logo...")
        logo_gen = LogoGenerator()
        return logo_gen.create_cascading_logo()

def parse_arguments():
    parser = argparse.ArgumentParser(description="Image generation and processing workflow")
    parser.add_argument('-t', '--test', action='store_true', help='Run in test mode with auto-filled answers')
    return parser.parse_args()

def main():
    args = parse_arguments()
    
    api_key = os.getenv('HF_API_KEY')
    if not api_key:
        print("Error: HF_API_KEY not found in environment variables")
        return

    workflow = WorkflowManager(api_key, test_mode=args.test)
    
    try:
        # Step 1: Ask questions and save answers
        workflow.ask_questions()
        
        # Step 2: Generate prompt
        workflow.generate_prompt()
        
        # Step 3: Generate image
        if workflow.generate_image():
            # Step 4: Analyze image and generate narrative
            narrative = workflow.analyze_image_and_generate_narrative()
            
            if narrative:
                # Step 5: Generate audio from narrative and create video
                workflow.generate_audio_from_narrative(narrative)
            
            # Step 6: Generate cascading logo
            workflow.generate_logo()
        
    except Exception as e:
        print(f"An error occurred during the workflow: {e}")

if __name__ == "__main__":
    main()