import streamlit as st
from PIL import Image
import yaml
import requests
import os
from datetime import datetime
import json
import secrets
from thaTube import thaTube

folders = {'out':'output',
           'mp3':'mp3',
          'assets':'assets'}
          

# Output folder
output_folder = "output"
for folder in folders.keys():
    os.makedirs(folders[folder], exist_ok=True)

def generate_random_key():
    return secrets.token_hex(8)  # Generates a 16-character hexadecimal string

# Function to load YAML files
def load_yaml_file(file_path):
    with open(file_path, "r") as file:
        return yaml.safe_load(file)

# Initialize app configuration
st.set_page_config(page_title="Multi-tab Image App", layout="wide")

# Placeholder for the image display
image_placeholder = st.empty()

# Load Chatbot system prompts
chatbot_prompts = load_yaml_file("system_prompts.yaml")

# Tabs
tab_names = ["Image Generation", "Chat/Prompt Enhancement", "AIvatar ChatBot", "Story Generation", "Gallery", "System Prompts Editor"]
tabs = st.tabs(tab_names)

# thaTube
playlist_url='https://www.youtube.com/playlist?list=PLUwVJKJgdARRdQ3Y3NFB9HMTjXTa5INcz'
tuber = thaTube(playlist_url)

# Image Generation Tab
with tabs[0]:
    st.header("Image Generation")
    
    # Input fields
    prompt = st.text_input("Prompt")
    seed = st.slider("Seed", min_value=0, max_value=1000000, value=0)
    randomize_seed = st.checkbox("Randomize seed", value=True)
    width = st.slider("Width", min_value=256, max_value=1024, value=1024)
    height = st.slider("Height", min_value=256, max_value=1024, value=1024)
    steps = st.slider("Inference steps", min_value=1, max_value=50, value=4)
    
    # Generate Image button
    if st.button("Generate Image"):
        st.info("Generating image...")
        # Simulate the image generation process
        try:
            # Replace with your own API call
            result = {"image_path": "image.png", "generated_seed": seed}
            
            image_path = result['image_path']
            image = Image.open(image_path)
            image_placeholder.image(image)
            
            # Save image
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            prompt_slug = "_".join(prompt.replace(':', '').replace("'", '').replace('"', '').replace(',', '').replace('.', '').replace('*', '##').strip().split()[:10]).lower()
            filename = f"{timestamp}_{prompt_slug}.png"
            file_path = os.path.join(output_folder, filename)
            image.save(file_path)
            st.success(f"Image saved as {file_path}")
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")

# Chat/Prompt Enhancement Tab
with tabs[1]:
    st.header("Chat/Prompt Enhancement")
    
    # Chat history and user input
    chat_history = st.text_area("Chat History", height=300)
    user_input = st.text_input("Your Input")
    
    # Buttons for sending message or listening
    if st.button("Send"):
        st.info("Sending message...")
        # Simulate sending a message to the chatbot
    if st.button("Listen"):
        st.info("Listening...")

# AIvatar ChatBot Tab
with tabs[2]:
    st.header("AIvatar ChatBot")
    
    # Chat history and user input
    chat_history_ai = st.text_area("Chat History", height=300, key=generate_random_key())
    user_input_ai = st.text_input("Your Input", key=generate_random_key())
    
    # Send button
    if st.button("Send AI Message"):
        st.info("Sending message to AIvatar ChatBot...")
        # Simulate sending a message to the chatbot

# Story Generation Tab
with tabs[3]:
    st.header("Story Generation")
    
    # Input fields for story generation
    story_prompt = st.text_input("Story Prompt", key=generate_random_key())
    story_style = st.text_input("Story Style", key=generate_random_key())
    
    # Generate Story button
    if st.button("Generate Story"):
        st.info("Generating story...")
        # Simulate story generation process
        st.text_area("Story Output", "Generated story will appear here...", height=300, key=generate_random_key())
        
    # Save to PDF button
    if st.button("Save to PDF"):
        st.info("Saving story to PDF...")

# Gallery Tab
with tabs[4]:
    st.header("Gallery")
    
    # Load and display images from the output folder
    images = [f for f in os.listdir(output_folder) if f.endswith(".png")]
    if images:
        st.image([os.path.join(output_folder, img) for img in images], width=150)
    else:
        st.info("No images found in the gallery.")

# System Prompts Editor Tab
with tabs[5]:
    st.header("System Prompts Editor")
    
    # Select and edit prompts
    prompt_keys = list(chatbot_prompts.keys())
    selected_prompt_key = st.selectbox("Select a prompt to edit", prompt_keys)
    selected_prompt_text = st.text_area("Prompt Text", chatbot_prompts[selected_prompt_key], height=200, key=generate_random_key())
    
    # Save prompt button
    if st.button("Save Prompt"):
        chatbot_prompts[selected_prompt_key] = selected_prompt_text
        # Save the updated prompts back to the YAML file
        with open("system_prompts.yaml", "w") as file:
            yaml.safe_dump(chatbot_prompts, file)
        st.success(f"Prompt '{selected_prompt_key}' saved successfully!")