from dash import dcc, html, Input, Output
import pandas as pd
import dataiku
import re
import json
from datetime import datetime
import traceback

from webapp.config import WebAppConfig
from webapp.utils import get_llm_outputs_by_note_id, get_note_metadata_by_note_id, read_note_id, create_evidence_details_map, build_styled_text_components
from dataiku_filter_listener import DataikuFilterListener


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=[
                html.H1(
                    "Discharge Note Viewer",
                    style={'color': '#343a40', 'borderBottom': '2px solid #dee2e6', 'paddingBottom': '0.5rem'}
                ),
                html.P(
                    style={'color': '#6c757d', 'marginTop': '1rem'},
                    children=[
                        "The text supporting the extracted medical concepts is highlighted. ",
                        html.Span("Primary diagnoses", style=diag_primary_style), " are dark green, ",
                        html.Span("related diagnoses", style=diag_other_style), " are light green, ",
                        html.Span("major procedure or operation", style=proc_major_style), " are dark purple, and ",
                        html.Span("relevant medical services and procedures", style=proc_other_style), " are light purple.",
                        html.Br(),
                        html.Br(),
                        html.Span("Hover over the highlighted evidence to see the medical concept it supports.")
                    ]
                ),
                DataikuFilterListener(id='dfl'),
                html.Hr(),
                dcc.Loading(
                    id="loading-indicator",
                    type="circle",
                    children=html.Div(id="report-container")
                )
            ]
        )
    ]
)


# --- Define the callback ---
@app.callback(
    Output("report-container", "children"),
    Input("dfl", "filters")
)

def update_report_from_filter(filters):
    if filters is None:
        initial_html = html.P("ℹ️ Please select a note from the dashboard filter.")
        return [initial_html]

    try:
        note_id_str = read_note_id(filters)
        if note_id_str is None:
            raise ValueError("Could not find 'note_id' in the received filter data.")
        note_id = int(note_id_str)
        selected_note_outputs_df = get_llm_outputs_by_note_id(note_id)
        selected_note_metadata = get_note_metadata_by_note_id(note_id)

        if selected_note_metadata.empty:
            raise ValueError(f"Data for Note ID {note_id} not found in the dataset.")

        # Collect the evidence for each category
        domain_style_map = {
            'primary diagnosis': diag_primary_style,
            'relevant diagnosis': diag_other_style,
            'major procedure or operation': proc_major_style,
            'relevant medical services and procedures': proc_other_style
        }

        # Pass the correct CLASS NAMES to the replacement function
        details_map = create_evidence_details_map(selected_note_outputs_df, domain_style_map)

        # Get the single note text as a string
        discharge_note_text = selected_note_metadata["note"].iloc[0]

        # Build the styled components
        corpus_body, tooltips = build_styled_text_components(discharge_note_text, details_map)

        return [
            html.H3(
                f"Showing details for Note ID: {note_id_str}",
                style={'color': '#495057'}),
            corpus_body,
            *tooltips
        ]

    except Exception as e:
        # Error handling remains the same
        tb_str = traceback.format_exc()
        return html.Div([
            html.H3("An Error Occurred", style={'color': 'red'}),
            html.Pre(f"Error: {e}\n\nTraceback:\n{tb_str}", style={'border': '1px solid red', 'padding': '10px'})
        ])
