from dash import dcc, html, Input, Output, dash_table
import pandas as pd
from pandas.api.types import CategoricalDtype
import dataiku
import re
import json
from datetime import datetime
import traceback

from webapp.config import WebAppConfig
from webapp.utils import get_committed_note_ids, get_note_summary_by_note_id, get_verified_codes_by_note_id, get_code_labels, get_edit_logs_by_note_id


committed_note_ids = get_committed_note_ids()

diag_primary_style = WebAppConfig.Diagnosis.get_primary_style() 
diag_other_style = WebAppConfig.Diagnosis.get_other_style() 
proc_major_style = WebAppConfig.Procedure.get_major_style()
proc_other_style = WebAppConfig.Procedure.get_other_style()

app.layout = html.Div(
    style={
        'fontFamily': 'Roboto, sans-serif',
        'backgroundColor': '#f8f9fa',
        'padding': '2rem'
    },
    children=[
        html.Div(
            style={
                'maxWidth': '900px', 'margin': '0 auto', 'padding': '2rem',
                'backgroundColor': '#ffffff', 'borderRadius': '8px',
                'boxShadow': '0 4px 6px rgba(0, 0, 0, 0.1)'
            },
            children=[
                dcc.Store(id='data-store'),
                
                html.H1(
                    "Note Summary & Verified Codes",
                    style={'color': '#343a40', 'borderBottom': '2px solid #dee2e6', 'paddingBottom': '0.5rem'}
                ),
                html.P(
                    style={'color': '#6c757d', 'marginTop': '1rem'},
                    children=[
                        "Select a note ID to review the note summary and verified codes. ",
                        html.Span("Hover over the billing codes to see the edit history.")
                    ]
                ),
                html.Div(
                    style={'marginTop': '2rem', 'marginBottom': '2rem'},
                    children=[
                        html.Label(
                            "Select Note ID:",
                            style={'fontWeight': 'bold', 'color': '#495057', 'display': 'block', 'marginBottom': '0.5rem'}
                        ),
                        dcc.Dropdown(
                            id='note-id-dropdown',
                            options=[{'label': i, 'value': i} for i in committed_note_ids],
                            value=committed_note_ids[0] if len(committed_note_ids) > 0 else None, # Select the first ID by default
                            placeholder="Select a Note ID...",
                            clearable=False
                        )
                    ]
                ),
                dcc.Loading(
                    id="loading-summary",
                    type="circle",
                    children=html.Div(id='summary-container', style={'marginTop': '2rem'})
                ),
                dcc.Loading(
                    id="loading-table",
                    type="circle",
                    children=html.Div(
                        id='table-container',
                        style={'marginTop': '2rem'},
                        children=[
                            html.H3("Verified Codes", style={'color': '#343a40', 'borderBottom': '1px solid #dee2e6', 'paddingBottom': '0.5rem'}),
                            html.Label("Filter by Concept Type:", style={'fontWeight': 'bold', 'marginTop': '1rem'}),
                            dcc.Dropdown(id='concept-type-filter', clearable=False),
                            html.Div(style={'marginTop': '1rem'}),
                            dash_table.DataTable(id='verified-codes-table', style_cell={'fontFamily': 'Roboto, sans-serif'})
                        ]
                    )
                )
            ]
        )
    ]
)

# --- Callback 1: Update data and component properties when Note ID changes ---
@app.callback(
    Output('summary-container', 'children'),
    Output('data-store', 'data'),
    Output('concept-type-filter', 'options'),
    Output('concept-type-filter', 'value'),
    Input('note-id-dropdown', 'value')
)
def update_on_note_selection(selected_note_id):
    if selected_note_id is None:
        return html.Div("Please select a Note ID."), None, [], None
    try:
        # Fetch the data for the selected note ID
        summary_df = get_note_summary_by_note_id(selected_note_id)
        codes_df = get_verified_codes_by_note_id(selected_note_id)
        log_df = get_edit_logs_by_note_id(selected_note_id)

        if summary_df.empty:
            summary_component = html.Div(f"No summary found. type: {type(selected_note_id)}", style={'color': 'red'})
        else:
            patient_info = summary_df.iloc[0]
            patient_id = patient_info['patient_id']
            sex = patient_info['sex']
            age = patient_info['age']
            specialty = patient_info['specialty']
            summary_text = patient_info['summary']
            summary_component = html.Div([
                html.H3("Patient Information", style={'color': '#343a40', 'borderBottom': '1px solid #dee2e6', 'paddingBottom': '0.5rem'}),
                html.Div(
                    style={'display': 'flex', 'justifyContent': 'space-around', 'padding': '1rem', 'backgroundColor': '#f8f9fa', 'borderRadius': '5px'},
                    children=[
                        html.Div([html.B("Patient ID: "), html.Span(patient_id)]),
                        html.Div([html.B("Sex: "), html.Span(sex)]),
                        html.Div([html.B("Age: "), html.Span(age)]),
                    ]
                ),
                html.H3("Note Summary", style={'marginTop': '2rem', 'color': '#343a40', 'borderBottom': '1px solid #dee2e6', 'paddingBottom': '0.5rem'}),
                html.Div([html.B("Medical Specialty: "), html.Span(specialty)]),
                html.P(summary_text, style={'lineHeight': '1.6', 'color': '#212529'})
            ])

        if codes_df.empty:
            return summary_component, None, [], 'All'

        codes = codes_df['Mapped billing code'].unique().tolist()
        labels_df = get_code_labels(codes)
        code_w_label_df = codes_df.merge(labels_df, left_on='Mapped billing code', right_on='billing_references').drop(columns=['billing_references'])

        concept_types = code_w_label_df['Concept type'].unique()
        filter_options = [{'label': 'All', 'value': 'All'}] + [{'label': i, 'value': i} for i in concept_types]

        stored_data = {
            'code': code_w_label_df.to_json(orient='split'),
            'log': log_df.to_json(orient='split')
        }

        return summary_component, stored_data, filter_options, 'All'
    
    except Exception as e:
        # --- DIAGNOSTIC 4: Catch and print any errors ---
        print("!!! AN ERROR OCCURRED IN THE CALLBACK !!!")
        print(traceback.format_exc())
        # Return an error message to the UI to show something went wrong
        error_message = f"An error occurred: {e}"
        return html.Div(error_message, style={'color': 'red'}), None, [], None
        
