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

custom_css = Style("""
    body {
        font-family: Arial, sans-serif;
        line-height: 1.6;
        color: #333;
        max-width: 800px;
        margin: 0 auto;
        padding: 20px;
        background-color: #f0f0f0;
    }
    h1 {
        color: #2c3e50;
        text-align: center;
        margin-bottom: 30px;
    }
    form {
        background-color: #ffffff;
        padding: 20px;
        border-radius: 8px;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
    label {
        display: block;
        margin-bottom: 5px;
        font-weight: bold;
        color: #2c3e50;
    }
    select, button {
        width: 100%;
        padding: 10px;
        margin-bottom: 15px;
        border: 1px solid #ddd;
        border-radius: 4px;
        background-color: #fff;
        color: #333;
    }
    button {
        background-color: #3498db;
        color: white;
        font-weight: bold;
        cursor: pointer;
        transition: background-color 0.3s;
    }
    button:hover {
        background-color: #2980b9;
    }
    #results {
        margin-top: 30px;
        background-color: #ffffff;
        padding: 20px;
        border-radius: 8px;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
    #results h2 {
        color: #2c3e50;
        border-bottom: 2px solid #3498db;
        padding-bottom: 10px;
    }
    #results p {
        margin-bottom: 10px;
        color: #333;
    }
    #results hr {
        border: none;
        border-top: 1px solid #eee;
        margin: 20px 0;
    }
    a {
        color: #3498db;
        text-decoration: none;
    }
    a:hover {
        text-decoration: underline;
    }
""")

# 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()

    except sqlite3.Error as e:
        print(f"An error occurred: {e}")
        return {
            'combinations': [],
            'drugs': [],
            'brands': [],
            'formulations': [],
            'indications': [],
            'treatment_phases': [],
            'hospital_types': []
        }

    # 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, ht.name as hospital_type
                      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
                      JOIN hospital_types ht ON c.hospital_type_id = ht.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()]
    
    cursor.execute('SELECT name FROM hospital_types')
    hospital_types = [row[0] for row in cursor.fetchall()]

    conn.close()

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

biologics_data = load_data()

app, rt = fast_app()

def search_biologics(drug, brand, formulation, indication, treatment_phase, hospital_type):
    field_indices = {
        'drug': 1,
        'brand': 2,
        'formulation': 3,
        'indication': 4,
        'treatment_phase': 5,
        'hospital_type': 9
    }
    
    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, 'hospital_type': hospital_type}.items()
    )]
    
    if not results:
        return "No results found."
    
    output = ""
    for item in results:
        output += f"""
        <div class="result-item">
            <h2>{item[1]} ({item[2]})</h2>
            <p><strong>PBS Code:</strong> <a href="https://www.pbs.gov.au/medicine/item/{item[0]}" target="_blank">{item[0]}</a></p>
            <p><strong>Formulation:</strong> {item[3]}</p>
            <p><strong>Indication:</strong> {item[4]}</p>
            <p><strong>Treatment Phase:</strong> {item[5]}</p>
            <p><strong>Streamlined Code:</strong> {item[6] or 'N/A'}</p>
            <p><strong>Authority Method:</strong> {item[8]}</p>
            <p><strong>Online Application:</strong> {'Yes' if item[7] else 'No'}</p>
            <p><strong>Hospital Type:</strong> {item[9]}</p>
        </div>
        <hr>
        """
    
    return output

def update_options(drug, brand, formulation, indication, treatment_phase, hospital_type):
    selected = {
        'drug': drug,
        'brand': brand,
        'formulation': formulation,
        'indication': indication,
        'treatment_phase': treatment_phase,
        'hospital_type': hospital_type
    }
    
    field_indices = {
        'drug': 1,
        'brand': 2,
        'formulation': 3,
        'indication': 4,
        'treatment_phase': 5,
        'hospital_type': 9
    }
    
    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)),
        'hospital_types': sorted(set(combo[9] for combo in filtered))
    }
    
    return options

@rt('/')
def get():
    return Div(
        custom_css,
        H1("Biologics Prescriber Helper"),
        Form(
            Div(
                Label("Drug:"),
                Select(Option("All", value=""), *[Option(drug, value=drug) for drug in biologics_data['drugs']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Brand:"),
                Select(Option("All", value=""), *[Option(brand, value=brand) for brand in biologics_data['brands']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Formulation:"),
                Select(Option("All", value=""), *[Option(formulation, value=formulation) for formulation in biologics_data['formulations']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Indication:"),
                Select(Option("All", value=""), *[Option(indication, value=indication) for indication in biologics_data['indications']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Treatment Phase:"),
                Select(Option("All", value=""), *[Option(phase, value=phase) for phase in biologics_data['treatment_phases']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Hospital Type:"),
                Select(Option("All", value=""), *[Option(ht, value=ht) for ht in biologics_data['hospital_types']], name="hospital_type", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase'],[name='hospital_type']")
            ),
            Div(
                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 Div(
        custom_css,
        Form(
            Div(
                Label("Drug:"),
                Select(Option("All", value=""), *[Option(drug, value=drug) for drug in biologics_data['drugs']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Brand:"),
                Select(Option("All", value=""), *[Option(brand, value=brand) for brand in biologics_data['brands']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Formulation:"),
                Select(Option("All", value=""), *[Option(formulation, value=formulation) for formulation in biologics_data['formulations']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Indication:"),
                Select(Option("All", value=""), *[Option(indication, value=indication) for indication in biologics_data['indications']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Treatment Phase:"),
                Select(Option("All", value=""), *[Option(phase, value=phase) for phase in biologics_data['treatment_phases']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Hospital Type:"),
                Select(Option("All", value=""), *[Option(ht, value=ht) for ht in biologics_data['hospital_types']], name="hospital_type", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase'],[name='hospital_type']")
            ),
            Div(
                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 = '', hospital_type: str = ''):
    options = update_options(drug, brand, formulation, indication, treatment_phase, hospital_type)
    return Div(
        custom_css,
        Form(
            Div(
                Label("Drug:"),
                Select(Option("All", value=""), *[Option(d, value=d, selected=(d == drug)) for d in options['drugs']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Brand:"),
                Select(Option("All", value=""), *[Option(b, value=b, selected=(b == brand)) for b in options['brands']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Formulation:"),
                Select(Option("All", value=""), *[Option(f, value=f, selected=(f == formulation)) for f in options['formulations']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Indication:"),
                Select(Option("All", value=""), *[Option(i, value=i, selected=(i == indication)) for i in options['indications']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Treatment Phase:"),
                Select(Option("All", value=""), *[Option(p, value=p, selected=(p == treatment_phase)) for p in options['treatment_phases']], 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'],[name='hospital_type']")
            ),
            Div(
                Label("Hospital Type:"),
                Select(Option("All", value=""), *[Option(ht, value=ht, selected=(ht == hospital_type)) for ht in options['hospital_types']], name="hospital_type", hx_get="/update_options", hx_target="#options", hx_trigger="change", hx_include="[name='drug'],[name='brand'],[name='formulation'],[name='indication'],[name='treatment_phase'],[name='hospital_type']")
            ),
            Div(
                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 = '', hospital_type: str = ''):
    results = search_biologics(drug, brand, formulation, indication, treatment_phase, hospital_type)
    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()