# File: pages/distribusi_kasus_demografi.py (Revisi Final dengan Bulan Dinamis)

from dash import dcc, html, Input, Output, callback, no_update, State
import dash_bootstrap_components as dbc
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from sqlalchemy import select, distinct, func, and_
import io

# Impor library yang dibutuhkan untuk PDF
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, PageBreak
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.lib.enums import TA_CENTER, TA_LEFT, TA_JUSTIFY
import plotly.io as pio

# Impor engine Anda dari file database.py
from database import engine, detail_penyakit

# -----------------------------------------------------------------------------
# PENGATURAN AWAL
# -----------------------------------------------------------------------------
try:
    # Pengaturan agar unduh PDF tidak error
    pio.kaleido.scope.mathjax = None
except (AttributeError, ImportError):
    pass

# -----------------------------------------------------------------------------
# FUNGSI HELPER (ANALISIS TEKS)
# -----------------------------------------------------------------------------
def get_filter_text(pusk, tahun, bulan):
    pusk_txt = "Seluruh Puskesmas" if not pusk else ", ".join(pusk)
    tahun_txt = "Seluruh Tahun" if not tahun else ", ".join(map(str, sorted(tahun)))
    bulan_txt = "Seluruh Bulan" if not bulan else ", ".join(sorted(bulan))
    return f"Menampilkan data dari: {pusk_txt} | Tahun: {tahun_txt} | Bulan: {bulan_txt}"

def generate_gender_analysis_text(total_lk, total_pr):
    if total_lk == 0 and total_pr == 0: return "Tidak ada data kasus untuk dianalisis."
    total_kasus = total_lk + total_pr
    persen_lk = (total_lk / total_kasus * 100) if total_kasus > 0 else 0
    persen_pr = (total_pr / total_kasus * 100) if total_kasus > 0 else 0
    kesimpulan = ""
    if persen_lk > persen_pr:
        kesimpulan = f"lebih banyak menyerang laki-laki ({persen_lk:.1f}%) dibandingkan perempuan ({persen_pr:.1f}%)"
    elif persen_pr > persen_lk:
        kesimpulan = f"lebih banyak menyerang perempuan ({persen_pr:.1f}%) dibandingkan laki-laki ({persen_lk:.1f}%)"
    else:
        kesimpulan = f"memiliki distribusi yang seimbang antara laki-laki dan perempuan"
    return f"Total kasus tercatat: **{int(total_kasus):,}**. Secara umum, kasus penyakit {kesimpulan}."

def generate_age_analysis_text(baru_data, lama_data):
    total_kasus_umur_baru = sum(baru_data.values())
    total_kasus_umur_lama = sum(lama_data.values())
    total_kasus_umur = total_kasus_umur_baru + total_kasus_umur_lama
    if total_kasus_umur == 0: return "Tidak ada data kasus untuk dianalisis."
    
    # Hitung total kasus per kelompok umur
    total_per_kelompok = {k: baru_data.get(k, 0) + lama_data.get(k, 0) for k in set(baru_data) | set(lama_data)}
    
    max_total_kelompok = max(total_per_kelompok, key=total_per_kelompok.get, default=None)
    max_baru_kelompok = max(baru_data, key=baru_data.get, default=None) if total_kasus_umur_baru > 0 else "Tidak ada"
    
    return f"Kelompok umur dengan total kasus tertinggi adalah **{max_total_kelompok}<**. Kasus baru terbanyak ditemukan pada kelompok **{max_baru_kelompok}**."

