import gradio as gr
import time

def calculate_price(payment_mode, tokens, plan, custom_price, file):
    if payment_mode == "Pay as you go":
        price = round(tokens * 0.01, 2)  # Example: $0.01 per token
        return f"{tokens:,} tokens\nPrice: ${price:.2f}", price
    elif payment_mode == "Plan":
        if plan == "Free":
            return "0 tokens\nPrice: $0", 0
        elif plan == "Starter":
            return "100,000 tokens\nPrice: $15", 15
        elif plan == "Pro":
            return "500,000 tokens\nPrice: $30", 30
        elif plan == "Custom":
            return f"Custom plan\nPrice: ${custom_price}", float(custom_price or 0)
    elif file is not None:
        # Simulate token count from file size
        tokens = 1000  # Replace it with real calculation
        price = round(tokens * 0.01, 2)
        return f"{tokens:,} tokens\nPrice: ${price:.2f}", price
    return "", 0

def generate_dataset(*args, **kwargs):
    for i in range(5):
        yield f"Generating... ({(i+1)*20}%)", None, (i+1)/5
        time.sleep(0.3)
    yield "Ready! Please pay to download.", "dataset.jsonl", 1.0

with gr.Blocks(
    title="Nexa Data Studio",
    css="""
    body, .gradio-container {
        min-height: 100vh;
        background: #111 !important;
        color: #fff !important;
    }
    .gradio-container {
        max-width: 900px !important;
        margin: 40px auto !important;
        box-shadow: 0 2px 16px #0008;
        border-radius: 16px;
        padding: 32px 32px 24px 32px !important;
        background: #111 !important;
        color: #fff !important;
        display: flex;
        flex-direction: column;
        align-items: center;
    }
    .footer {margin-top: 2em; color: #bbb; font-size: 0.9em; text-align: center;}
    #header {text-align: center;}
    """
) as demo:
    gr.Markdown(
        """
        <div style="display:flex;align-items:center;gap:16px;justify-content:center;">
            <img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" height="40"/>
            <h1 style="margin-bottom:0;">Nexa Data Studio</h1>
        </div>
        <p style="text-align:center;">
        <b>Generate or label scientific datasets for ML research.</b>
        </p>
        """,
        elem_id="header"
    )

    payment_mode = gr.Radio(
        ["Pay as you go", "Plan"],
        label="Payment Mode",
        value="Pay as you go"
    )

    with gr.Row() as payg_row:
        tokens = gr.Slider(100, 100000, value=1000, step=100, label="Tokens Requested")
    with gr.Row(visible=False) as plan_row:
        plan = gr.Dropdown(
            ["Free", "Starter", "Pro", "Custom"],
            label="Plan",
            value="Free"
        )
        custom_price = gr.Number(label="Custom Price ($)", visible=False)

    job_type = gr.Radio(
        ["Generate Dataset", "Label Uploaded Data"],
        label="Job Type",
        value="Generate Dataset"
    )

    with gr.Column(visible=False) as label_col:
        file = gr.File(label="Upload Dataset (.txt or .jsonl)")

    price_info = gr.Textbox(label="Summary", interactive=False)
    download = gr.File(label="Download")
    progress = gr.Slider(0, 1, value=0, step=0.01, label="Progress", interactive=False)
    status = gr.Text(label="Status", interactive=False)

    def update_payment_ui(payment_mode_val, plan_val):
        return (
            gr.update(visible=payment_mode_val == "Pay as you go"),
            gr.update(visible=payment_mode_val == "Plan"),
            gr.update(visible=payment_mode_val == "Plan" and plan_val == "Custom")
        )

    payment_mode.change(
        update_payment_ui,
        inputs=[payment_mode, plan],
        outputs=[payg_row, plan_row, custom_price]
    )
    plan.change(
        lambda p: gr.update(visible=p == "Custom"),
        inputs=plan,
        outputs=custom_price
    )

    def update_label_ui(job_type_val):
        return gr.update(visible=job_type_val == "Label Uploaded Data")
    job_type.change(update_label_ui, inputs=job_type, outputs=label_col)

    def update_summary(payment_mode, tokens, plan, custom_price, file, job_type):
        if job_type == "Label Uploaded Data" and file is not None:
            return calculate_price("Label", tokens, plan, custom_price, file)[0]
        return calculate_price(payment_mode, tokens, plan, custom_price, file)[0]

    inputs = [payment_mode, tokens, plan, custom_price, file, job_type]
    gr.Button("Generate", elem_id="generate-btn", variant="primary").click(
        generate_dataset,
        inputs=inputs,
        outputs=[status, download, progress]
    )
    gr.Button("Update Summary").click(
        update_summary,
        inputs=inputs,
        outputs=price_info
    )

    gr.Markdown(
        f"""
        <div class="footer">
        &copy; {time.strftime("%Y")} Nexa Data Studio &mdash; Powered by Hugging Face Spaces<br>
        For support, contact <a href="mailto:support@nexadatastudio.com">support@nexadatastudio.com</a>
        </div>
        """
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
    print("Nexa Data Studio is running at http://localhost:7860")