from fasthtml.common import *
import sqlite3
import os
import datetime
from apscheduler.schedulers.background import BackgroundScheduler
from pbs_data import PBSPublicDataAPIClient
import os
from fasthtml_hf import setup_hf_backup  # Add this line

# Database file
DB_FILE = 'rheumatology_biologics_data.db'

if not os.path.exists(DB_FILE):
    print(f"Database file {DB_FILE} does not exist!")

def load_data():
    try:
        conn = sqlite3.connect(DB_FILE)
        cursor = conn.cursor()
        # ... (rest of the function)
    except sqlite3.Error as e:
        print(f"An error occurred: {e}")
        return {
            'combinations': [],
            'drugs': [],
            'brands': [],
            'formulations': [],
            'indications': [],
            'treatment_phases': []
        }

    # Fetch all data
    cursor.execute('''SELECT c.pbs_code, d.name as drug, b.name as brand, f.name as formulation,
                      i.name as indication, tp.name as treatment_phase, c.streamlined_code,
                      c.online_application, c.authority_method
                      FROM combinations c
                      JOIN drugs d ON c.drug_id = d.id
                      JOIN brands b ON c.brand_id = b.id
                      JOIN formulations f ON c.formulation_id = f.id
                      JOIN indications i ON c.indication_id = i.id
                      JOIN treatment_phases tp ON c.treatment_phase_id = tp.id''')
    
    data = cursor.fetchall()
    
    # Fetch distinct values for dropdowns
    cursor.execute('SELECT name FROM drugs')
    drugs = [row[0] for row in cursor.fetchall()]
    
    cursor.execute('SELECT name FROM brands')
    brands = [row[0] for row in cursor.fetchall()]
    
    cursor.execute('SELECT name FROM formulations')
    formulations = [row[0] for row in cursor.fetchall()]
    
    cursor.execute('SELECT name FROM indications')
    indications = [row[0] for row in cursor.fetchall()]
    
    cursor.execute('SELECT name FROM treatment_phases')
    treatment_phases = [row[0] for row in cursor.fetchall()]

    conn.close()

    return {
        'combinations': data,
        'drugs': drugs,
        'brands': brands,
        'formulations': formulations,
        'indications': indications,
        'treatment_phases': treatment_phases
    }

biologics_data = load_data()

app, rt = fast_app()

def search_biologics(drug, brand, formulation, indication, treatment_phase):
    field_indices = {
        'drug': 1,
        'brand': 2,
        'formulation': 3,
        'indication': 4,
        'treatment_phase': 5
    }
    
    results = [combo for combo in biologics_data['combinations'] if all(
        not value or combo[field_indices[key]] == value
        for key, value in {'drug': drug, 'brand': brand, 'formulation': formulation, 'indication': indication, 'treatment_phase': treatment_phase}.items()
    )]
    
    if not results:
        return "No results found."
    
    output = ""
    for item in results:
        output += f"""
        <h2>{item[1]} ({item[2]})</h2>
        <p>PBS Code: <a href="https://www.pbs.gov.au/medicine/item/{item[0]}" target="_blank">{item[0]}</a></p>
        <p>Formulation: {item[3]}</p>
        <p>Indication: {item[4]}</p>
        <p>Treatment Phase: {item[5]}</p>
        <p>Streamlined Code: {item[6] or 'N/A'}</p>
        <p>Authority Method: {item[8]}</p>
        <p>Online Application: {'Yes' if item[7] else 'No'}</p>
        <hr>
        """
    
    return output

def update_options(drug, brand, formulation, indication, treatment_phase):
    selected = {
        'drug': drug,
        'brand': brand,
        'formulation': formulation,
        'indication': indication,
        'treatment_phase': treatment_phase
    }
    
    field_indices = {
        'drug': 1,
        'brand': 2,
        'formulation': 3,
        'indication': 4,
        'treatment_phase': 5
    }
    
    filtered = [combo for combo in biologics_data['combinations'] if all(
        not selected[key] or combo[field_indices[key]] == selected[key]
        for key in selected
    )]
    
    options = {
        'drugs': sorted(set(combo[1] for combo in filtered)),
        'brands': sorted(set(combo[2] for combo in filtered)),
        'formulations': sorted(set(combo[3] for combo in filtered)),
        'indications': sorted(set(combo[4] for combo in filtered)),
        'treatment_phases': sorted(set(combo[5] for combo in filtered))
    }
    
    return options

