import os
import gradio as gr
from gradio_client import Client, handle_file

# ---------------------------------------------------
# Environment
# ---------------------------------------------------

HF_TOKEN = os.getenv("HF_TOKEN")

if not HF_TOKEN:
    raise ValueError(
        "HF_TOKEN environment variable not set."
    )

# ---------------------------------------------------
# Coldet Client
# ---------------------------------------------------

client = Client(
    "Coldet/backend",
    headers={
        "Authorization": f"Bearer {HF_TOKEN}"
    }
)

# ---------------------------------------------------
# Model List
# ---------------------------------------------------

AVAILABLE_MODELS = [
    "coldet-v1-mini.coldet",
    "coldet-v2-mini.coldet",
]

# ---------------------------------------------------
# Model Credit Costs
# ---------------------------------------------------

MODEL_COSTS = {
    "coldet-v1-mini.coldet": 1,
    "coldet-v2-mini.coldet": 2,
}

# ---------------------------------------------------
# Helpers
# ---------------------------------------------------

def format_error(e):
    return f"Error:\n\n{str(e)}"

# ---------------------------------------------------
# API FUNCTIONS
# ---------------------------------------------------

def create_account(username: str, password: str):

    try:

        result = client.predict(
            username=username,
            password=password,
            api_name="/create_account",
        )

        return f"Account created!\n\nSession ID:\n{result}"

    except Exception as e:
        return format_error(e)

# ---------------------------------------------------

def login(username: str, password: str):

    try:

        result = client.predict(
            username=username,
            password=password,
            api_name="/login",
        )

        return f"Logged in!\n\nSession ID:\n{result}"

    except Exception as e:
        return format_error(e)

# ---------------------------------------------------

def change_password(
    session_id: str,
    old_password: str,
    new_password: str
):

    try:

        result = client.predict(
            session_id=session_id,
            old_password=old_password,
            new_password=new_password,
            api_name="/change_password",
        )

        return result

    except Exception as e:
        return format_error(e)

# ---------------------------------------------------

def create_token(session_id: str):

    try:

        result = client.predict(
            session_id=session_id,
            api_name="/create_token",
        )

        return f"API Key:\n{result}"

    except Exception as e:
        return format_error(e)

# ---------------------------------------------------

def list_keys(session_id: str):

    try:

        result = client.predict(
            session_id=session_id,
            api_name="/list_keys",
        )

        return result

    except Exception as e:

        return {
            "error": str(e)
        }

# ---------------------------------------------------

def delete_key(
    session_id: str,
    api_key: str
):

    try:

        result = client.predict(
            session_id=session_id,
            api_key=api_key,
            api_name="/delete_key",
        )

        return result

    except Exception as e:
        return format_error(e)

# ---------------------------------------------------

def predict(
    image,
    api_key: str,
    model: str
):

    if image is None:

        return {
            "error": "Please upload an image."
        }

    try:

        result = client.predict(
            image=handle_file(image),
            api_key=api_key,
            model=model,
            api_name="/predict",
        )

        result["estimated_cost"] = MODEL_COSTS.get(model, 1)

        return result

    except Exception as e:

        return {
            "error": str(e)
        }

# ---------------------------------------------------
# UI
# ---------------------------------------------------

