import gradio as gr
import torch
import spaces

from transformers import AutoTokenizer, AutoModelForCausalLM


# ============================================================
# MODEL
# ============================================================

MODEL_ID = "InserloftResearch/Parrot-Coder"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.float16 if DEVICE == "cuda" else torch.float32
)

model.to(DEVICE)
model.eval()

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token


# ============================================================
# GENERATION
# ============================================================

@spaces.GPU
def generate_code(prompt, max_new_tokens, temperature, top_p):

    if not prompt or not prompt.strip():
        return ""

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=1024
    )

    inputs = {
        k: v.to(DEVICE)
        for k, v in inputs.items()
    }

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )

    text = tokenizer.decode(
        output[0],
        skip_special_tokens=True
    )

    if text.startswith(prompt):
        text = text[len(prompt):]

    return text.strip()


# ============================================================
# CSS
# ============================================================

CSS = r"""

/* ==========================================================
   ROOT
========================================================== */

:root {
    color-scheme: dark;
}

html,
body {
    margin: 0 !important;
    padding: 0 !important;

    width: 100% !important;
    height: 100% !important;

    background: #050505 !important;
    color: #f2f2f2 !important;

    overflow: hidden !important;
}

body {
    font-family:
        Inter,
        -apple-system,
        BlinkMacSystemFont,
        "Segoe UI",
        sans-serif !important;
}

.gradio-container {
    width: 100vw !important;
    max-width: none !important;

    height: 100vh !important;
    min-height: 100vh !important;

    margin: 0 !important;
    padding: 0 !important;

    background: #050505 !important;
}

.gradio-container > .main {
    width: 100% !important;
    max-width: none !important;

    padding: 0 !important;
    margin: 0 !important;
}

.contain {
    width: 100% !important;
    max-width: none !important;
}

footer,
.built-with {
    display: none !important;
}


/* ==========================================================
   RESET GRADIO
========================================================== */

.block,
.form,
.gr-box,
.gr-panel,
.gr-group,
.gr-input,
.gr-button {
    box-shadow: none !important;
    border-radius: 0 !important;
}


/* ==========================================================
   APP
========================================================== */

#app {
    width: 100%;
    height: 100vh;

    display: grid;

    grid-template-rows:
        68px
        minmax(0, 1fr)
        42px;
}


/* ==========================================================
   HEADER
========================================================== */

#header {
    padding: 0 42px;

    display: flex;
    align-items: center;
    justify-content: space-between;

    border-bottom: 1px solid #1c1c1c;
}

.header-left {
    display: flex;
    align-items: center;
    gap: 16px;

    font-size: 11px;
    font-weight: 600;

    letter-spacing: .13em;
    text-transform: uppercase;
}

.header-separator {
    width: 1px;
    height: 15px;

    background: #303030;
}

.header-muted {
    color: #555;
    font-weight: 400;
}

.header-right {
    color: #505050;

    font-size: 9px;
    letter-spacing: .14em;
    text-transform: uppercase;
}


/* ==========================================================
   CENTER
========================================================== */

#center {
    min-height: 0 !important;

    display: grid;

    grid-template-rows:
        190px
        minmax(0, 1fr);
}


/* ==========================================================
   HERO
========================================================== */

#hero {
    padding: 30px 42px 34px;

    display: flex;
    align-items: flex-end;
    justify-content: space-between;

    border-bottom: 1px solid #1c1c1c;
}

.hero-kicker {
    margin-bottom: 15px;

    color: #565656;

    font-size: 9px;
    letter-spacing: .17em;
    text-transform: uppercase;
}

.hero-title {
    margin: 0;

    font-size: clamp(58px, 6.2vw, 88px);
    line-height: .86;

    font-weight: 600;
    letter-spacing: -.07em;
}

.hero-copy {
    width: 350px;
    padding-bottom: 4px;

    color: #626262;

    font-size: 11px;
    line-height: 1.7;

    text-align: right;
}


/* ==========================================================
   WORKSPACE
========================================================== */

#workspace {
    min-height: 0 !important;

    display: grid !important;

    grid-template-columns: 1fr 1fr !important;

    gap: 1px !important;

    background: #1d1d1d !important;
}


/* ==========================================================
   PANELS
========================================================== */

.panel {
    min-width: 0 !important;
    min-height: 0 !important;

    background: #080808 !important;

    display: grid !important;

    grid-template-rows:
        48px
        minmax(0, 1fr)
        54px !important;
}

.output-panel {
    grid-template-rows:
        48px
        minmax(0, 1fr) !important;
}


/* ==========================================================
   PANEL HEADER
========================================================== */

.panel-header {
    padding: 0 18px;

    display: flex;
    align-items: center;
    justify-content: space-between;

    border-bottom: 1px solid #1a1a1a;
}

.panel-label {
    color: #999;

    font-size: 9px;
    font-weight: 600;

    letter-spacing: .15em;
    text-transform: uppercase;
}

.panel-info {
    color: #444;

    font-size: 9px;
    letter-spacing: .1em;
    text-transform: uppercase;
}


/* ==========================================================
   TEXT EDITORS
========================================================== */

#prompt-box,
#output-box {
    min-height: 0 !important;
    height: 100% !important;

    border: 0 !important;
}

#prompt-box textarea,
#output-box textarea {
    width: 100% !important;
    height: 100% !important;

    min-height: 0 !important;

    padding: 30px 28px !important;

    background: #080808 !important;

    border: 0 !important;
    border-radius: 0 !important;

    color: #ededed !important;

    font-family:
        "SFMono-Regular",
        "Cascadia Code",
        Consolas,
        monospace !important;

    font-size: 13px !important;
    line-height: 1.85 !important;

    resize: none !important;
}

#prompt-box textarea::placeholder,
#output-box textarea::placeholder {
    color: #3d3d3d !important;
}


/* ==========================================================
   GENERATE AREA
========================================================== */

#button-row {
    padding: 9px !important;

    border-top: 1px solid #1a1a1a;
}

#generate {
    width: 100% !important;
    height: 36px !important;

    border: 1px solid #dedede !important;
    border-radius: 2px !important;

    background: #ededed !important;
    color: #050505 !important;

    font-size: 10px !important;
    font-weight: 650 !important;

    letter-spacing: .09em;
    text-transform: uppercase;

    transition:
        opacity .15s ease,
        transform .15s ease;
}

#generate:hover {
    opacity: .82 !important;
}

#generate:active {
    transform: translateY(1px);
}


/* ==========================================================
   FOOTER
========================================================== */

#footer {
    padding: 0 42px;

    display: flex;
    align-items: center;
    justify-content: space-between;

    color: #414141;

    font-size: 9px;
    letter-spacing: .1em;

    text-transform: uppercase;
}

.footer-right {
    color: #4a4a4a;
}


/* ==========================================================
   HIDDEN SETTINGS
========================================================== */

#settings {
    display: none !important;
}


/* ==========================================================
   MOBILE
========================================================== */

@media (max-width: 800px) {

    html,
    body {
        overflow: auto !important;
    }

    .gradio-container {
        width: 100% !important;
        height: auto !important;
        min-height: 100vh !important;
    }

    #app {
        height: auto;
        min-height: 100vh;

        display: block;
    }

    #header {
        height: 60px;
        padding: 0 18px;
    }

    .header-right {
        display: none;
    }

    #center {
        display: block;
    }

    #hero {
        padding: 48px 18px 44px;

        display: block;
    }

    .hero-title {
        font-size: 54px;
    }

    .hero-copy {
        width: auto;

        margin-top: 24px;

        text-align: left;
    }

    #workspace {
        display: flex !important;
        flex-direction: column !important;
    }

    .panel {
        height: 430px !important;
    }

    #prompt-box textarea,
    #output-box textarea {
        height: 100% !important;
        padding: 24px 20px !important;
    }

    #footer {
        padding: 24px 18px;
    }
}

"""