# -----------------------------------------------------------------------------
# LAYOUT HALAMAN
# -----------------------------------------------------------------------------
layout = dbc.Container([
    dcc.Store(id='dkd-data-store'),
    dcc.Download(id="dkd-download-pdf"),

    # Header dan Tombol Unduh
    dbc.Row([
        dbc.Col(html.H3("Distribusi Kasus Demografi", className="mt-4 mb-4"), md=9),
        dbc.Col(dbc.Button("Unduh Halaman (PDF)", id="dkd-btn-download", color="primary", className="mt-4 float-end", disabled=True), md=3)
    ], align="center"),
    
    # Panel Filter Utama
    dbc.Card(dbc.CardBody([
        dbc.Row([
            dbc.Col([dbc.Label("Pilih Puskesmas:"), dcc.Dropdown(id='dkd-pusk-filter', multi=True, placeholder="Pilih...")], md=4),
            dbc.Col([dbc.Label("Pilih Tahun:"), dcc.Dropdown(id='dkd-tahun-filter', multi=True, placeholder="Pilih...")], md=4),
            dbc.Col([
                dbc.Label("Pilih Bulan:"), 
                dcc.Dropdown(id='dkd-bulan-filter', multi=True, placeholder="Pilih Tahun dulu...", disabled=True)
            ], md=4)
        ])
    ]), className="mb-3 shadow-sm"),
    
    html.Div(id='dkd-filter-summary-text', className="text-center text-muted fst-italic mb-3"),
    
    dcc.Loading(id="dkd-loading-main", type="dot", children=[
        dbc.Tabs(id="dkd-tabs", active_tab='tab-gender', children=[
            dbc.Tab(label="Distribusi Jenis Kelamin", tab_id="tab-gender", children=html.Div(id='tab-content-gender')),
            dbc.Tab(label="Distribusi Kelompok Umur", tab_id="tab-age", children=html.Div(id='tab-content-age')),
        ])
    ])
], fluid=True)

# -----------------------------------------------------------------------------
# CALLBACKS
# -----------------------------------------------------------------------------

# Callback 1: Memuat filter Puskesmas dan Tahun saat halaman dibuka
@callback(
    Output('dkd-pusk-filter', 'options'),
    Output('dkd-tahun-filter', 'options'),
    Input('url', 'pathname')
)
def dkd_load_main_filters(pathname):
    if pathname != '/distribusi_kasus_demografi':
        return no_update, no_update
    try:
        with engine.connect() as conn:
            pusk_list = [row[0] for row in conn.execute(select(distinct(detail_penyakit.c.kode_pusk)).where(detail_penyakit.c.kode_pusk.is_not(None), detail_penyakit.c.kode_pusk != '').order_by(detail_penyakit.c.kode_pusk)).fetchall() if row[0]]
            tahun_list = [row[0] for row in conn.execute(select(distinct(detail_penyakit.c.tahun)).where(detail_penyakit.c.tahun.is_not(None)).order_by(detail_penyakit.c.tahun.desc())).fetchall() if row[0]]
        pusk_options = [{'label': p, 'value': p} for p in pusk_list]
        tahun_options = [{'label': str(t), 'value': t} for t in tahun_list]
        return pusk_options, tahun_options
    except Exception as e:
        print(f"Error saat memuat filter demografi: {e}")
        return [], []

# Callback 2: Memperbarui filter Bulan berdasarkan Tahun yang dipilih (CHAINED CALLBACK)
@callback(
    Output('dkd-bulan-filter', 'options'),
    Output('dkd-bulan-filter', 'disabled'),
    Output('dkd-bulan-filter', 'value'), # Otomatis reset value bulan jika tahun berubah
    Input('dkd-tahun-filter', 'value')
)
def dkd_update_bulan_filter(selected_tahun):
    if not selected_tahun:
        return [], True, []

    nama_bulan = {
        '01': 'Januari', '02': 'Februari', '03': 'Maret', '04': 'April',
        '05': 'Mei', '06': 'Juni', '07': 'Juli', '08': 'Agustus',
        '09': 'September', '10': 'Oktober', '11': 'November', '12': 'Desember'
    }
    try:
        with engine.connect() as conn:
            stmt = select(distinct(detail_penyakit.c.bulan)).where(detail_penyakit.c.tahun.in_(selected_tahun)).order_by(detail_penyakit.c.bulan)
            bulan_list = [row[0] for row in conn.execute(stmt).fetchall() if row[0]]
        
        bulan_options = [{'label': nama_bulan.get(b, b), 'value': b} for b in bulan_list]
        return bulan_options, False, []

    except Exception as e:
        print(f"Error saat memuat filter bulan dinamis: {e}")
        return [], True, []

