import os
import re
import json

import streamlit as st
# from streamviz import gauge

from utils import validate_pdf
from styles import apply_custom_styles
from policy_analyser.analyse import Health

if 'GPT_KEY' not in os.environ or os.environ.get('GPT_KEY') in [None, '']:
    os.environ['GPT_KEY'] = st.secrets['GPT_KEY']

if 'AZURE_LAYOUT_KEY' not in os.environ.get('AZURE_LAYOUT_KEY') or os.environ.get('AZURE_LAYOUT_KEY') in [None, '']:
    os.environ['AZURE_LAYOUT_KEY'] = st.secrets['AZURE_LAYOUT_KEY']

if 'health_analyser' not in st.session_state:
    st.session_state.health_analyser = Health(ocr_engine = 'azure/layout')

def markdown_table_to_json(markdown):
    lines = markdown.strip().split("\n")
    
    # Extract headers
    headers = [h.strip() for h in lines[0].split("|") if h.strip()]
    
    # Extract rows
    rows = []
    for line in lines[2:]:  # Skip header and separator line
        values = [v.strip() for v in line.split("|") if v.strip()]
        row_dict = dict(zip(headers, values))
        rows.append(row_dict)
    
    return rows

def visualise_pie_chart(analysis):
    verdicts = {}
    score = 0
    total = 0
    for verdict in ['GOOD', 'AVERAGE', 'BAD']:
        table = analysis.split(f'<{verdict}>')[-1].split(f'</{verdict}>')[0]
        table = markdown_table_to_json(table)
        if len(table) > 0:
            verdicts[verdict] = table
            if verdict == 'GOOD':
                score += 5 * len(table)
            if verdict == 'AVERAGE':
                score += 3 * len(table)
            elif verdict == 'BAD':
                score += len(table)
            total += 5 * len(table)
    gauge(gVal = total, gTitle = '', gMode = 'gauge+number',
          grLow = total // 3,
          grMid = 2 * (total // 3))

def main():
    # Apply custom styles
    apply_custom_styles()

    # Header
    st.markdown("""
        <div class="header-container">
            <img src="https://acko-brand.ackoassets.com/brand/vector-svg/gradient/horizontal-reverse.svg" height=50 width=100>
                <h1>Insurance Policy Analyzer</h1>
            <p>Upload and compare insurance policies</p>
        </div>
    """, unsafe_allow_html=True)

    # File upload section
    st.markdown('<div class="upload-container">', unsafe_allow_html=True)
    uploaded_files = st.file_uploader("Choose policy PDF files", type="pdf", accept_multiple_files=True)
    print(uploaded_files)
    lob = st.selectbox(
        'Type of insurance',
        options = ['Health', 'Life', 'Auto'],
        index = 0
    )
    st.markdown('</div>', unsafe_allow_html=True)

    if uploaded_files and st.button('Analyse'):
        # Create tabs for different views
        tab1, tab2 = st.tabs(["Summary View", "Detailed Comparison"])

        # Store analysis results
        all_analyses = []

        # Process each uploaded file
        for uploaded_file in uploaded_files:
            # Read PDF content
            pdf_bytes = uploaded_file.read()
            # displayPDF(pdf_bytes)

            # Validate PDF
            if not validate_pdf(pdf_bytes):
                st.error(f"Invalid PDF file: {uploaded_file.name}")
                continue

            # Show loading state
            with st.spinner(f"Analyzing {uploaded_file.name}..."):
                try:
                    # Make API call
                    response = st.session_state.health_analyser(pdf_bytes)
                    analysis = next(
                        (item for item in response if item.get("stage") == "ANALYSE"), None
                    )['response']
                    analysis = analysis.split('<CUSTOMER_RESPONSE>')[-1].split('</CUSTOMER_RESPONSE>')[0]
                    suggestion = next(
                        (item for item in response if item.get("stage") == "SUGGEST"), None
                    )['response']
                    # suggestion = suggestion.split('<POLICY_PITCH>')[-1].split('</POLICY_PITCH>')[0]
                    suggestion = suggestion.replace('<POLICY_PITCH>', '').replace('</POLICY_PITCH>', '')
                    text = next(
                        (item for item in response if item.get("stage") == "OCR"), None
                    )['response']
                    # Store results
                    all_analyses.append({
                        'name': uploaded_file.name,
                        'text' : text,
                        'analysis' : analysis,
                        'suggestion' : suggestion
                    })

                except Exception as e:
                    st.error(f"Error analyzing {uploaded_file.name}: {str(e)}")

        # Summary View Tab
        with tab1:
            for idx, analysis in enumerate(all_analyses):
                with st.expander(f"### Policy {idx + 1}: {analysis['name']}"):
                    with st.container():
                        st.markdown(f'Policy Document : {analysis["text"]}')
                    with st.container():
                        st.markdown(re.sub(r'\<\/?(GOOD|AVERAGE|BAD|FINAL_VERDICT)\>', '', analysis['analysis']))
                    with st.container():
                        st.markdown('# Why Acko? 🚀')
                        st.markdown(analysis['suggestion'])
                        # visualise_pie_chart(analysis['analysis'])

        # Detailed Comparison Tab
        with tab2:
            st.warning('Coming Soon')
            # if len(all_analyses) > 1:
            #     # Create comparison matrix
            #     factors_to_compare = set()
            #     for analysis in all_analyses:
            #         factors_to_compare.update(
            #             [f.split(':')[0] for f in analysis['good_factors'] + 
            #              analysis['average_factors'] + analysis['bad_factors']]
            #         )

            #     # Create comparison table
            #     st.markdown("### Policy Comparison Matrix")

            #     comparison_data = []
            #     for factor in sorted(factors_to_compare):
            #         row = {'Factor': factor}
            #         for idx, analysis in enumerate(all_analyses):
            #             policy_name = f"Policy {idx + 1}"
            #             verdict = 'Not Found'
            #             for category in ['good_factors', 'average_factors', 'bad_factors']:
            #                 for item in analysis[category]:
            #                     if item.split(':')[0] == factor:
            #                         verdict = category.split('_')[0].title()
            #                         break
            #             row[policy_name] = verdict
            #         comparison_data.append(row)

            #     # Display comparison table
            #     st.table(comparison_data)
            # else:
            #     st.info("Upload multiple policies to see comparison")

    # Footer
    st.markdown("""
        <div style="margin-top: 50px; text-align: center; color: #666;">
            <p>Upload one or more insurance policy PDFs to get detailed analysis and comparison.</p>
            <p>We support all major insurance providers.</p>
        </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    st.set_page_config(
        page_title="Insurance Policy Analyzer",
        page_icon="📋",
        layout="wide"
    )
    main()