import streamlit as st
import requests

API_URL = "http://localhost:7860"  # Change to your backend URL

st.title("Financial Statements Frontend")

main_section = st.sidebar.radio(
    "Section",
    ["Financial Statement Generation", "RLHF Management"]
)

if main_section == "Financial Statement Generation":
    endpoint = st.sidebar.selectbox(
        "Choose Statement Endpoint",
        ["Balance Sheet", "Profit & Loss", "Cash Flow", "Notes", "Interactive Notes"]
    )

    if endpoint == "Interactive Notes":
        st.header("🎯 Interactive Notes Generation with Feedback")
        
        # Session Management
        col1, col2 = st.columns([3, 1])
        with col1:
            default_session_id = st.session_state.get('session_id', '')
            session_id = st.text_input("Session ID (leave empty for new session)", 
                                    value=default_session_id,
                                    placeholder="Auto-generated for new sessions")
        with col2:
            if st.button("🔍 Check Session Status"):
                if session_id:
                    try:
                        resp = requests.get(f"{API_URL}/notes-llm/session/{session_id}")
                        if resp.status_code == 200:
                            session_data = resp.json()
                            st.success(f"Session Status: {session_data['status']}")
                            st.info(f"Iterations: {session_data['current_iteration']}")
                            st.info(f"Feedbacks: {len(session_data['feedback_history'])}")
                        else:
                            st.error("Session not found")
                    except Exception as e:
                        st.error(f"Error checking session: {e}")
        
        # File Upload
        st.subheader("📁 Upload Trial Balance")
        f = st.file_uploader("Upload Excel file", type=["xlsx"])
        
        # Workflow Tabs
        tab1, tab2, tab3, tab4 = st.tabs(["🚀 Initial Generation", "💬 Submit Feedback", 
                                        "🔄 Generate Improved", "✅ Final Approval"])
        
        with tab1:
            st.subheader("Initial Notes Generation")
            use_rlhf = st.checkbox("Use RLHF (if available)", value=False)
            
            if 'initial_notes_generated' in st.session_state and st.session_state['initial_notes_generated']:
                st.success("✅ Initial notes have been generated!")
                st.info(f"📋 Session ID: {st.session_state['session_id']}")
                
                st.download_button(
                    "📥 Download Initial Notes",
                    data=st.session_state['notes_content'],
                    file_name=st.session_state['notes_filename'],
                    mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                )
                
                st.subheader("📊 Generation Info")
                col1, col2, col3 = st.columns(3)
                with col1:
                    st.metric("Validation Score", st.session_state.get('validation_score', 'N/A'))
                with col2:
                    st.metric("Attempts", st.session_state.get('attempts_made', 'N/A'))
                with col3:
                    st.metric("Interactive", st.session_state.get('interactive_enabled', 'N/A'))
                    
                if st.button("🔄 Start New Generation"):
                    for key in ['initial_notes_generated', 'session_id', 'notes_content', 'notes_filename', 
                              'validation_score', 'attempts_made', 'interactive_enabled']:
                        if key in st.session_state:
                            del st.session_state[key]
                    st.rerun()
            
            elif f and st.button("🎯 Generate Initial Notes", type="primary"):
                with st.spinner("Generating initial notes..."):
                    url = f"{API_URL}/notes-llm"
                    if use_rlhf:
                        url += "?use_rlhf=true"
                    
                    try:
                        resp = requests.post(url, files={"file": (f.name, f, f.type)})
                        
                        if resp.status_code == 200:
                            new_session_id = resp.headers.get('X-Session-ID')
                            if new_session_id:
                                st.session_state['session_id'] = new_session_id
                                st.session_state['initial_notes_generated'] = True
                                st.session_state['notes_content'] = resp.content
                                st.session_state['notes_filename'] = "notes_initial.xlsx"
                                st.session_state['validation_score'] = resp.headers.get('X-Validation-Score', 'N/A')
                                st.session_state['attempts_made'] = resp.headers.get('X-Attempts-Made', 'N/A')
                                st.session_state['interactive_enabled'] = resp.headers.get('X-Interactive-Enabled', 'N/A')
                                
                                cd = resp.headers.get("content-disposition", "")
                                if "filename=" in cd:
                                    st.session_state['notes_filename'] = cd.split("filename=")[-1].strip().strip('"').strip("'")
                                
                                st.success(f"✅ Initial generation complete!")
                                st.info(f"📋 Session ID: {new_session_id}")
                                
                                st.download_button(
                                    "📥 Download Initial Notes",
                                    data=resp.content,
                                    file_name=st.session_state['notes_filename'],
                                    mime=resp.headers.get("content-type", "application/octet-stream")
                                )
                                
                                st.subheader("📊 Generation Info")
                                col1, col2, col3 = st.columns(3)
                                with col1:
                                    st.metric("Validation Score", st.session_state['validation_score'])
                                with col2:
                                    st.metric("Attempts", st.session_state['attempts_made'])
                                with col3:
                                    st.metric("Interactive", st.session_state['interactive_enabled'])
                                
                        else:
                            st.error(f"❌ Generation failed: {resp.text}")
                            
                    except Exception as e:
                        st.error(f"❌ Error: {e}")
        
        with tab2:
            st.subheader("Submit Feedback for Improvement")
            
            if not session_id:
                st.warning("⚠ Please enter a Session ID from Step 1")
            else:
                feedback_text = st.text_area(
                    "💬 Describe what you'd like to improve:",
                    placeholder="e.g., Add more detailed depreciation notes with asset categories",
                    height=100
                )
                
                feedback_type = st.selectbox(
                    "📋 Feedback Type:",
                    ["text", "numeric", "formula", "suggestion"],
                    help="""
                    • text: Additional explanations
                    • numeric: Number-related changes  
                    • formula: Excel formulas
                    • suggestion: General improvements
                    """
                )
                
                if st.button("📤 Submit Feedback", type="primary"):
                    if not feedback_text.strip():
                        st.error("❌ Please enter feedback text")
                    else:
                        with st.spinner("Submitting feedback..."):
                            try:
                                data = {
                                    "session_id": session_id,
                                    "feedback_text": feedback_text,
                                    "feedback_type": feedback_type
                                }
                                
                                resp = requests.post(f"{API_URL}/notes-llm/feedback", data=data)
                                
                                if resp.status_code == 200:
                                    result = resp.json()
                                    st.success("✅ Feedback submitted successfully!")
                                    st.info(f"🔢 Iteration: {result['iteration']}")
                                    st.info(f"🏷 UDF Version: {result['udf_version']}")
                                    
                                    st.info("💡 Next: Go to 'Generate Improved' tab to see the enhancements!")
                                    
                                else:
                                    st.error(f"❌ Feedback submission failed: {resp.text}")
                                    
                            except Exception as e:
                                st.error(f"❌ Error: {e}")
        
        with tab3:
            st.subheader("Generate Improved Notes")
            
            if not session_id:
                st.warning("⚠ Please enter a Session ID")
            elif not f:
                st.warning("⚠ Please upload the trial balance file again")
            else:
                improved_key = f"improved_notes_{session_id}"
                if improved_key in st.session_state and st.session_state[improved_key]:
                    st.success("✅ Improved notes have been generated!")
                    
                    st.download_button(
                        "📥 Download Improved Notes",
                        data=st.session_state[improved_key]['content'],
                        file_name=st.session_state[improved_key]['filename'],
                        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                    )
                    
                    col1, col2, col3 = st.columns(3)
                    with col1:
                        st.metric("Iteration", st.session_state[improved_key]['iteration'])
                    with col2:
                        st.metric("Feedbacks Applied", st.session_state[improved_key]['feedbacks_applied'])
                    with col3:
                        st.metric("UDFs Used", st.session_state[improved_key]['udfs_used'])
                        
                    if st.button("🔄 Regenerate Improved Notes"):
                        if improved_key in st.session_state:
                            del st.session_state[improved_key]
                        st.rerun()
                
                else:
                    st.info("🔄 This will apply all submitted feedback to generate improved notes")
                    
                    if st.button("🚀 Generate Improved Notes", type="primary"):
                        with st.spinner("Applying feedback and generating improved notes..."):
                            try:
                                files = {"file": (f.name, f, f.type)}
                                data = {"session_id": session_id}
                                
                                resp = requests.post(f"{API_URL}/notes-llm/generate", 
                                                   files=files, data=data)
                                
                                if resp.status_code == 200:
                                    st.success("✅ Improved notes generated!")
                                    
                                    iteration = resp.headers.get('X-Iteration', '1')
                                    filename = f"notes_improved_v{iteration}.xlsx"
                                    
                                    st.session_state[improved_key] = {
                                        'content': resp.content,
                                        'filename': filename,
                                        'iteration': iteration,
                                        'feedbacks_applied': resp.headers.get('X-Feedbacks-Applied', 'N/A'),
                                        'udfs_used': resp.headers.get('X-UDFs-Archived', 'N/A')
                                    }
                                    
                                    col1, col2, col3 = st.columns(3)
                                    with col1:
                                        st.metric("Iteration", iteration)
                                    with col2:
                                        st.metric("Feedbacks Applied", st.session_state[improved_key]['feedbacks_applied'])
                                    with col3:
                                        st.metric("UDFs Used", st.session_state[improved_key]['udfs_used'])
                                    
                                    st.download_button(
                                        "📥 Download Improved Notes",
                                        data=resp.content,
                                        file_name=filename,
                                        mime=resp.headers.get("content-type", "application/octet-stream")
                                    )
                                    
                                else:
                                    st.error(f"❌ Generation failed: {resp.text}")
                                    
                            except Exception as e:
                                st.error(f"❌ Error: {e}")
        
        with tab4:
            st.subheader("Final Approval")
            
            if not session_id:
                st.warning("⚠ Please enter a Session ID")
            else:
                st.info("✅ This will finalize the session and set the final UDF")
                st.warning("⚠ After approval, no more feedback can be added to this session")
                
                if st.button("🏁 Approve Final Result", type="primary"):
                    with st.spinner("Approving session..."):
                        try:
                            data = {"session_id": session_id}
                            resp = requests.post(f"{API_URL}/notes-llm/approve", data=data)
                            
                            if resp.status_code == 200:
                                result = resp.json()
                                st.success("✅ Session approved successfully!")
                                st.info(f"📋 Session: {result['session_id']}")
                                st.info(f"🏷 Final UDF: {result['final_udf']}")
                                
                                st.balloons()
                                st.success("🎉 Interactive feedback session completed!")
                                
                            else:
                                st.error(f"❌ Approval failed: {resp.text}")
                                
                        except Exception as e:
                            st.error(f"❌ Error: {e}")
        
        # Session History Sidebar
        if session_id:
            st.sidebar.subheader("📋 Session History")
            try:
                resp = requests.get(f"{API_URL}/notes-llm/session/{session_id}")
                if resp.status_code == 200:
                    session_data = resp.json()
                    
                    st.sidebar.metric("Status", session_data['status'].upper())
                    st.sidebar.metric("Iterations", session_data['current_iteration'])
                    st.sidebar.metric("Feedbacks", len(session_data['feedback_history']))
                    
                    if session_data['feedback_history']:
                        st.sidebar.subheader("Recent Feedback")
                        for fb in session_data['feedback_history'][-3:]:
                            st.sidebar.text(f"#{fb['iteration']}: {fb['feedback_type']}")
                            st.sidebar.caption(fb['feedback_text'][:50] + "...")
                            
                else:
                    st.sidebar.error("Session not found")
            except Exception as e:
                st.sidebar.error(f"Error loading session: {e}")

    else:
        # Standard endpoints: Balance Sheet, P&L, Cash Flow, Notes
        st.header(f"{endpoint} Generation")
        f = st.file_uploader("Upload Excel file", type=["xlsx"])
        
        # Show RLHF option only for endpoints that support it
        if endpoint != "Notes":
            use_rlhf = st.checkbox("Use RLHF (Human Feedback Loop)?", value=False)
        else:
            use_rlhf = False
            st.info("📝 This is the simple notes generation endpoint (no RLHF/feedback)")
        
        if f and st.button(f"Generate {endpoint}"):
            endpoint_map = {
                "Balance Sheet": "bs",
                "Profit & Loss": "pnl", 
                "Cash Flow": "cf",
                "Notes": "notes"
            }
            url = f"{API_URL}/{endpoint_map[endpoint]}"
            if use_rlhf and endpoint != "Notes":
                url += "?use_rlhf=true"
            
            with st.spinner(f"Generating {endpoint}..."):
                try:
                    resp = requests.post(url, files={"file": (f.name, f, f.type)})
                    if resp.status_code == 200:
                        cd = resp.headers.get("content-disposition", "")
                        filename = f"{endpoint_map[endpoint]}_output.xlsx"
                        if "filename=" in cd:
                            filename = cd.split("filename=")[-1].strip().strip('"').strip("'")
                        
                        st.success(f"✅ {endpoint} generated successfully!")
                        st.download_button(
                            f"📥 Download {endpoint}",
                            data=resp.content,
                            file_name=filename,
                            mime=resp.headers.get("content-type", "application/octet-stream")
                        )
                        
                        # Show RLHF metadata if available
                        if use_rlhf and endpoint != "Notes":
                            st.subheader("📊 RLHF Information")
                            col1, col2, col3 = st.columns(3)
                            with col1:
                                st.metric("Statement ID", resp.headers.get('X-RLHF-Statement-ID', 'N/A'))
                            with col2:
                                st.metric("Quality Score", resp.headers.get('X-RLHF-Quality-Score', 'N/A'))
                            with col3:
                                st.metric("Confidence", resp.headers.get('X-RLHF-Confidence', 'N/A'))
                    else:
                        st.error(f"❌ Generation failed: {resp.text}")
                except Exception as e:
                    st.error(f"❌ Error: {e}")

