import argparse
import os
from pathlib import Path
import subprocess
import sys
import time
from fabriq.rag_pipeline import RAGPipeline
from fabriq.indexing import DocumentIndexer
from fabriq.config_parser import ConfigParser
import streamlit as st
st.set_page_config(page_title="Chatbot", page_icon="🤖", layout='wide', initial_sidebar_state='collapsed')

@st.cache_resource(show_spinner=False)
def load_chatbot(_config: ConfigParser):
    """Cache chatbot initialization to avoid reloading"""
    rag_pipeline = RAGPipeline(_config)
    return rag_pipeline

@st.cache_data(show_spinner=False)
def load_config(config_path: str):
    """Cache config loading to avoid repeated file reads"""
    return ConfigParser(config_path)

def index_documents(file_paths: list, config: ConfigParser):
    """Index documents from a given file path"""
    indexer = DocumentIndexer(config)
    indexer.index_documents(file_paths)

def stream_data(chatbot, prompt):
    """Generator for streaming response chunks"""
    with st.spinner("Thinking..."):
        response = chatbot.get_response(prompt, stream=True)
    for chunk in response["text"]:
        yield chunk.content
        time.sleep(0.01)


def app(config_path="config/config.yaml"):
    script_path = os.path.join(os.path.dirname(__file__), "chatbot.py")
    subprocess.run([sys.executable, "-m", "streamlit", "run", script_path])
    
    # Load cached config
    config = load_config(config_path)
    chatbot = load_chatbot(config)
    
    st.title("Chatbot")

    uploaded_files = st.sidebar.file_uploader("Upload documents to index", accept_multiple_files=True, type=['pdf', 'docx', 'pptx', 'xlsx'])
    if uploaded_files:
        file_paths = [f.name for f in uploaded_files]
        index_documents(file_paths, config)
        st.success("Documents indexed successfully!")
    
    # Initialize messages once
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # Render existing messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Handle new input
    if prompt := st.chat_input("Ask your query"):
        # Append and display user message
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Stream assistant response using st.write_stream
        with st.chat_message("assistant"):
            full_response = st.write_stream(stream_data(chatbot, prompt))
        
        # Append assistant response
        st.session_state.messages.append({"role": "assistant", "content": full_response})

# app("/Users/aaryan/AI Agent Builder/config/config.yaml")
def main():
    """Entry point for the Chat UI tool"""
    parser = argparse.ArgumentParser(
        description="Chat with your Documents using Fabriq Chat UI.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    
    parser.add_argument(
        "--config",
        "-c",
        default="config/config.yaml",
        help="Path to configuration file (default: config/config.yaml)"
    )
    
    args = parser.parse_args()
    
    # Check if config file exists
    if not Path(args.config).exists():
        print(f"❌ Config file not found: {args.config}")
        sys.exit(1)
    
    # Initialize and run chatbot
    app(config_path=args.config)


if __name__ == "__main__":
    main()