# Callback 3: Memperbarui semua konten berdasarkan filter
@callback(
    Output('tab-content-gender', 'children'),
    Output('tab-content-age', 'children'),
    Output('dkd-filter-summary-text', 'children'),
    Output('dkd-data-store', 'data'),
    Output('dkd-btn-download', 'disabled'),
    Input('dkd-pusk-filter', 'value'),
    Input('dkd-tahun-filter', 'value'),
    Input('dkd-bulan-filter', 'value')
)
def update_demografi_tabs(selected_pusk, selected_tahun, selected_bulan):
    if not all([selected_pusk, selected_tahun, selected_bulan]):
        msg = html.P("Silakan pilih Puskesmas, Tahun, dan Bulan untuk memulai analisis.", className="text-center text-primary mt-5")
        return msg, msg, "", None, True

    filter_summary_text = get_filter_text(selected_pusk, selected_tahun, selected_bulan)
    filters = [
        detail_penyakit.c.kode_pusk.in_(selected_pusk), 
        detail_penyakit.c.tahun.in_(selected_tahun), 
        detail_penyakit.c.bulan.in_(selected_bulan)
    ]
    
    usia_cols = ['laki_laki', 'perempuan', 'usia_0_7_hr_baru', 'usia_0_7_hr_lama', 'usia_8_28_hr_baru', 'usia_8_28_hr_lama', 'usia_1bl_1th_baru', 'usia_1bl_1th_lama', 'usia_1_4th_baru', 'usia_1_4th_lama', 'usia_5_9th_baru', 'usia_5_9th_lama', 'usia_10_14th_baru', 'usia_10_14th_lama', 'usia_15_19th_baru', 'usia_15_19th_lama', 'usia_20_44th_baru', 'usia_20_44th_lama', 'usia_45_54th_baru', 'usia_45_54th_lama', 'usia_55_59th_baru', 'usia_55_59th_lama', 'usia_60_69th_baru', 'usia_60_69th_lama', 'usia_70pl_baru', 'usia_70pl_lama']
    stmt = select(*[func.sum(getattr(detail_penyakit.c, col)).label(col) for col in usia_cols]).where(and_(*filters))
    
    with engine.connect() as conn:
        result = conn.execute(stmt).fetchone()

    if not result or all(val is None for val in result):
        msg = dbc.Alert("Tidak ada data ditemukan untuk kriteria yang dipilih.", color="warning", className="m-4")
        return msg, msg, filter_summary_text, None, True

    # --- KONTEN TAB 1: JENIS KELAMIN ---
    total_lk = result.laki_laki or 0
    total_pr = result.perempuan or 0
    fig_gender = go.Figure(data=[go.Pie(labels=['Laki-laki', 'Perempuan'], values=[total_lk, total_pr], hole=.4, marker_colors=['royalblue', 'crimson'])])
    fig_gender.update_layout(title_text='</b>Distribusi Kasus Berdasarkan Jenis Kelamin</b>', template='plotly_white')
    analysis_gender = generate_gender_analysis_text(total_lk, total_pr)
    tab_gender_content = dbc.Row(dbc.Col([dcc.Graph(figure=fig_gender), dbc.Card(dbc.CardBody(dcc.Markdown(analysis_gender)), className="mt-2")]), className="mt-4")

    # --- KONTEN TAB 2: KELOMPOK UMUR ---
    kelompok_map = {'Bayi & Balita (<5th)':['usia_0_7_hr_baru','usia_0_7_hr_lama','usia_8_28_hr_baru','usia_8_28_hr_lama','usia_1bl_1th_baru','usia_1bl_1th_lama','usia_1_4th_baru','usia_1_4th_lama'],'Anak (5-9th)':['usia_5_9th_baru','usia_5_9th_lama'],'Remaja (10-19th)':['usia_10_14th_baru','usia_10_14th_lama','usia_15_19th_baru','usia_15_19th_lama'],'Dewasa (20-59th)':['usia_20_44th_baru','usia_20_44th_lama','usia_45_54th_baru','usia_45_54th_lama','usia_55_59th_baru','usia_55_59th_lama'],'Lansia (60+th)':['usia_60_69th_baru','usia_60_69th_lama','usia_70pl_baru','usia_70pl_lama']}
    baru_data, lama_data = {}, {}
    for kelompok, cols in kelompok_map.items():
        baru_data[kelompok] = sum(getattr(result, col, 0) or 0 for col in cols if 'baru' in col)
        lama_data[kelompok] = sum(getattr(result, col, 0) or 0 for col in cols if 'lama' in col)
    
    fig_age = go.Figure(data=[
        go.Bar(name='Kasus Baru', x=list(baru_data.keys()), y=list(baru_data.values())),
        go.Bar(name='Kasus Lama', x=list(lama_data.keys()), y=list(lama_data.values()))
    ]).update_layout(barmode='group', title_text="</b>Distribusi Kasus Berdasarkan Kelompok Umur</b>", template='plotly_white')
    analysis_age = generate_age_analysis_text(baru_data, lama_data)
    tab_age_content = dbc.Row(dbc.Col([dcc.Graph(figure=fig_age), dbc.Card(dbc.CardBody(dcc.Markdown(analysis_age)), className="mt-2")]), className="mt-4")

    data_to_store = {'gender': {'Laki-laki': total_lk, 'Perempuan': total_pr}, 'age_new': baru_data, 'age_old': lama_data}
    
    return tab_gender_content, tab_age_content, filter_summary_text, data_to_store, False

