from tinysql_model_demo import model_demo
from tinysql_dataset_viewer import dataset_viewer
import gradio as gr

custom_css = """
/* Force black background on EVERYTHING first */
* {
    background-color: #0A0A0A !important;
}

/* Then selectively override what needs different colors */

/* Dark gray for inputs and tables */
input, textarea, select {
    background-color: #1A1A1A !important;
    color: #E0E0E0 !important;
    border: 1px solid #3A3A3A !important;
}

/* Tables - dark gray background */
table, .dataframe {
    background-color: #1A1A1A !important;
}

table td, table th, .dataframe td, .dataframe th {
    background-color: #1A1A1A !important;
    color: #E0E0E0 !important;
}

/* CRITICAL: Text inside tables should be transparent */
table td div, table td span, table td p,
table th div, table th span, table th p {
    background-color: transparent !important;
}

/* Code editors - dark gray */
.code, pre, code {
    background-color: #1A1A1A !important;
}

/* Text inside code blocks transparent */
.code *, pre *, code * {
    background-color: transparent !important;
}

/* Keep gradient boxes with their inline styles */
div[style*="background: linear-gradient"],
div[style*="background: #2A2A2A"],
div[style*="background: #3A3A3A"] {
    background: revert !important;
}

/* Text inside gradient boxes transparent */
div[style*="background"] span,
div[style*="background"] div:not([style*="background"]),
div[style*="background"] p {
    background-color: transparent !important;
}

/* All text white */
*, h1, h2, h3, h4, h5, h6, p, span, div {
    color: #E0E0E0 !important;
}

/* Tabs */
.tab-nav {
    background: #0A0A0A !important;
}

.tab-nav button {
    font-size: 1.8rem !important;
    font-weight: 900 !important;
    color: #888 !important;
    background: transparent !important;
}

.tab-nav button.selected {
    background: #FF6B4A !important;
    color: white !important;
}

/* Orange labels */
label span {
    background: #FF6B4A !important;
    color: white !important;
    padding: 0.4rem 0.8rem !important;
    border-radius: 6px !important;
    font-weight: 700 !important;
}

/* Orange buttons */
.primary {
    background: #FF6B4A !important;
    color: white !important;
}

/* Orange sliders */
input[type="range"]::-webkit-slider-thumb {
    background: #FF6B4A !important;
}

input[type="range"]::-moz-range-thumb {
    background: #FF6B4A !important;
}
"""

with gr.Blocks(css=custom_css, title="TinySQL Demo", theme=gr.themes.Base()) as demo:
    gr.HTML("""
        <div style="text-align: center; padding: 2rem 0; background: #0A0A0A;">
            <h1 style="font-size: 2.5rem; font-weight: 700; color: #FF6B4A; margin: 0;">TinySQL</h1>
            <p style="color: #999; font-size: 1rem; margin-top: 0.5rem;">Mechanistic Interpretability for Text-to-SQL</p>
        </div>
    """)
    
    shared_instruction = gr.State("")
    shared_schema = gr.State("")
    
    with gr.Tabs():
        with gr.Tab("Dataset Viewer"):
            viewer_components = dataset_viewer(shared_instruction, shared_schema)
        
        with gr.Tab("Model Demo"):
            model_components = model_demo(shared_instruction, shared_schema)
    
    gr.HTML("""
        <div style="background: #2A2A2A; border-radius: 12px; padding: 2rem; margin: 3rem auto; max-width: 900px;">
            <div style="color: #FF6B4A; font-weight: 700; font-size: 1.2rem; text-align: center; margin-bottom: 1.5rem;">Cite Our Work</div>
            <pre style="background: #1A1A1A; color: #D0D0D0; padding: 1.5rem; border-radius: 8px; overflow-x: auto; font-family: monospace; font-size: 0.85rem; margin: 0;">@misc{harrasse2025tinysqlprogressivetexttosqldataset,
    title={TinySQL: A Progressive Text-to-SQL Dataset for Mechanistic Interpretability Research}, 
    author={Abir Harrasse and Philip Quirke and Clement Neo and Dhruv Nathawani and Luke Marks and Amir Abdullah},
    year={2025},
    eprint={2503.12730},
    archivePrefix={arXiv},
    primaryClass={cs.LG},
    url={https://arxiv.org/abs/2503.12730}
}</pre>
        </div>
    """)
    
    gr.HTML("""
        <div style="text-align: center; padding: 2rem 0; background: #0A0A0A;">
            <p style="color: #999; font-size: 0.9rem;">
                <a href="https://arxiv.org/abs/2503.12730" style="color: #FF6B4A; text-decoration: none;">Paper</a> • 
                <a href="https://github.com/withmartian/TinySQL" style="color: #FF6B4A; text-decoration: none;">Code</a> • 
                <a href="https://huggingface.co/collections/withmartian/tinysql-6760e92748b63fa56a6ffc9f" style="color: #FF6B4A; text-decoration: none;">Dataset</a>
            </p>
        </div>
    """)

if __name__ == "__main__":
    demo.launch(ssr_mode=False)