"""
Streamlit app for AI Project Assistant.
"""
import streamlit as st
from pathlib import Path
import os
from datetime import datetime
from dotenv import load_dotenv
from src.rag import ProjectRAG
from src.agent import ProjectAgent
from src.parsers import load_meetings_from_directory
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from langchain_core.messages import SystemMessage, HumanMessage

# Load environment variables
load_dotenv()

# Page config
st.set_page_config(
    page_title="AI Project Assistant",
    page_icon="🤖",
    layout="wide"
)

# Custom CSS
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        font-weight: bold;
        margin-bottom: 1rem;
    }
    .project-card {
        padding: 1rem;
        border-radius: 0.5rem;
        background-color: #f0f2f6;
        margin: 0.5rem 0;
    }
    .action-item {
        padding: 0.5rem;
        margin: 0.25rem 0;
        border-left: 3px solid #1f77b4;
        background-color: #e8f4f8;
    }
    .blocker {
        padding: 0.5rem;
        margin: 0.25rem 0;
        border-left: 3px solid #d62728;
        background-color: #ffe8e8;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'rag' not in st.session_state:
    st.session_state.rag = None
if 'agent' not in st.session_state:
    st.session_state.agent = None
if 'messages' not in st.session_state:
    st.session_state.messages = []
if 'initialized' not in st.session_state:
    st.session_state.initialized = False


def initialize_system():
    """Initialize RAG and Agent systems."""
    data_dir = Path("./data")
    
    if not data_dir.exists():
        data_dir.mkdir(parents=True)
        st.warning("Created data directory. Please add your meeting notes to 'data/project_name/meetings/'")
        return False
    
    with st.spinner("Loading and indexing meetings..."):
        st.session_state.rag = ProjectRAG(data_dir)
        st.session_state.rag.load_and_index()
        
        if not st.session_state.rag.meetings:
            return False
        
        st.session_state.agent = ProjectAgent(st.session_state.rag)
        st.session_state.initialized = True
    
    return True


def main():
    """Main app function."""
    
    # Header
    st.markdown('<div class="main-header">🤖 AI Project Assistant</div>', unsafe_allow_html=True)
    st.markdown("Your intelligent assistant for managing multiple projects through meeting summaries")
    
    # Sidebar
    with st.sidebar:
        st.header("⚙️ Settings")

        # Project filter
        if st.session_state.rag and st.session_state.initialized:
            projects = st.session_state.rag.get_all_projects()
            selected_project = st.selectbox(
                "Select Project",
                options=["All Projects"] + projects,
                key="selected_project"
            )
            st.session_state.project_filter = None if selected_project == "All Projects" else selected_project

        if st.button("🔄 Reload Meetings", use_container_width=True):
            st.session_state.initialized = False
            st.rerun()

        st.divider()

        st.header("📊 Quick Stats")
        if st.session_state.rag and st.session_state.initialized:
            current_filter = st.session_state.get("project_filter")

            if current_filter:
                st.info(f"Showing: **{current_filter}**")
                action_items = st.session_state.rag.get_open_action_items(project=current_filter)
                blockers = st.session_state.rag.get_blockers(project=current_filter)
            else:
                projects = st.session_state.rag.get_all_projects()
                st.metric("Total Projects", len(projects))
                action_items = st.session_state.rag.get_open_action_items()
                blockers = st.session_state.rag.get_blockers()

            st.metric("Total Meetings", len(st.session_state.rag.meetings))
            st.metric("Open Action Items", len(action_items))
            st.metric("Current Blockers", len(blockers))

        st.divider()

        st.header("💡 Example Queries")
        st.markdown("""
        - What are the open action items?
        - What blockers do we have?
        - What decisions were made?
        - What should I focus on next?
        - Summarize the project status
        """)
    
    # Check HF Token (auto-provided on Spaces, optional locally)
    if not os.getenv("HF_TOKEN"):
        st.warning("⚠️ HF_TOKEN not found. Running with limited functionality.")
        st.info("On Spaces: Token is automatically provided")
        st.info("Locally: Set HF_TOKEN in .env file (optional for free API)")
    
    # Initialize system
    if not st.session_state.initialized:
        if not initialize_system():
            st.warning("No meetings found. Add your meeting notes to get started!")
            
            with st.expander("📝 How to add meetings"):
                st.markdown("""
                1. Create a folder structure: `data/your_project_name/meetings/`
                2. Add markdown files with your meeting notes
                3. Use this format:
                
                ```markdown
                # Meeting: Project Kickoff
                Date: 2025-01-15
                Participants: Alice, Bob
                
                ## Discussion
                Your meeting notes here
                
                ## Decisions
                - Decision 1
                - Decision 2
                
                ## Action Items
                - [ ] Alice: Task 1 by Jan 20
                - [x] Bob: Task 2 (completed)
                
                ## Blockers
                - Waiting for approval
                ```
                """)
            return
        
        st.success(f"Loaded {len(st.session_state.rag.meetings)} meetings!")

    # Main content tabs - ONLY 2 TABS
    tab1, tab2 = st.tabs(["💬 Chat", "📤 Upload Meeting"])

    with tab1:
        st.header("💬 Ask Questions About Your Projects")

        # Project selection BEFORE chat
        if st.session_state.rag and st.session_state.initialized:
            projects = st.session_state.rag.get_all_projects()

            st.markdown("### Select a Project to Chat About")

            # Create columns for project buttons
            cols = st.columns(len(projects) + 1)

            # "All Projects" button
            with cols[0]:
                if st.button("🌐 All Projects", use_container_width=True, type="secondary"):
                    st.session_state.selected_chat_project = None
                    st.rerun()

            # Individual project buttons
            for i, project in enumerate(projects, 1):
                with cols[i]:
                    if st.button(f"📁 {project}", use_container_width=True, type="primary"):
                        st.session_state.selected_chat_project = project
                        st.rerun()

            st.divider()

            # Show selected project
            selected_project = st.session_state.get("selected_chat_project")
            if selected_project:
                st.success(f"💬 Chatting about: **{selected_project}**")
            else:
                st.info("💬 Chatting about: **All Projects**")

        # Only show chat if a selection has been made
        if "selected_chat_project" in st.session_state:
            # Display chat messages
            for message in st.session_state.messages:
                with st.chat_message(message["role"]):
                    st.markdown(message["content"])
        else:
            st.warning("👆 Please select a project above to start chatting")
            return

    # Chat input (must be outside tabs/columns/expanders) - only show if project selected
    if tab1 and "selected_chat_project" in st.session_state:
        prompt = st.chat_input("What would you like to know about your projects?")

        # Process chat input
        if prompt:
            # Add user message
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(prompt)

            # Get agent response with project filter
            with st.chat_message("assistant"):
                with st.spinner("Thinking..."):
                    # Add project context to query if specific project selected
                    selected_project = st.session_state.get("selected_chat_project")
                    if selected_project:
                        enhanced_prompt = f"[Project: {selected_project}] {prompt}"
                    else:
                        enhanced_prompt = prompt

                    response = st.session_state.agent.query(enhanced_prompt)
                    st.markdown(response)

            st.session_state.messages.append({"role": "assistant", "content": response})
            st.rerun()

    with tab2:
        st.header("📤 Upload Meeting Notes")
        st.markdown("Upload plain text meeting notes and let AI structure them for you!")

        col1, col2 = st.columns(2)

        with col1:
            project_name = st.text_input("Project Name", placeholder="e.g., mobile_app_redesign")
            meeting_date = st.date_input("Meeting Date", value=datetime.now())
            meeting_title = st.text_input("Meeting Title", placeholder="e.g., Sprint Planning")

        with col2:
            participants = st.text_input("Participants (comma-separated)", placeholder="e.g., Alice, Bob, Charlie")

        st.markdown("### Paste or Upload Meeting Notes")

        # Option 1: Text area
        meeting_text = st.text_area(
            "Paste your meeting notes here (plain text)",
            height=300,
            placeholder="""Example:
We discussed the new feature requirements.
Alice will implement the login page by next Friday.
Bob raised a concern about the database migration.
We decided to use PostgreSQL instead of MySQL.
Charlie is blocked waiting for API credentials.
"""
        )

        # Option 2: File upload
        uploaded_file = st.file_uploader("Or upload a text file", type=['txt', 'md'])

        if uploaded_file is not None:
            meeting_text = uploaded_file.read().decode('utf-8')
            st.info(f"Loaded {len(meeting_text)} characters from file")

        if st.button("🤖 Structure Meeting with AI", type="primary", disabled=not meeting_text or not project_name):
            with st.spinner("AI is structuring your meeting notes..."):
                try:
                    # Use HF Inference API to structure the meeting
                    endpoint = HuggingFaceEndpoint(
                        repo_id="meta-llama/Llama-3.2-3B-Instruct",
                        temperature=0.3,
                        max_new_tokens=1024,
                        huggingfacehub_api_token=os.getenv("HF_TOKEN")
                    )
                    llm = ChatHuggingFace(llm=endpoint)

                        system_prompt = """You are a meeting notes structuring assistant.
Convert unstructured meeting notes into a well-formatted markdown document with these sections:
1. # Meeting: [title]
2. Date: [date]
3. Participants: [list]
4. ## Discussion (key points discussed)
5. ## Decisions (decisions made)
6. ## Action Items (as checkboxes with assignee and deadline if mentioned)
7. ## Blockers (any blockers or issues raised)

Format action items as:
- [ ] Person: Task description by deadline
or
- [ ] Task description (if no person/deadline mentioned)

Extract all relevant information from the raw notes."""

                        user_prompt = f"""Structure these meeting notes:

Raw Notes:
{meeting_text}

Meeting Details:
- Title: {meeting_title or 'Meeting'}
- Date: {meeting_date}
- Participants: {participants or 'Not specified'}
"""

                        messages = [
                            SystemMessage(content=system_prompt),
                            HumanMessage(content=user_prompt)
                        ]

                        response = llm.invoke(messages)
                        structured_md = response.content

                        # Display preview
                        st.success("✅ Meeting structured successfully!")
                        st.markdown("### Preview")
                        st.markdown(structured_md)

                        # Save option
                        st.markdown("### Save Meeting")

                        save_col1, save_col2 = st.columns([3, 1])

                        with save_col1:
                            filename = st.text_input(
                                "Filename",
                                value=f"{meeting_date.strftime('%Y-%m-%d')}-{meeting_title.lower().replace(' ', '-') if meeting_title else 'meeting'}.md"
                            )

                        with save_col2:
                            st.markdown("<br>", unsafe_allow_html=True)
                            if st.button("💾 Save to Project"):
                                # Create project directory if needed
                                project_dir = Path("data") / project_name / "meetings"
                                project_dir.mkdir(parents=True, exist_ok=True)

                                # Save file
                                file_path = project_dir / filename
                                with open(file_path, 'w') as f:
                                    f.write(structured_md)

                                st.success(f"✅ Saved to `{file_path}`")
                                st.info("💡 Refresh the page to reload meetings into the RAG system")

                                # Download option
                                st.download_button(
                                    label="📥 Download Markdown",
                                    data=structured_md,
                                    file_name=filename,
                                    mime="text/markdown"
                                )

                except Exception as e:
                    st.error(f"Error: {str(e)}")
                    if "quota" in str(e).lower() or "rate" in str(e).lower():
                        st.warning("⚠️ API rate limit reached. Please wait a moment and try again.")


if __name__ == "__main__":
    main()