import streamlit as st
from sentence_transformers import SentenceTransformer
import torch
import json
from streamlit_lottie import st_lottie
from notion_client import Client
from datetime import datetime

# Set page config at the very beginning
st.set_page_config(page_title="Prompt easz", layout="centered")

# Initialize Notion client
notion=Client(auth=st.secrets["NOTION_API_KEY"])
NOTION_DATABASE_ID = st.secrets["notiondb"]

# Custom CSS to increase text size, beautify the app, and highlight the final message
st.markdown("""
    <style>
    .big-font {
        font-size:60px !important;
        text-align: center;
    }
    .slider-label {
        font-size:25px !important;
        font-weight: bold;
    }
    .small-text {
        font-size:12px !important;
    }
    .medium-text {
        font-size:16px !important;
    }
    .center-text {
        text-align: center;
    }
    .highlight {
        font-family: 'Courier New', Courier, monospace;
        background-color: #f0f0f0;
        padding: 10px;
        border-radius: 5px;
    }
    .stTextInput > div > div > input, .stSelectbox > div > div > select, .stButton > button { font-size: 16px; }
    .stMarkdown { font-size: 14px; }
    </style>
    """, unsafe_allow_html=True)

@st.cache_resource
def load_model():
    return SentenceTransformer('all-MiniLM-L6-v2')

@st.cache_data
def load_prompts():
    try:
        with open('prompts.json', 'r') as file:
            return json.load(file)
    except json.JSONDecodeError as e:
        st.error(f"Error loading prompts.json: {e}")
        return {}

@st.cache_data
def load_categories():
    try:
        with open('categories.json', 'r') as file:
            return json.load(file)
    except json.JSONDecodeError as e:
        st.error(f"Error loading categories.json: {e}")
        return {}

def load_lottiefile(filepath: str):
    with open(filepath, "r") as f:
        return json.load(f)

def find_prompt_types(user_input, num_types=3):
    model = load_model()
    prompt_types = load_prompts()
    user_embedding = model.encode(user_input, convert_to_tensor=True)
    similarities = {}
    for pt, data in prompt_types.items():
        all_keywords = " ".join([" ".join(variation["keywords"]) for variation in data])
        pt_embedding = model.encode(all_keywords, convert_to_tensor=True)
        similarity = torch.cosine_similarity(user_embedding, pt_embedding, dim=0).item()
        similarities[pt] = similarity
    return sorted(similarities, key=similarities.get, reverse=True)[:num_types]

def generate_prompt(prompt_type, topic, category, complexity, topic2=None, constraint=None):
    model = load_model()
    prompt_types = load_prompts()
    categories_info = load_categories()
    template = prompt_types[prompt_type][0]["template"]
    complexity_instructions = {
        "low": "Use simple language and basic concepts. Avoid technical jargon and provide elementary explanations suitable for beginners.",
        "medium": "Include more detailed explanations and introduce some field-specific terminology. Provide a balance between depth and accessibility.",
        "high": "Delve into advanced concepts and use specialized terminology. Assume a high level of prior knowledge and provide in-depth analysis."
    }
    category_info = categories_info.get(category, {})
    analogy = category_info.get("analogy", "analogous")
    expert_role = category_info.get("expert_role", "an expert")
    related_topics = ", ".join(category_info.get("related_topics", []))
    return f"""
    # {prompt_type} Prompt: {topic} in {category}

    {template.format(topic=topic, category=category, topic2=topic2, constraint=constraint)}

    **Complexity: {complexity.capitalize()}**
    {complexity_instructions[complexity]}

    **Category Details:**
    - **Analogy:** {analogy}
    - **Expert Role:** {expert_role}
    - **Related Topics:** {related_topics}

    **Additional Instructions:**
    - Ensure your response is well-structured and easy to follow.
    - Use relevant examples to illustrate key points.
    - Tailor your language to the specified complexity level.
    """.strip()

def save_feedback_to_notion(user_prompt, improvement_suggestion, timestamp):
    try:
        notion.pages.create(
            parent={"database_id": NOTION_DATABASE_ID},
            properties={
                "User Prompt": {"title": [{"text": {"content": user_prompt}}]},
                "Improvement Suggestion": {"rich_text": [{"text": {"content": improvement_suggestion}}]}
            }
        )
        return True
    except Exception as e:
        # Log the error to the console instead of displaying it in the Streamlit app
        print(f"Error saving to Notion: {str(e)}")
        return False

def main():
    # Display the title and animation
    st.markdown('<h1 class="big-font">Prompt easz</h1>', unsafe_allow_html=True)
    st.markdown('<h2 class="center-text">AI Prompt Generator</h2>', unsafe_allow_html=True)
    
    # Load and display Lottie animation
    lottie_animation = load_lottiefile("animation.json")  # Replace with your Lottie file path
    col1, col2, col3 = st.columns([1,2,1])
    with col2:
        st_lottie(lottie_animation, speed=1, height=400, key="initial")

    # Add About link with icon below the Lottie animation
    col1, col2, col3 = st.columns([1,2,1])
    with col2:
        st.markdown(
            '<div style="text-align: center;">'
            '<a href="/About" target="_blank" style="text-decoration: none;">'
            '<span style="font-size: 24px;">ℹ️</span> '
            '<span style="vertical-align: middle;">About</span>'
            '</a>'
            '</div>',
            unsafe_allow_html=True
        )

    user_input = st.text_input("What would you like a prompt for?")

    if user_input:
        suggested_types = find_prompt_types(user_input)
        prompt_types = load_prompts()
        categories_info = load_categories()

        unique_types = list(set(list(prompt_types.keys()) + suggested_types))
        
        topic = st.text_input("📌 Main Topic:", value=user_input)
        category = st.selectbox("🏷️ Category or Field:", options=list(categories_info.keys()), index=0)
        complexity = st.select_slider("🧠 Complexity Level:", options=["low", "medium", "high"], value="medium")
        
        st.write("Choose Your Prompt")
        
        selected_tab = st.selectbox("Select Prompt Type:", unique_types)

        if selected_tab in prompt_types:
            topic2 = st.text_input(f"📌 Second Topic for {selected_tab}:", key=f"{selected_tab}_topic2") if selected_tab == "Comparative Analysis" else None
            constraint = st.number_input(f"🔢 Word Limit for {selected_tab}:", min_value=10, max_value=500, value=50, key=f"{selected_tab}_constraint") if selected_tab == "Constraint Addition" else None
            
            prompt = generate_prompt(selected_tab, topic, category, complexity, topic2, constraint)
            st.markdown(prompt)

            # Add feedback collection section
            st.subheader("Provide Feedback")
            improvement_suggestion = st.text_area("What can be improved in this prompt?", height=100)
            
            if st.button("Submit Feedback"):
                if improvement_suggestion:
                    timestamp = datetime.now().isoformat()
                    if save_feedback_to_notion(prompt, improvement_suggestion, timestamp):
                        st.success("Feedback saved to Notion successfully!")
                    else:
                        st.error("Failed to save feedback to Notion.")
                else:
                    st.warning("Please provide feedback before submitting.")

    st.markdown("---")
    st.markdown("""
    ### 🌟 How to Use AI Prompt Wizard Ultimate
    1. Enter your desired topic or question.
    2. Adjust the main topic, category, and complexity if needed.
    3. Select the prompt type from the dropdown.
    4. View the generated prompt.
    5. Provide feedback on what can be improved in the prompt.
    6. Submit your feedback to help us enhance our prompt generation.
    """)

if __name__ == "__main__":
    main()