# ============================================================
# UI
# ============================================================

with gr.Blocks(
    title="Parrot-Coder — Inserloft Research",
    css=CSS,
    theme=gr.themes.Base()
) as demo:

    with gr.Column(elem_id="app"):

        # ----------------------------------------------------
        # HEADER
        # ----------------------------------------------------

        gr.HTML(
            """
            <header id="header">

                <div class="header-left">

                    <span>Inserloft Research</span>

                    <span class="header-separator"></span>

                    <span class="header-muted">
                        Parrot-Coder
                    </span>

                </div>

                <div class="header-right">
                    Experimental Model · 2026
                </div>

            </header>
            """
        )

        with gr.Column(elem_id="center"):

            # ------------------------------------------------
            # HERO
            # ------------------------------------------------

            gr.HTML(
                """
                <section id="hero">

                    <div>

                        <div class="hero-kicker">
                            Research / Models / 001
                        </div>

                        <h1 class="hero-title">
                            Parrot-Coder
                        </h1>

                    </div>

                    <div class="hero-copy">
                        Compact code generation for fast
                        experimentation and research.
                        <br>
                        Developed by Inserloft Research.
                    </div>

                </section>
                """
            )

            # ------------------------------------------------
            # WORKSPACE
            # ------------------------------------------------

            with gr.Row(elem_id="workspace"):

                # PROMPT
                with gr.Column(
                    elem_classes=["panel"]
                ):

                    gr.HTML(
                        """
                        <div class="panel-header">

                            <span class="panel-label">
                                Prompt
                            </span>

                            <span class="panel-info">
                                Input
                            </span>

                        </div>
                        """
                    )

                    prompt = gr.Textbox(
                        show_label=False,
                        show_copy_button=False,
                        placeholder=(
                            "Write an instruction...\n\n"
                            "Example:\n"
                            "Create a Python function that checks "
                            "whether a number is prime."
                        ),
                        elem_id="prompt-box"
                    )

                    with gr.Row(elem_id="button-row"):
                        generate = gr.Button(
                            "Generate",
                            elem_id="generate"
                        )

                # OUTPUT
                with gr.Column(
                    elem_classes=["panel", "output-panel"]
                ):

                    gr.HTML(
                        """
                        <div class="panel-header">

                            <span class="panel-label">
                                Parrot
                            </span>

                            <span class="panel-info">
                                Output
                            </span>

                        </div>
                        """
                    )

                    output = gr.Textbox(
                        show_label=False,
                        show_copy_button=True,
                        placeholder="Generated code...",
                        elem_id="output-box"
                    )

        # ----------------------------------------------------
        # HIDDEN SETTINGS
        # ----------------------------------------------------

        max_tokens = gr.Slider(
            minimum=16,
            maximum=512,
            value=192,
            step=8,
            visible=False,
            elem_id="settings"
        )

        temperature = gr.Slider(
            minimum=0.1,
            maximum=1.5,
            value=0.7,
            step=0.05,
            visible=False
        )

        top_p = gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            visible=False
        )

        # ----------------------------------------------------
        # FOOTER
        # ----------------------------------------------------

        gr.HTML(
            """
            <footer id="footer">

                <span>
                    Parrot-Coder · Inserloft Research
                </span>

                <span class="footer-right">
                    Hugging Face Space
                </span>

            </footer>
            """
        )


    # ========================================================
    # EVENTS
    # ========================================================

    generate.click(
        fn=generate_code,
        inputs=[
            prompt,
            max_tokens,
            temperature,
            top_p
        ],
        outputs=output
    )

    prompt.submit(
        fn=generate_code,
        inputs=[
            prompt,
            max_tokens,
            temperature,
            top_p
        ],
        outputs=output
    )


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