elif main_section == "RLHF Management":
    rlhf_endpoint = st.sidebar.selectbox(
        "Choose RLHF Endpoint",
        [
            "Submit Feedback",
            "Pending Reviews", 
            "Model Info",
            "Stats",
            "Manual Retrain",
            "Review Statement"
        ]
    )

    if rlhf_endpoint == "Submit Feedback":
        st.header("📝 Submit Text-Based RLHF Feedback")
        st.info("💡 This system focuses on qualitative text feedback instead of numeric scores")
        
        statement_id = st.text_input("Statement ID", placeholder="Enter statement ID to review")
        reviewer_id = st.text_input("Reviewer ID", value="streamlit_user")
        
        st.subheader("📋 Text-Based Feedback (Required)")
        specific_errors = st.text_area(
            "Specific Errors Found",
            placeholder="Describe any specific errors, calculation mistakes, or inaccuracies...",
            height=100
        )
        missing_items = st.text_area(
            "Missing Items", 
            placeholder="List any items that should be included but are missing...",
            height=100
        )
        improvement_suggestions = st.text_area(
            "Improvement Suggestions",
            placeholder="Provide specific suggestions for improvement...",
            height=100
        )
        
        st.subheader("📊 Additional Information")
        would_accept_for_audit = st.checkbox("Would Accept for Audit?", value=True)
        complexity_level = st.selectbox("Complexity Level", ["low", "medium", "high"], index=1)

        if st.button("📤 Submit Text-Based Feedback", type="primary"):
            if not statement_id.strip():
                st.error("❌ Please enter a Statement ID")
            elif not any([specific_errors.strip(), missing_items.strip(), improvement_suggestions.strip()]):
                st.error("❌ Please provide at least one type of text feedback")
            else:
                with st.spinner("Submitting feedback..."):
                    try:
                        data = {
                            "statement_id": statement_id,
                            "reviewer_id": reviewer_id,
                            "specific_errors": specific_errors,
                            "missing_items": missing_items,
                            "improvement_suggestions": improvement_suggestions,
                            "would_accept_for_audit": would_accept_for_audit,
                            "complexity_level": complexity_level
                        }
                        resp = requests.post(f"{API_URL}/rlhf/feedback", data=data)
                        if resp.status_code == 200:
                            st.success("✅ Text-based feedback submitted successfully!")
                            result = resp.json()
                            st.json(result)
                        else:
                            st.error(f"❌ Submission failed: {resp.text}")
                    except Exception as e:
                        st.error(f"❌ Error: {e}")

    elif rlhf_endpoint == "Pending Reviews":
        st.header("📋 Pending Reviews")
        limit = st.number_input("Number of Reviews to Fetch", min_value=1, max_value=50, value=10)
        if st.button("🔍 Get Pending Reviews"):
            try:
                resp = requests.get(f"{API_URL}/rlhf/pending-reviews?limit={limit}")
                if resp.status_code == 200:
                    data = resp.json()
                    if data.get('pending_reviews'):
                        st.success(f"Found {len(data['pending_reviews'])} pending reviews")
                        for review in data['pending_reviews']:
                            with st.expander(f"Statement ID: {review.get('statement_id', 'N/A')}"):
                                st.write(f"*Type:* {review.get('statement_type', 'N/A')}")
                                st.write(f"*Created:* {review.get('created_at', 'N/A')}")
                                if st.button(f"Review {review.get('statement_id', 'N/A')}", key=review.get('statement_id')):
                                    st.info(f"Use 'Review Statement' tab with ID: {review.get('statement_id')}")
                    else:
                        st.info("No pending reviews found")
                else:
                    st.error(f"Error: {resp.text}")
            except Exception as e:
                st.error(f"❌ Error: {e}")

    elif rlhf_endpoint == "Model Info":
        st.header("🤖 RLHF Model Information")
        if st.button("📊 Get Model Info"):
            try:
                resp = requests.get(f"{API_URL}/rlhf/model-info")
                if resp.status_code == 200:
                    data = resp.json()
                    st.success("✅ Model information retrieved")
                    
                    # Display key metrics
                    col1, col2, col3 = st.columns(3)
                    with col1:
                        st.metric("Total Feedback", data.get('total_feedback_count', 'N/A'))
                    with col2:
                        st.metric("Model Type", data.get('model_type', 'N/A'))
                    with col3:
                        st.metric("Last Updated", data.get('last_updated', 'N/A'))
                    
                    st.json(data)
                else:
                    st.error(f"Error: {resp.text}")
            except Exception as e:
                st.error(f"❌ Error: {e}")

    elif rlhf_endpoint == "Stats":
        st.header("📈 RLHF Statistics")
        if st.button("📊 Get RLHF Stats"):
            try:
                resp = requests.get(f"{API_URL}/rlhf/stats")
                if resp.status_code == 200:
                    data = resp.json()
                    st.success("✅ Statistics retrieved")
                    
                    # Display key stats
                    if 'feedback_stats' in data:
                        stats = data['feedback_stats']
                        col1, col2, col3, col4 = st.columns(4)
                        with col1:
                            st.metric("Total Feedback", stats.get('total_feedback', 'N/A'))
                        with col2:
                            st.metric("Unique Reviewers", stats.get('unique_reviewers', 'N/A'))
                        with col3:
                            st.metric("Acceptance Rate", f"{stats.get('acceptance_rate', 0):.1%}")
                        with col4:
                            st.metric("Avg Complexity", stats.get('avg_complexity_score', 'N/A'))
                    
                    st.json(data)
                else:
                    st.error(f"Error: {resp.text}")
            except Exception as e:
                st.error(f"❌ Error: {e}")

    elif rlhf_endpoint == "Manual Retrain":
        st.header("🔄 Manual RLHF Retrain")
        st.info("💡 This will analyze current feedback patterns (no ML training involved)")
        if st.button("🚀 Trigger Retrain", type="primary"):
            try:
                resp = requests.post(f"{API_URL}/rlhf/retrain")
                if resp.status_code == 200:
                    st.success("✅ Retrain completed successfully!")
                    st.json(resp.json())
                else:
                    st.error(f"Error: {resp.text}")
            except Exception as e:
                st.error(f"❌ Error: {e}")

    elif rlhf_endpoint == "Review Statement":
        st.header("🔍 Review Specific Statement")
        statement_id = st.text_input("Statement ID", placeholder="Enter statement ID to review")
        if statement_id and st.button("📋 Get Review Form"):
            try:
                resp = requests.get(f"{API_URL}/rlhf/review/{statement_id}")
                if resp.status_code == 200:
                    if "text/html" in resp.headers.get("content-type", ""):
                        st.success("✅ Review form loaded")
                        st.components.v1.html(resp.text, height=800, scrolling=True)
                    else:
                        st.json(resp.json())
                else:
                    st.error(f"Statement not found: {resp.text}")
            except Exception as e:
                st.error(f"❌ Error: {e}")

# Clear session button
if st.sidebar.button("🗑 Clear All Session Data"):
    for key in list(st.session_state.keys()):
        del st.session_state[key]
    st.sidebar.success("✅ Session data cleared!")
    st.rerun()

st.caption("FinRyver Financial Statements Generator with Interactive Feedback & RLHF Support")