import streamlit as st
import PyPDF2
import docx
import pandas as pd
from langchain_core.prompts import ChatPromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.output_parsers import StrOutputParser

# Secure API key handling
api_key = "AIzaSyBuhWQHGfMQewIY_gdnmbzzYoori6faeUo"  # Store in .streamlit/secrets.toml

# Database options
db_options = ["MySQL", "PostgreSQL", "SQLite", "MSSQL", "Oracle", "MongoDB", "Firebase", "Pandas"]

# Function to read PDF files
def read_pdf(file):
    reader = PyPDF2.PdfReader(file)
    text = "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
    return text

# Function to read DOCX files
def read_docx(file):
    doc = docx.Document(file)
    text = "\n".join([para.text for para in doc.paragraphs])
    return text

# System prompt for SQL, NoSQL, and Pandas problem solving
prompt_template = ChatPromptTemplate.from_messages([
    ("system", """
    You are an expert in database and data analysis problem solving. Given a problem statement, 
    provide the most efficient solution using SQL, NoSQL, or Pandas depending on the selected system. 
    Ensure the query follows best practices, is optimized for performance, and works across multiple RDBMS like MySQL, PostgreSQL, SQLite, MSSQL, and Oracle.
    Additionally, support NoSQL databases such as MongoDB and Firebase by providing appropriate queries.
    If Pandas is selected, provide optimized Python code using Pandas functions.
    
    Output format:
    **Solution:**
    ```sql, NoSQL, or Python
    -- Query or Pandas code
    ```
    
    **Explanation:**
    - Step-by-step breakdown of the logic.
    - Performance considerations and optimizations if applicable.
    - Compatibility notes for the selected system.
    """),
    ("human", """
    Problem Statement: {sql_problem}
    Selected System: {db_type}
    """)
])

chat_model = ChatGoogleGenerativeAI(google_api_key=api_key, model='models/gemini-2.0-flash-exp')
parser = StrOutputParser()
chain = prompt_template | chat_model | parser

def solve_problem(sql_problem, db_type):
    """Generates an AI-powered solution based on the selected system."""
    try:
        response = chain.invoke({"sql_problem": sql_problem, "db_type": db_type})
        return response if response else "No solution available."
    except Exception as e:
        return f"Error: {str(e)}"

# Streamlit UI
st.title("🛠️ AI-Powered Data & Database Problem Solver")
st.write("Enter a problem statement, select the system type (SQL, NoSQL, or Pandas), and AI will generate an optimized solution for you.")

col1, col2 = st.columns(2)
with col1:
    sql_problem = st.text_area("Problem Statement", placeholder="Describe your problem here...")
    uploaded_file = st.file_uploader("Upload a DOCX or PDF file", type=["pdf", "docx"])
    if uploaded_file is not None:
        if uploaded_file.type == "application/pdf":
            sql_problem = read_pdf(uploaded_file)
        elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
            sql_problem = read_docx(uploaded_file)
        st.success("File uploaded successfully!")

with col2:
    db_type = st.selectbox("Select System Type", db_options)

if st.button("Generate Solution") and sql_problem:
    with st.spinner("Generating solution..."):
        solution = solve_problem(sql_problem, db_type)
    
    st.subheader("AI-Generated Solution")
    st.markdown(solution, unsafe_allow_html=True)