import streamlit as st
import os
from setup import page_setup
from utilities import save_notebook, save_datasets
from grammar_exec import execute
from notebook import main_call
import pandas as pd
from pathlib import Path


col1, col2, col3 = st.columns([2, 4, 1])
with col2:
    st.title(":blue[Great] Lens 🕵️‍♂️")

# File upload section
st.subheader("📁 Upload Files")

col1, col2 = st.columns(2)

with col1:
    st.markdown("**Upload Notebook**")
    notebook = st.file_uploader(
        label='Select Jupyter Notebook',
        accept_multiple_files=False,
        type=['ipynb'],
        help="Upload a .ipynb file to analyze"
    )

with col2:
    st.markdown("**Upload Datasets (Optional)**")
    datasets = st.file_uploader(
        label='Select Dataset Files',
        accept_multiple_files=True,
        type=['csv', 'xlsx', 'xls', 'json', 'txt', 'parquet'],
        help="Upload datasets that your notebook references"
    )

# Display uploaded files info
if datasets:
    st.info(f"📊 {len(datasets)} dataset(s) uploaded: {', '.join([f.name for f in datasets])}")

if notebook:
    st.success(f"📓 Notebook uploaded: {notebook.name}")
    
    # Save files to /tmp/Notebook
    save_notebook(notebook)
    if datasets:
        save_datasets(datasets)
        st.info("✅ Datasets saved to notebook directory")
    
    results_tab, grammar_tab = st.tabs(['Execution', 'Grammar/Fact']) 
    
    # with results_tab:
    #     with st.spinner("🔄 Executing notebook..."):
    #         try:
    #             notebook_dir_path = Path("/tmp/Notebook")
    #             notebook_files = [f for f in notebook_dir_path.iterdir() if f.suffix == '.ipynb']
                
    #             if not notebook_files:
    #                 st.error("No notebook found in directory")
    #             else:
    #                 notebook_path = notebook_files[0]
    #                 st.write(f'🚀 Processing notebook: {notebook_path.name}')
                    
    #                 # Show available datasets
    #                 dataset_files = [f for f in notebook_dir_path.iterdir() 
    #                                if f.suffix.lower() in ['.csv', '.xlsx', '.xls', '.json', '.txt', '.parquet']]
                    
    #                 if dataset_files:
    #                     st.info(f"📁 Available datasets: {', '.join([f.name for f in dataset_files])}")
                    
    #                 results = main_call(notebook_path)
                    
    #                 # Display results in a more user-friendly way
    #                 if isinstance(results, dict):
    #                     col1, col2 = st.columns([1, 3])
    #                     with col1:
    #                         if results['status'] == 'Pass':
    #                             st.success("✅ **Status: PASSED**")
    #                         else:
    #                             st.error("❌ **Status: FAILED**")
                        
    #                     with col2:
    #                         st.write(f"**Notebook:** {results['notebook']}")
    #                         if results['error_message']:
    #                             st.error(f"**Error:** {results['error_message']}")
    #                 else:
    #                     st.dataframe(results)
                        
    #         except Exception as e:
    #             st.error(f"❌ Error processing notebook: {str(e)}")
    
    with grammar_tab:
        try:
            with st.spinner("🔍 Analyzing grammar and facts..."):
                results = execute("/tmp/Notebook")
                
                if not results.empty:
                    # Display grammar results in a more readable format
                    st.subheader("📝 Grammar & Style Analysis")
                    
                    if 'Grammar_Text' in results.columns and len(results['Grammar_Text'].dropna()) > 0:
                        grammar_issues = results[results['Grammar_Text'].notna()]
                        
                        for idx, row in grammar_issues.iterrows():
                            if row['Is Grammar Error?']:
                                st.warning(f"**Grammar Error:** {row['Grammar_Text']}")
                                st.info(f"**Suggestion:** {row['Grammar_Suggestions']}")
                            else:
                                st.info(f"**Style Suggestion:** {row['Grammar_Text']}")
                                st.success(f"**Improvement:** {row['Grammar_Suggestions']}")
                            st.divider()
                    
                    st.subheader("🎯 Factual Accuracy Analysis")
                    
                    if 'Fact_Text' in results.columns and len(results['Fact_Text'].dropna()) > 0:
                        fact_issues = results[results['Fact_Text'].notna()]
                        
                        for idx, row in fact_issues.iterrows():
                            st.error(f"**Factual Error:** {row['Fact_Text']}")
                            st.success(f"**Correction:** {row['Fact_Suggestions']}")
                            st.divider()
                    
                    # Show raw dataframe as well
                    with st.expander("📊 View Raw Results"):
                        st.dataframe(results)
                else:
                    st.success("✅ No grammar or factual issues found!")
                    
        except Exception as e:
            st.error(f"❌ Unable to process grammar/facts: {str(e)}")

# Add some helpful information
st.sidebar.markdown("## 💡 How to Use")
st.sidebar.markdown("""
1. **Upload Notebook**: Select your .ipynb file
2. **Upload Datasets**: Add any CSV, Excel, or other data files your notebook uses
3. **Execution Tab**: See if your notebook runs successfully
4. **Grammar/Fact Tab**: Check for text quality and factual accuracy

### 🔧 Colab Support
The tool automatically handles Google Colab specific code:
- Replaces Drive mounts with local file access
- Uses your uploaded datasets instead of Colab file uploads
- Skips Colab-specific imports that won't work locally
""")

st.sidebar.markdown("## 📋 Supported Formats")
st.sidebar.markdown("""
**Notebooks:** .ipynb  
**Datasets:** .csv, .xlsx, .xls, .json, .txt, .parquet
""")