import os
import re
import yaml
import datetime
import streamlit as st
import streamlit_chat as stc
from huggingface_hub import InferenceClient

# Configuration and constants
CONFIG_DIRECTORY = os.path.dirname(__file__).replace('\\', '/')  # Normalize directory path
CONFIG_FILES = [os.path.join(CONFIG_DIRECTORY, "config/config.yaml")]
OUTPUT_DIRECTORY = os.path.join(CONFIG_DIRECTORY, "generated_scripts")

# Ensure the output directory exists
os.makedirs(OUTPUT_DIRECTORY, exist_ok=True)

class ConfigurationManager:
    def __init__(self, config_directory, config_files):
        self.config_directory = config_directory
        self.config_files = config_files
        self.configs = {}

    def load_configs(self):
        for file in self.config_files:
            with open(file, 'r') as f:
                config = yaml.safe_load(f)
                config_name = os.path.splitext(os.path.basename(file))[0]
                self.configs[config_name] = config

    def get_full_context(self):
        return self.configs.get('config', {}).get('context', {}).get('context', '')

def initialize_ai_client(model_name):
    return InferenceClient(model_name)

def clean_script_content(scripts):
    clean_scripts = []
    for script in scripts:
        clean_scripts.append(re.sub(r'^.*?:\s(.*)$\n', r'\1', script))
    return clean_scripts

def extract_script_description(script):
    if not script:
        return "script"
    return script.split('\n')[0]

def generate_script_filename(description):
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{timestamp}_{description.replace(' ', '_')}.py"
    return os.path.join(OUTPUT_DIRECTORY, filename)

def chat_with_ai(question, client, max_tokens, context):
    prompt = f"context: {context}\nquestion: {question}\n"
    try:
        response = client(prompt, max_length=max_tokens)
        return response.get('generated_text', "").split('\n')
    except Exception as e:
        st.error(f"Error in AI response: {e}")
        return []

def main():
    st.title("Python Script Generator with Chat")
    st.write("Generate Python scripts based on your prompts using an AI chat interface.")

    # Load configurations
    config_manager = ConfigurationManager(CONFIG_DIRECTORY, CONFIG_FILES)
    config_manager.load_configs()
    config = config_manager.configs.get('config', {})

    if not config:
        st.error("Missing 'config.yaml' configurations.")
        return

    model_name = config.get('ai', {}).get('model', '')
    max_tokens = config.get('ai', {}).get('max_tokens', 100)

    # Initialize AI Client
    client = initialize_ai_client(model_name)

    # Streamlit Chat History
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []

    # User input and chat interface
    user_input = st.text_input("Enter your prompt:")
    if st.button("Send"):
        if user_input.strip() == "":
            st.error("Please enter a valid prompt.")
        else:
            # Display user input in chat
            st.session_state.chat_history.append({"role": "user", "content": user_input})

            # Get context and AI response
            context = config_manager.get_full_context()
            response = chat_with_ai(user_input, client, max_tokens, context)
            cleaned_scripts = clean_script_content(response)

            if cleaned_scripts:
                ai_reply = cleaned_scripts[0]
                st.session_state.chat_history.append({"role": "assistant", "content": ai_reply})
                
                # Display generated script
                description = extract_script_description(ai_reply)
                filename = generate_script_filename(description)
                
                # Save the script to a file
                try:
                    with open(filename, 'w') as script_file:
                        script_file.write("\n".join(cleaned_scripts))
                    st.success(f"Script saved to: {filename}")
                    st.code("\n".join(cleaned_scripts), language='python')
                except Exception as e:
                    st.error(f"Failed to save script: {e}")
            else:
                st.write("No valid script content generated.")

    # Display chat history
    if st.session_state.chat_history:
        for chat in st.session_state.chat_history:
            if chat["role"] == "user":
                stc.message(chat["content"], is_user=True)
            else:
                stc.message(chat["content"], is_user=False)

if __name__ == "__main__":
    main()


make this work in streamlit please using their most recent formatting please