# Callback 4: Menghasilkan dan Mengunduh PDF
@callback(
    Output("dkd-download-pdf", "data"),
    Input("dkd-btn-download", "n_clicks"),
    State("dkd-data-store", "data"),
    State("dkd-filter-summary-text", "children"),
    prevent_initial_call=True
)
def generate_demografi_pdf(n_clicks, stored_data, filter_text):
    if not stored_data:
        return no_update

    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=(8.5*inch, 11*inch), rightMargin=0.75*inch, leftMargin=0.75*inch, topMargin=0.75*inch, bottomMargin=0.75*inch)
    styles = getSampleStyleSheet()
    style_h1 = ParagraphStyle(name='H1', parent=styles['h1'], alignment=TA_CENTER, fontSize=16, spaceAfter=14)
    style_h2 = ParagraphStyle(name='H2', parent=styles['h2'], alignment=TA_LEFT, fontSize=14, spaceBefore=12, spaceAfter=6, fontName='Helvetica-Bold')
    style_analysis = ParagraphStyle(name='Analysis', parent=styles['Normal'], alignment=TA_JUSTIFY, fontSize=10, leading=14, spaceAfter=10)
    story = []

    story.append(Paragraph("Laporan Distribusi Kasus Demografi", style_h1))
    story.append(Paragraph(filter_text, styles['Normal']))
    story.append(Spacer(1, 0.3*inch))

    # Membuat ulang grafik dan analisis dari data yang disimpan
    gender_data = stored_data['gender']; total_lk, total_pr = gender_data['Laki-laki'], gender_data['Perempuan']
    fig_gender = go.Figure(data=[go.Pie(labels=['Laki-laki', 'Perempuan'], values=[total_lk, total_pr], hole=.4, marker_colors=['royalblue', 'crimson'])])
    fig_gender.update_layout(title_text='Distribusi Kasus Berdasarkan Jenis Kelamin', template='plotly_white')
    analysis_gender = generate_gender_analysis_text(total_lk, total_pr)

    baru_data, lama_data = stored_data['age_new'], stored_data['age_old']
    fig_age = go.Figure(data=[go.Bar(name='Kasus Baru', x=list(baru_data.keys()), y=list(baru_data.values())), go.Bar(name='Kasus Lama', x=list(lama_data.keys()), y=list(lama_data.values()))]).update_layout(barmode='group', title_text="Distribusi Kasus Berdasarkan Kelompok Umur", template='plotly_white')
    analysis_age = generate_age_analysis_text(baru_data, lama_data)

    def add_content_to_story(fig, analysis_text):
        try:
            img_bytes = pio.to_image(fig, format="png", width=800, height=450, scale=2)
            img = Image(io.BytesIO(img_bytes), width=7*inch, height=(7*450/800)*inch); img.hAlign = 'CENTER'
            story.append(img)
            story.append(Spacer(1, 0.1*inch))
            story.append(Paragraph(analysis_text, style_analysis))
            story.append(Spacer(1, 0.2*inch))
        except Exception as e:
            story.append(Paragraph(f"<i>Gagal memuat grafik: {e}</i>", styles['Italic']))

    story.append(Paragraph("Distribusi Jenis Kelamin", style_h2)); add_content_to_story(fig_gender, analysis_gender); story.append(PageBreak())
    story.append(Paragraph("Distribusi Kelompok Umur", style_h2)); add_content_to_story(fig_age, analysis_age)
    
    doc.build(story)
    buffer.seek(0)
    return dcc.send_bytes(buffer.getvalue(), "Laporan_Demografi.pdf")