import streamlit as st
import os
import json
from datetime import datetime
from crewai import Agent, Task, Crew
from langchain_openai import OpenAI
from openai import OpenAI as OpenAIClient
import logging
import requests
import replicate
import random
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO
import textwrap

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Constants for directory paths
BASE_DIR = "data"
QUOTES_DIR = os.path.join(BASE_DIR, "quotes")
IMAGES_DIR = os.path.join(BASE_DIR, "images")
FINAL_IMAGES_DIR = os.path.join(BASE_DIR, "final_images")
MUSIC_DIR = os.path.join(BASE_DIR, "music")  # New directory for storing music files

# Ensure the directories exist
try:
    os.makedirs(QUOTES_DIR, exist_ok=True)
    os.makedirs(IMAGES_DIR, exist_ok=True)
    os.makedirs(FINAL_IMAGES_DIR, exist_ok=True)
    os.makedirs(MUSIC_DIR, exist_ok=True)  # Create music directory
    logger.info(f"Directories successfully created or verified: {BASE_DIR}, {QUOTES_DIR}, {IMAGES_DIR}, {FINAL_IMAGES_DIR}, {MUSIC_DIR}")
except Exception as e:
    logger.error(f"Error creating directories: {str(e)}")

def log_message(message, type="info"):
    if 'process_logs' not in st.session_state:
        st.session_state.process_logs = []
    timestamp = datetime.now().strftime("%H:%M:%S")
    st.session_state.process_logs.append({
        "timestamp": timestamp,
        "message": message,
        "type": type
    })

class SpiritualThemes:
    TRADITIONS = {
        "Eastern": ["Buddhism", "Hinduism", "Taoism", "Zen", "Vedanta", "Yoga Philosophy"],
        "Western": ["Christian Mysticism", "Sufi Wisdom", "Kabbalah", "Greek Philosophy"],
        "Indigenous": ["Native American Wisdom", "Aboriginal Spirituality", "African Traditional"],
        "Modern": ["Mindfulness", "Contemporary Philosophy", "Integrative Spirituality"]
    }
    THEMES = ["Inner Peace", "Mindfulness", "Compassion", "Wisdom", "Self-Discovery", 
              "Unity", "Balance", "Transformation", "Enlightenment", "Purpose"]
    QUOTE_STYLES = ["contemplative", "inspirational", "philosophical", "practical wisdom", 
                   "mystical insight", "ancient teaching", "modern interpretation"]