# --- Callback 2: EXPANDED to be the single controller for the table ---
@app.callback(
    Output('verified-codes-table', 'data'),
    Output('verified-codes-table', 'columns'),
    Output('verified-codes-table', 'tooltip_data'),
    Output('verified-codes-table', 'style_table'),
    Output('verified-codes-table', 'style_header'),
    Output('verified-codes-table', 'style_cell'),
    Output('verified-codes-table', 'style_data_conditional'),
    Input('concept-type-filter', 'value'),
    Input('data-store', 'data')
)
def filter_and_style_table(selected_type, stored_data):
    # If there's no stored data, return empty values for everything
    if not stored_data:
        return [], [], [], {}, {}, {}, []

    # Read data from the store
    codes_df = pd.read_json(stored_data['code'], orient='split')
    log_df = pd.read_json(stored_data['log'], orient='split')

    # Filter the DataFrame based on the dropdown
    if selected_type == 'All' or selected_type is None:
        filtered_df = codes_df
    else:
        filtered_df = codes_df[codes_df['Concept type'] == selected_type]
        
    if not filtered_df.empty:
        filtered_df['key'] = '(' + filtered_df['Note ID'].astype(str) + ', ' + filtered_df['No'].astype(str) + ')'

    # Create tooltips to highlight cells with edit history
    tooltips = []
    styles = [{'if': {'row_index': 'odd'}, 'backgroundColor': 'rgb(248, 248, 248)'}]
    
    if not log_df.empty:
        history_lookup = {(row['key'], row['column_name']): row for index, row in log_df.iterrows()}
        for i, row in filtered_df.iterrows():
            row_tooltips = {}
            for col in filtered_df.columns:
                history_item = history_lookup.get((row['key'], col))
                if history_item is not None:
                    try:
                        # dateutil.parser is more robust for various formats
                        date_obj = pd.to_datetime(history_item['date'])
                        formatted_date = date_obj.strftime('%Y-%m-%d %H:%M')
                    except:
                        formatted_date = "Invalid Date"
                    row_tooltips[col] = {
                        'value': f"Last Edit by: {history_item['user']}\n\nDate: {formatted_date}",
                        'type': 'markdown'
                    }
            tooltips.append(row_tooltips)
        
        for i, row in filtered_df.reset_index(drop=True).iterrows():
            for col in filtered_df.columns:
                if (row['key'], col) in history_lookup:
                    styles.append({
                        'if': {'row_index': i, 'column_id': col},
                        'backgroundColor': '#fff3cd',
                        'fontWeight': 'bold'
                    })
    concept_order = ['primary diagnosis', 'relevant diagnosis', 'major procedure or operation', 'relevant medical services and procedures']
    cat_type = CategoricalDtype(categories=concept_order, ordered=True)
    filtered_df['Concept type'] = filtered_df['Concept type'].astype(cat_type)
    table_data = filtered_df.sort_values(by="Concept type").to_dict('records')
    columns_to_show = ['Concept type', 'Mapped billing code', 'label']
    table_columns = [{"name": i, "id": i, "resizable": True, "width": "150px"} for i in filtered_df.columns if i in columns_to_show]
    
    style_table = {'overflowX': 'auto'}
    style_header = {'backgroundColor': 'rgb(230, 230, 230)', 'fontWeight': 'bold'}
    style_cell = {'padding': '10px', 'textAlign': 'left', 'border': '1px solid grey'}

    return table_data, table_columns, tooltips, style_table, style_header, style_cell, styles