with gr.Blocks(
    title="Coldet Backend Interface"
) as demo:

    gr.Markdown("# 🚀 Coldet Backend Interface")

    gr.Markdown(
        "Frontend for the Coldet backend API."
    )

    # ---------------------------------------------------
    # ACCOUNT TAB
    # ---------------------------------------------------

    with gr.Tab("🔑 Account"):

        with gr.Row():

            # -----------------------------
            # Create Account
            # -----------------------------

            with gr.Column():

                gr.Markdown("## Create Account")

                create_username = gr.Textbox(
                    label="Username"
                )

                create_password = gr.Textbox(
                    label="Password",
                    type="password"
                )

                create_btn = gr.Button(
                    "Create Account"
                )

                create_output = gr.Textbox(
                    label="Session",
                    lines=4,
                    interactive=False
                )

                create_btn.click(
                    fn=create_account,
                    inputs=[
                        create_username,
                        create_password
                    ],
                    outputs=create_output,
                )

            # -----------------------------
            # Login
            # -----------------------------

            with gr.Column():

                gr.Markdown("## Login")

                login_username = gr.Textbox(
                    label="Username"
                )

                login_password = gr.Textbox(
                    label="Password",
                    type="password"
                )

                login_btn = gr.Button(
                    "Login"
                )

                login_output = gr.Textbox(
                    label="Session",
                    lines=4,
                    interactive=False
                )

                login_btn.click(
                    fn=login,
                    inputs=[
                        login_username,
                        login_password
                    ],
                    outputs=login_output,
                )

    # ---------------------------------------------------
    # PASSWORD TAB
    # ---------------------------------------------------

    with gr.Tab("🔒 Password"):

        gr.Markdown("## Change Password")

        password_session = gr.Textbox(
            label="Session ID"
        )

        old_password = gr.Textbox(
            label="Old Password",
            type="password"
        )

        new_password = gr.Textbox(
            label="New Password",
            type="password"
        )

        password_btn = gr.Button(
            "Change Password"
        )

        password_output = gr.Textbox(
            label="Result",
            interactive=False
        )

        password_btn.click(
            fn=change_password,
            inputs=[
                password_session,
                old_password,
                new_password
            ],
            outputs=password_output,
        )

    # ---------------------------------------------------
    # API KEYS TAB
    # ---------------------------------------------------

    with gr.Tab("🔐 API Keys"):

        gr.Markdown("## API Key Management")

        key_session = gr.Textbox(
            label="Session ID",
            lines=2
        )

        with gr.Row():

            create_key_btn = gr.Button(
                "Create API Key"
            )

            refresh_keys_btn = gr.Button(
                "Refresh Keys"
            )

        token_output = gr.Textbox(
            label="New API Key",
            interactive=False,
            lines=3
        )

        keys_output = gr.JSON(
            label="Your Keys"
        )

        create_key_btn.click(
            fn=create_token,
            inputs=key_session,
            outputs=token_output,
        )

        refresh_keys_btn.click(
            fn=list_keys,
            inputs=key_session,
            outputs=keys_output,
        )

        gr.Markdown("## Delete API Key")

        delete_key_input = gr.Textbox(
            label="API Key"
        )

        delete_key_btn = gr.Button(
            "Delete Key"
        )

        delete_output = gr.Textbox(
            label="Delete Result",
            interactive=False
        )

        delete_key_btn.click(
            fn=delete_key,
            inputs=[
                key_session,
                delete_key_input
            ],
            outputs=delete_output,
        )

    # ---------------------------------------------------
    # INFERENCE TAB
    # ---------------------------------------------------

    with gr.Tab("🤖 Inference"):

        gr.Markdown("## Run Inference")

        gr.Markdown(
            "\n".join(
                [
                    f"- `{model}` → {cost} credit(s)"
                    for model, cost in MODEL_COSTS.items()
                ]
            )
        )

        with gr.Row():

            with gr.Column():

                image_input = gr.Image(
                    label="Image",
                    type="filepath"
                )

                api_key_input = gr.Textbox(
                    label="API Key"
                )

                model_input = gr.Dropdown(
                    label="Model",
                    choices=AVAILABLE_MODELS,
                    value=AVAILABLE_MODELS[0],
                )

                predict_btn = gr.Button(
                    "Run Inference"
                )

            with gr.Column():

                predict_output = gr.JSON(
                    label="Result"
                )

        predict_btn.click(
            fn=predict,
            inputs=[
                image_input,
                api_key_input,
                model_input
            ],
            outputs=predict_output,
        )

# ---------------------------------------------------
# Launch
# ---------------------------------------------------

demo.queue()

if __name__ == "__main__":

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False
    )