class ContentGenerator:
    def __init__(self, openai_api_key, replicate_api_key):
        self.llm = OpenAI(api_key=openai_api_key, temperature=0.7)
        self.openai_client = OpenAIClient(api_key=openai_api_key)
        self.replicate_client = replicate.Client(api_token=replicate_api_key)

    def _generate_prompt_variation(self):
        tradition_category = random.choice(list(SpiritualThemes.TRADITIONS.keys()))
        themes = random.sample(SpiritualThemes.THEMES, 2)
        return f"""Generate a spiritual quote from {tradition_category} wisdom tradition.
        Focus on themes of {themes[0]} and {themes[1]}.
        Make it profound yet accessible.
        Format: Return only the quote text on first line, then author name on second line."""

    def _get_fallback_quote(self):
        fallback_quotes = [
            {"text": "Be the change you wish to see in the world.",
             "author": "Mahatma Gandhi", "tradition": "Modern Spirituality"},
            {"text": "Peace comes from within. Do not seek it without.",
             "author": "Buddha", "tradition": "Buddhism"},
            {"text": "Love is the bridge between you and everything.",
             "author": "Rumi", "tradition": "Sufi Wisdom"}
        ]
        quote = random.choice(fallback_quotes)
        quote.update({
            "theme": random.choice(SpiritualThemes.THEMES),
            "generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "source": "fallback"
        })
        log_message("ℹ️ Using fallback quote system", "info")
        return quote

    def generate_quote(self):
        try:
            log_message("🎯 Starting quote generation")
            response = self.openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{
                    "role": "system",
                    "content": "Generate a spiritual quote with attribution."
                }, {
                    "role": "user",
                    "content": self._generate_prompt_variation()
                }],
                temperature=0.7
            )
            
            text = response.choices[0].message.content
            lines = text.split('\n')
            quote_text = lines[0] if lines else ""
            author = lines[1].replace('-', '').strip() if len(lines) > 1 else "Unknown"
            tradition = random.choice(list(SpiritualThemes.TRADITIONS.keys()))
            theme = random.choice(SpiritualThemes.THEMES)
            
            result = {
                "text": quote_text,
                "author": author,
                "tradition": tradition,
                "theme": theme,
                "generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }
            
            # Save the quote to a file
            filename = f"quote_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
            file_path = os.path.join(QUOTES_DIR, filename)
            try:
                with open(file_path, 'w') as f:
                    json.dump(result, f, indent=4)
                log_message(f"✨ Generated quote saved at {file_path}")
                logger.info(f"Quote saved successfully at: {file_path}")
            except Exception as e:
                log_message(f"❌ Failed to save quote to file: {str(e)}", "error")
                logger.error(f"Error saving quote: {str(e)}")
            
            return result
            
        except Exception as e:
            log_message(f"❌ Quote generation failed: {str(e)}")
            return self._get_fallback_quote()

    def generate_image(self, quote, tradition, theme):
        try:
            log_message("🎨 Starting image generation")
            style_prompts = {
                "Buddhism": "zen minimalism, misty mountains, lotus flowers",
                "Hinduism": "mandala patterns, sacred geometry, vibrant cosmos",
                "Taoism": "flowing water, yin-yang harmony, natural elements",
                "default": "serene minimalism, sacred geometry, natural elements"
            }
            style = style_prompts.get(tradition, style_prompts["default"])
            prompt = f"""Create a spiritual and contemplative image inspired by this theme: '{theme}'
            Style inspiration: {style}
            Mood: serene, inspiring, contemplative
            Requirements: Suitable for Instagram, space for text overlay, balanced composition"""
            
            response = self.openai_client.images.generate(
                model="dall-e-3",
                prompt=prompt,
                size="1024x1024",
                quality="standard",
                n=1
            )
            image_url = response.data[0].url

            # Download and save the generated image
            image_response = requests.get(image_url)
            if image_response.status_code == 200:
                filename = f"image_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
                file_path = os.path.join(IMAGES_DIR, filename)
                try:
                    with open(file_path, 'wb') as f:
                        f.write(image_response.content)
                    log_message(f"✨ Image generated and saved at {file_path}")
                    logger.info(f"Image saved successfully at: {file_path}")
                except Exception as e:
                    log_message(f"❌ Failed to save image to file: {str(e)}", "error")
                    logger.error(f"Error saving image: {str(e)}")
            else:
                raise Exception("Failed to download image")

            return image_url
        except Exception as e:
            log_message(f"❌ Image generation failed: {str(e)}", "error")
            logger.error(f"Image generation failed: {str(e)}")
            raise

    def generate_audio(self, tradition, theme):
        try:
            log_message("🎵 Starting audio generation")
            music_prompts = {
                "Buddhism": "Tibetan singing bowls and gentle bells with peaceful ambient drones",
                "Hinduism": "Indian bansuri flute with gentle tabla rhythms",
                "Modern Spirituality": "Ambient synthesizer with crystal bowls",
                "default": "Peaceful ambient music with gentle bells"
            }
            prompt = music_prompts.get(tradition, music_prompts["default"])
            
            output = self.replicate_client.run(
                "meta/musicgen:671ac645ce5e552cc63a54a2bbff63fcf798043055d2dac5fc9e36a837eedcfb",
                input={
                    "prompt": prompt,
                    "model_version": "stereo-large",
                    "output_format": "mp3",
                    "duration": 15
                }
            )
            audio_url = output if isinstance(output, str) else str(output)
            
            # Download and save the generated music
            music_response = requests.get(audio_url)
            if music_response.status_code == 200:
                filename = f"music_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp3"
                file_path = os.path.join(MUSIC_DIR, filename)
                try:
                    with open(file_path, 'wb') as f:
                        f.write(music_response.content)
                    log_message(f"✨ Music generated and saved at {file_path}")
                    logger.info(f"Music saved successfully at: {file_path}")
                except Exception as e:
                    log_message(f"❌ Failed to save music to file: {str(e)}", "error")
                    logger.error(f"Error saving music: {str(e)}")
            else:
                raise Exception("Failed to download audio")

            return audio_url
        except Exception as e:
            log_message(f"❌ Music generation failed: {str(e)}", "error")
            logger.error(f"Music generation failed: {str(e)}")
            raise

    def create_image_with_quote(self, image_url, quote_text, author):
        try:
            response = requests.get(image_url)
            img = Image.open(BytesIO(response.content))
            width, height = img.size
    
            # Create blank image for text overlay
            txt = Image.new('RGBA', (width, height), (255, 255, 255, 0))
            d = ImageDraw.Draw(txt)
    
            # Semi-transparent background for text
            d.rectangle([0, height / 2 - 100, width, height / 2 + 100], fill=(0, 0, 0, 127))
    
            try:
                quote_font = ImageFont.truetype("DejaVuSans-Bold", 40)
                author_font = ImageFont.truetype("DejaVuSans", 30)
            except:
                quote_font = author_font = ImageFont.load_default()
    
            # Word wrap the quote text
            max_width = width - 40  # Leave some padding on both sides
            lines = textwrap.wrap(quote_text, width=40)
    
            # Draw each line of the wrapped text
            y_offset = height / 2 - 100  # Starting Y position for the quote
            for line in lines:
                # Calculate the text size using textbbox
                bbox = d.textbbox((0, 0), line, font=quote_font)
                line_width = bbox[2] - bbox[0]
                line_height = bbox[3] - bbox[1]
    
                x_position = (width - line_width) / 2  # Center the line horizontally
                d.text((x_position, y_offset), line, font=quote_font, fill=(255, 255, 255, 255))
                y_offset += line_height + 5  # Move to the next line, with a small line gap
    
            # Draw the author text below the quote
            author_y_position = y_offset + 20  # Leave some space below the quote for the author
            author_text = f"- {author}"
            author_bbox = d.textbbox((0, 0), author_text, font=author_font)
            author_width = author_bbox[2] - author_bbox[0]
            author_x_position = (width - author_width) / 2
            d.text((author_x_position, author_y_position), author_text, font=author_font, fill=(255, 255, 255, 255))
    
            # Combine the original image with the overlay
            out = Image.alpha_composite(img.convert('RGBA'), txt)
    
            # Convert the image to bytes and save to a file
            buffer = BytesIO()
            out.convert('RGB').save(buffer, format='JPEG')
            final_image_bytes = buffer.getvalue()
            
            filename = f"final_image_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
            file_path = os.path.join(FINAL_IMAGES_DIR, filename)
            try:
                with open(file_path, 'wb') as f:
                    f.write(final_image_bytes)
                log_message(f"✨ Final image saved at {file_path}")
                logger.info(f"Final image saved successfully at: {file_path}")
            except Exception as e:
                log_message(f"❌ Failed to save final image to file: {str(e)}", "error")
                logger.error(f"Error saving final image: {str(e)}")
            
            return final_image_bytes
    
        except Exception as e:
            log_message(f"Error creating overlay: {str(e)}", "error")
            logger.error(f"Error creating image with quote: {str(e)}")
            return None

def generate_hashtags(tradition, theme):
    base_tags = ["#spirituality", "#mindfulness", "#wisdom", "#inspiration"]
    tradition_tags = {
        "Buddhism": ["#buddhism", "#zen", "#buddhist"],
        "Hinduism": ["#hinduism", "#vedanta", "#yoga"],
        "Taoism": ["#taoism", "#tao", "#balance"],
        "Modern Spirituality": ["#spiritual", "#awakening", "#growth"]
    }
    theme_tags = {
        "Inner Peace": ["#peace", "#calm", "#serenity"],
        "Mindfulness": ["#mindful", "#present", "#awareness"],
        "Wisdom": ["#wisdom", "#truth", "#knowledge"]
    }
    custom_tags = tradition_tags.get(tradition, []) + theme_tags.get(theme, [])
    return list(set(base_tags + custom_tags))[:10]

def main():
    st.set_page_config(page_title="Spiritual Content Generator", page_icon="🕉️", layout="wide")

    # Initialize session states
    if 'process_logs' not in st.session_state:
        st.session_state.process_logs = []
    for state_var in ['generated_content', 'generated_image', 'generated_audio']:
        if state_var not in st.session_state:
            st.session_state[state_var] = None

    if not all(key in st.secrets for key in ['OPENAI_API_KEY', 'REPLICATE_API_TOKEN']):
        st.error("Missing API keys in secrets!")
        st.stop()

    # Create columns for content and debug log
    col1, col2 = st.columns([2, 1])

    with col1:
        st.title("🕉️ Spiritual Content Generator")
        
        generator = ContentGenerator(st.secrets["OPENAI_API_KEY"], st.secrets["REPLICATE_API_TOKEN"])

        # Step 1: Generate Quote
        if st.button("Generate Quote"):
            try:
                quote_data = generator.generate_quote()
                st.session_state.generated_content = quote_data
                st.success("✨ Quote generated!")
                st.json(quote_data)
            except Exception as e:
                st.error(f"Error generating quote: {str(e)}")

        # Step 2: Create Image with Quote
        if st.session_state.generated_content and st.button("Create Image"):
            try:
                image_url = generator.generate_image(
                    st.session_state.generated_content["text"],
                    st.session_state.generated_content["tradition"],
                    st.session_state.generated_content["theme"]
                )
                st.session_state.generated_image = image_url
                st.success("✨ Image created!")
                st.image(image_url)
            except Exception as e:
                st.error(f"Error generating image: {str(e)}")

        # Step 3: Generate Music
        if st.session_state.generated_content and st.button("Generate Music"):
            try:
                audio_url = generator.generate_audio(
                    st.session_state.generated_content["tradition"],
                    st.session_state.generated_content["theme"]
                )
                st.session_state.generated_audio = audio_url
                st.success("✨ Music generated!")
                audio_response = requests.get(audio_url)
                if audio_response.status_code == 200:
                    st.audio(audio_response.content, format='audio/mp3')
            except Exception as e:
                st.error(f"Error generating audio: {str(e)}")

        # Step 4: Preview Quote on Image
        if st.session_state.generated_content and st.session_state.generated_image:
            st.subheader("Final Preview")
            final_image = generator.create_image_with_quote(
                st.session_state.generated_image,
                st.session_state.generated_content['text'],
                st.session_state.generated_content['author']
            )
            if final_image:
                st.image(final_image)
                st.download_button("📥 Download Image", final_image,
                                   "spiritual_quote.jpg", "image/jpeg")

            # Generate Instagram Caption
            hashtags = generate_hashtags(
                st.session_state.generated_content['tradition'],
                st.session_state.generated_content.get('theme', 'Wisdom')
            )
            caption = f"""✨ Wisdom from {st.session_state.generated_content['tradition']} ✨

"{st.session_state.generated_content['text']}"
- {st.session_state.generated_content['author']}

{' '.join(hashtags)}"""
            st.text_area("Instagram Caption", caption, height=200)

    # Live debug log
    with col2:
        st.markdown("### 🔍 Live Debug Log")
        placeholder = st.empty()
        
        def update_logs():
            with placeholder.container():
                for log in st.session_state.process_logs:
                    if log["type"] == "error":
                        st.error(f"{log['timestamp']} - {log['message']}")
                    else:
                        st.info(f"{log['timestamp']} - {log['message']}")
        
        update_logs()

if __name__ == "__main__":
    main()