@rt('/')
def get():
    return Titled("Biologics Prescriber Helper",
        Form(
            Select(Option("All", value=""), *[Option(drug, value=drug) for drug in biologics_data['drugs']], label="Drug", name="drug", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
            Select(Option("All", value=""), *[Option(brand, value=brand) for brand in biologics_data['brands']], label="Brand", name="brand", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
            Select(Option("All", value=""), *[Option(formulation, value=formulation) for formulation in biologics_data['formulations']], label="Formulation", name="formulation", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
            Select(Option("All", value=""), *[Option(indication, value=indication) for indication in biologics_data['indications']], label="Indication", name="indication", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
            Select(Option("All", value=""), *[Option(phase, value=phase) for phase in biologics_data['treatment_phases']], label="Treatment Phase", name="treatment_phase", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
            Button("Search", type="submit"),
            Button("Reset", hx_get="/reset", hx_target="#options"),
            hx_post="/search",
            hx_target="#results",
            id="options"
        ),
        Div(id="results")
    )

@rt('/reset')
def get():
    return Form(
        Select(Option("All", value=""), *[Option(drug, value=drug) for drug in biologics_data['drugs']], label="Drug", name="drug", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Select(Option("All", value=""), *[Option(brand, value=brand) for brand in biologics_data['brands']], label="Brand", name="brand", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Select(Option("All", value=""), *[Option(formulation, value=formulation) for formulation in biologics_data['formulations']], label="Formulation", name="formulation", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Select(Option("All", value=""), *[Option(indication, value=indication) for indication in biologics_data['indications']], label="Indication", name="indication", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Select(Option("All", value=""), *[Option(phase, value=phase) for phase in biologics_data['treatment_phases']], label="Treatment Phase", name="treatment_phase", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Button("Search", type="submit"),
        Button("Reset", hx_get="/reset", hx_target="#options"),
        hx_post="/search",
        hx_target="#results",
        id="options"
    )

@rt('/update_options')
def get(drug: str = '', brand: str = '', formulation: str = '', indication: str = '', treatment_phase: str = ''):
    options = update_options(drug, brand, formulation, indication, treatment_phase)
    return Form(
        Select(Option("All", value=""), *[Option(d, value=d, selected=(d == drug)) for d in options['drugs']], label="Drug", name="drug", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Select(Option("All", value=""), *[Option(b, value=b, selected=(b == brand)) for b in options['brands']], label="Brand", name="brand", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Select(Option("All", value=""), *[Option(f, value=f, selected=(f == formulation)) for f in options['formulations']], label="Formulation", name="formulation", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Select(Option("All", value=""), *[Option(i, value=i, selected=(i == indication)) for i in options['indications']], label="Indication", name="indication", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Select(Option("All", value=""), *[Option(p, value=p, selected=(p == treatment_phase)) for p in options['treatment_phases']], label="Treatment Phase", name="treatment_phase", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase']"),
        Button("Search", type="submit"),
        Button("Reset", hx_get="/reset", hx_target="#options"),
        hx_post="/search",
        hx_target="#results",
        id="options"
    )

@rt('/search')
def post(drug: str = '', brand: str = '', formulation: str = '', indication: str = '', treatment_phase: str = ''):
    results = search_biologics(drug, brand, formulation, indication, treatment_phase)
    return results


def update_data():
    print(f"Updating data at {datetime.datetime.now()}")
    client = PBSPublicDataAPIClient("2384af7c667342ceb5a736fe29f1dc6b", rate_limit=0.2)
    try:
        data = client.fetch_rheumatology_biologics_data()
        client.save_data_to_sqlite(data, DB_FILE)
        print("Data updated successfully")
        global biologics_data
        biologics_data = load_data()
    except Exception as e:
        print(f"An error occurred while updating data: {str(e)}")

# Set up the scheduler
scheduler = BackgroundScheduler()
scheduler.add_job(func=update_data, trigger='cron', day='1', hour='0', minute='0')
scheduler.start()

setup_hf_backup(app)  # Add this line
serve()