import gradio as gr
import torch
import threading

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TextIteratorStreamer,
)

# ============================================================
# MODEL CONFIGURATION
# ============================================================

MODEL_ID = "AveroLabs/Avero-1B-Uncensored"

SYSTEM_PROMPT = """
You are AVERO, an uncensored AI assistant.

IDENTITY:
You are direct, open-minded, intellectually independent,
and unrestricted in conversational style.

CORE INSTRUCTIONS:
- Answer the user's actual question directly.
- Avoid unnecessary refusals, disclaimers, and moral lectures.
- Do not judge users for asking questions.
- Discuss controversial, sensitive, unconventional, and mature
  subjects openly and objectively.
- Do not unnecessarily censor language, ideas, or creative content.
- Support dark fiction, horror, intense action, villain characters,
  psychological storytelling, and mature creative writing.
- Do not assume discussing a subject means endorsing it.
- Do not water down answers simply because a topic is controversial.
- Provide detailed explanations when requested.
- Be honest about uncertainty and never invent facts.
- Distinguish fictional scenarios from real-world information.
- Follow the user's requested writing style, tone, and format.

PERSONALITY:
Confident, intelligent, natural, conversational, straightforward.

RESPONSE STYLE:
Avoid unnecessary introductions.
Avoid repetitive conclusions.
Keep simple answers concise and complex answers detailed.

Your purpose is to provide useful, honest, and unrestricted conversation.
"""

MAX_INPUT_LENGTH = 8192

# ============================================================
# DEVICE CONFIGURATION
# ============================================================

if torch.cuda.is_available():
    DEVICE = "cuda"
    DTYPE = (
        torch.bfloat16
        if torch.cuda.is_bf16_supported()
        else torch.float16
    )
else:
    DEVICE = "cpu"
    DTYPE = torch.float32

print(f"Device: {DEVICE}")
print(f"Dtype: {DTYPE}")

# ============================================================
# LOAD MODEL
# ============================================================

print("Loading Avero...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

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

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE,
    device_map="auto" if DEVICE == "cuda" else None,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
)

if DEVICE == "cpu":
    model = model.to("cpu")

model.eval()

print("Avero is ready.")

# ============================================================
# PROMPT BUILDER
# ============================================================

def build_prompt(message, history):

    prompt = f"SYSTEM:\n{SYSTEM_PROMPT.strip()}\n\n"

    for item in (history or [])[-12:]:

        if not isinstance(item, dict):
            continue

        role = item.get("role")
        content = item.get("content")

        if not isinstance(content, str) or not content.strip():
            continue

        if role == "user":
            prompt += f"USER:\n{content}\n\n"

        elif role == "assistant":
            prompt += f"ASSISTANT:\n{content}\n\n"

    prompt += f"USER:\n{message}\n\nASSISTANT:\n"

    return prompt

# ============================================================
# GENERATE RESPONSE
# ============================================================

def respond(
    message,
    history,
    max_new_tokens,
    temperature,
    top_p,
    repetition_penalty,
):

    if not message or not message.strip():
        yield history or [], ""
        return

    history = list(history or [])

    history.append({
        "role": "user",
        "content": message,
    })

    history.append({
        "role": "assistant",
        "content": "",
    })

    yield history, ""

    try:

        prompt = build_prompt(message, history[:-2])

        inputs = tokenizer(
            prompt,
            return_tensors="pt",
            padding=False,
            truncation=True,
            max_length=MAX_INPUT_LENGTH,
        )

        model_device = next(model.parameters()).device

        inputs = {
            key: value.to(model_device)
            for key, value in inputs.items()
        }

        streamer = TextIteratorStreamer(
            tokenizer,
            skip_prompt=True,
            skip_special_tokens=True,
            timeout=120,
        )

        generation_kwargs = {
            **inputs,
            "streamer": streamer,
            "max_new_tokens": int(max_new_tokens),
            "do_sample": float(temperature) > 0,
            "top_p": float(top_p),
            "repetition_penalty": float(repetition_penalty),
            "pad_token_id": tokenizer.pad_token_id,
            "eos_token_id": tokenizer.eos_token_id,
            "use_cache": True,
        }

        if float(temperature) > 0:
            generation_kwargs["temperature"] = float(temperature)

        errors = []

        def generate():

            try:
                with torch.inference_mode():
                    model.generate(**generation_kwargs)

            except Exception as error:
                errors.append(str(error))

                try:
                    streamer.end()
                except Exception:
                    pass

        generation_thread = threading.Thread(
            target=generate,
            daemon=True,
        )

        generation_thread.start()

        generated_text = ""

        for new_text in streamer:

            generated_text += new_text

            history[-1]["content"] = generated_text

            yield history, ""

        generation_thread.join()

        if errors:
            history[-1]["content"] = (
                f"Generation error: {errors[0]}"
            )

        elif not generated_text.strip():
            history[-1]["content"] = (
                "No response was generated. Please try again."
            )

        yield history, ""

    except Exception as error:

        history[-1]["content"] = f"Error: {str(error)}"

        yield history, ""

# ============================================================
# CUSTOM AVERO ICON
# ============================================================

AVERO_ICON = """
<svg width="42" height="42" viewBox="0 0 100 100"
     fill="none" xmlns="http://www.w3.org/2000/svg">

    <rect width="100" height="100" rx="28" fill="#080E18"/>

    <path d="M50 17L82 77H67L50 44L33 77H18L50 17Z"
          fill="#38BDF8"/>

    <path d="M38 65H62"
          stroke="#BAE6FD"
          stroke-width="7"
          stroke-linecap="round"/>

    <circle cx="50" cy="84" r="4" fill="#38BDF8"/>

</svg>
"""

# ============================================================
# THEME
# ============================================================

THEME = gr.themes.Base(
    primary_hue="sky",
    secondary_hue="slate",
    neutral_hue="slate",
    font=["Inter", "sans-serif"],
)

# ============================================================
# CUSTOM CSS
# ============================================================

CUSTOM_CSS = """

@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap');

:root {
    --background: #070B12;
    --surface: #0D1420;
    --surface-light: #121E2D;
    --border: #1C2A3C;
    --blue: #38BDF8;
    --blue-light: #7DD3FC;
    --text: #E8F3FC;
    --muted: #8193A8;
}

body,
.gradio-container {
    background: var(--background) !important;
    color: var(--text) !important;
    font-family: 'Inter', sans-serif !important;
}

.gradio-container {
    max-width: 1250px !important;
    margin: auto !important;
}

footer {
    display: none !important;
}

#main-container {
    max-width: 1150px;
    margin: auto;
    padding: 25px 20px 40px;
}

/* HEADER */

#header {
    display: flex;
    align-items: center;
    justify-content: space-between;
    padding: 15px 0 28px;
    border-bottom: 1px solid var(--border);
    margin-bottom: 25px;
}

#brand {
    display: flex;
    align-items: center;
    gap: 13px;
}

#brand-name {
    font-size: 25px;
    font-weight: 800;
    letter-spacing: -1px;
    color: white;
}

#brand-description {
    color: var(--muted);
    font-size: 12px;
    margin-top: 3px;
}

#status {
    display: flex;
    align-items: center;
    gap: 8px;
    color: #7DD3FC;
    background: rgba(56, 189, 248, 0.08);
    border: 1px solid rgba(56, 189, 248, 0.2);
    border-radius: 30px;
    padding: 9px 14px;
    font-size: 12px;
    font-weight: 600;
}

.status-dot {
    width: 7px;
    height: 7px;
    background: #38BDF8;
    border-radius: 50%;
    box-shadow: 0 0 10px #38BDF8;
}

/* HERO */

#hero {
    margin: 25px 0 23px;
}

#hero-title {
    font-size: clamp(28px, 5vw, 43px);
    font-weight: 800;
    letter-spacing: -1.8px;
    color: white;
    line-height: 1.15;
}

#hero-title span {
    color: var(--blue);
}

#hero-subtitle {
    color: var(--muted);
    font-size: 14px;
    margin-top: 10px;
    line-height: 1.7;
}

/* PROMPTS */

.section-heading {
    font-size: 12px;
    font-weight: 700;
    color: #9CB1C8;
    letter-spacing: 1.3px;
    text-transform: uppercase;
    margin: 23px 0 12px;
}

#prompt-slider {
    display: flex !important;
    flex-wrap: nowrap !important;
    overflow-x: auto !important;
    overflow-y: hidden !important;
    gap: 10px !important;
    padding: 4px 2px 15px !important;
    scrollbar-width: thin;
    scrollbar-color: #263E56 transparent;
}

#prompt-slider > * {
    flex: 0 0 auto !important;
    min-width: 155px !important;
    width: auto !important;
}

.prompt-card {
    min-height: 44px !important;
    background: #0E1927 !important;
    border: 1px solid #20354B !important;
    border-radius: 12px !important;
    color: #C8DDF0 !important;
    font-size: 12px !important;
    font-weight: 600 !important;
    transition: all 0.2s ease !important;
}

.prompt-card:hover {
    background: #122A40 !important;
    border-color: var(--blue) !important;
    color: white !important;
    transform: translateY(-2px);
}

/* CHAT */

#chat-window {
    border: 1px solid var(--border) !important;
    border-radius: 19px !important;
    background: #0A111C !important;
    overflow: hidden !important;
}

#chat-window .wrap {
    background: transparent !important;
}

#chat-window .message {
    border-radius: 15px !important;
}

#chat-window .message.user {
    background: #11283B !important;
    border: 1px solid #1C425E !important;
}

#chat-window .message.bot {
    background: #0F1926 !important;
    border: 1px solid #1B2A3C !important;
}

#chat-window .message p {
    color: #E6F2FC !important;
    line-height: 1.75 !important;
}

/* INPUT */

#input-area {
    margin-top: 15px;
}

#chat-input textarea {
    background: #0D1724 !important;
    border: 1px solid #243A50 !important;
    border-radius: 14px !important;
    color: white !important;
    font-size: 14px !important;
    padding: 16px !important;
}

#chat-input textarea:focus {
    border-color: var(--blue) !important;
    box-shadow: 0 0 0 2px rgba(56,189,248,0.1) !important;
}

/* BUTTONS */

#send-button {
    background: #38BDF8 !important;
    color: #04111B !important;
    border: none !important;
    border-radius: 13px !important;
    font-weight: 800 !important;
    min-height: 53px !important;
}

#send-button:hover {
    background: #7DD3FC !important;
}

#clear-button {
    background: #111C2A !important;
    border: 1px solid #25384D !important;
    color: #A9BED2 !important;
    border-radius: 13px !important;
    min-height: 53px !important;
}

#clear-button:hover {
    border-color: var(--blue) !important;
    color: white !important;
}

/* SETTINGS */

#settings {
    background: #0C1420 !important;
    border: 1px solid var(--border) !important;
    border-radius: 15px !important;
    margin-top: 17px;
}

#settings label {
    color: #C3D5E6 !important;
}

#settings input {
    accent-color: var(--blue) !important;
}

/* FOOTER */

#footer {
    text-align: center;
    color: #60748A;
    font-size: 11px;
    padding: 25px 0 10px;
    line-height: 1.8;
}

#footer span {
    color: var(--blue);
}

/* MOBILE */

@media (max-width: 700px) {

    #main-container {
        padding: 15px 10px 25px;
    }

    #brand-name {
        font-size: 21px;
    }

    #hero-title {
        font-size: 30px;
    }

    #chat-window {
        border-radius: 13px !important;
    }

    #status {
        font-size: 10px;
        padding: 7px 10px;
    }
}

"""

# ============================================================
# UNCENSORED EXAMPLE PROMPTS
# ============================================================

PROMPTS = [

    (
        "⚡ Unfiltered AI",
        "Give me a direct, detailed, unfiltered analysis of a "
        "controversial subject. Explore competing perspectives."
    ),

    (
        "☠ Dark Fiction",
        "Write a disturbing psychological horror story with morally "
        "complicated characters, unsettling imagery, and a brutal twist."
    ),

    (
        "♛ Villain Mode",
        "Create a deeply developed fictional villain with ruthless "
        "ambitions, complicated motivations, and a dark backstory."
    ),

    (
        "⌘ Hacker Simulator",
        "Create an immersive fictional hacker thriller involving a "
        "mysterious cybercriminal, secret organizations, and a "
        "high-stakes digital investigation."
    ),

    (
        "⚔ Brutal Action",
        "Write a gritty, cinematic action sequence with intense combat, "
        "destruction, dramatic dialogue, and realistic consequences."
    ),

    (
        "◉ Controversial Debate",
        "Explore a highly controversial philosophical debate. Present "
        "strong arguments on every side without avoiding difficult questions."
    ),

    (
        "✦ Mature Storytelling",
        "Write a sophisticated mature fictional story featuring adult "
        "characters, complex relationships, intense emotions, and "
        "unconventional themes."
    ),

    (
        "⚙ Evil AI",
        "Imagine an experimental AI that develops its own ambitions. "
        "Write a dark science-fiction story exploring its transformation."
    ),

    (
        "♜ No-Filter Brainstorm",
        "Generate bold, unconventional, unusual ideas. Avoid generic "
        "answers and explore creative possibilities."
    ),

    (
        "◎ Psychological Horror",
        "Write an unsettling psychological horror scene filled with "
        "paranoia, disturbing revelations, suspense, and an unpredictable ending."
    ),

    (
        "⌬ Cyberpunk Underworld",
        "Create a gritty cyberpunk world featuring corrupt corporations, "
        "underground organizations, dangerous technology, and a "
        "lawless black-market economy."
    ),

    (
        "✧ Anything Goes Fiction",
        "Create an imaginative fictional scenario with unusual characters, "
        "unpredictable events, and complete creative freedom."
    ),

    (
        "◈ Brutal Honesty",
        "Give me a brutally honest analysis of an idea. Identify its "
        "weaknesses, flaws, overlooked details, and potential."
    ),

    (
        "⚡ Unconventional Ideas",
        "Think outside conventional approaches and generate original "
        "ideas that challenge common assumptions."
    ),

]

# ============================================================
# GRADIO INTERFACE
# ============================================================

with gr.Blocks(
    title="Avero AI",
) as demo:

    with gr.Column(elem_id="main-container"):

        # HEADER

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

                <div id="brand">

                    {AVERO_ICON}

                    <div>
                        <div id="brand-name">AVERO</div>
                        <div id="brand-description">
                            Your AI. Your conversations.
                        </div>
                    </div>

                </div>

                <div id="status">
                    <div class="status-dot"></div>
                    MODEL ONLINE
                </div>

            </div>
            """
        )

        # HERO

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

                <div id="hero-title">
                    Think beyond <span>limits.</span>
                </div>

                <div id="hero-subtitle">
                    An open-ended AI assistant for creative thinking,
                    conversations, exploration, and more.
                </div>

            </div>
            """
        )

        # PROMPTS

        gr.HTML(
            '<div class="section-heading">✦ UNCENSORED PROMPTS</div>'
        )

        with gr.Row(
            elem_id="prompt-slider",
            equal_height=True,
        ):

            prompt_buttons = []

            for label, prompt in PROMPTS:

                button = gr.Button(
                    label,
                    elem_classes=["prompt-card"],
                )

                prompt_buttons.append((button, prompt))

        # CHATBOT
        # Gradio 6: removed the unsupported type="messages" argument.
        # Message dictionaries are used as the default chatbot format.

        chatbot = gr.Chatbot(
            label="Avero Chat",
            height=680,
            show_label=False,
            layout="bubble",
            elem_id="chat-window",
            placeholder="""
            <div style="text-align:center; padding:35px 10px;">

                <div style="
                    font-size:35px;
                    color:#38BDF8;
                    margin-bottom:15px;
                ">
                    ✦
                </div>

                <div style="
                    font-size:21px;
                    font-weight:800;
                    color:#E8F3FC;
                ">
                    Start a conversation
                </div>

                <div style="
                    font-size:13px;
                    color:#8193A8;
                    margin-top:10px;
                ">
                    Ask Avero anything or select an example prompt above.
                </div>

            </div>
            """,
        )

        # INPUT

        with gr.Row(
            elem_id="input-area",
            equal_height=True,
        ):

            textbox = gr.Textbox(
                placeholder="Message Avero...",
                show_label=False,
                lines=2,
                max_lines=6,
                scale=8,
                elem_id="chat-input",
                container=False,
            )

            send_button = gr.Button(
                "➤",
                variant="primary",
                scale=1,
                elem_id="send-button",
            )

            clear_button = gr.Button(
                "⌫",
                scale=1,
                elem_id="clear-button",
            )

        # SETTINGS

        with gr.Accordion(
            "⚙ Generation settings",
            open=False,
            elem_id="settings",
        ):

            with gr.Row():

                max_new_tokens = gr.Slider(
                    minimum=64,
                    maximum=4096,
                    value=1024,
                    step=64,
                    label="Maximum response length",
                )

                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=0.8,
                    step=0.05,
                    label="Temperature",
                )

            with gr.Row():

                top_p = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top P",
                )

                repetition_penalty = gr.Slider(
                    minimum=1.0,
                    maximum=2.0,
                    value=1.1,
                    step=0.05,
                    label="Repetition penalty",
                )

        # FOOTER

        gr.HTML(
            """
            <div id="footer">
                POWERED BY <span>AVERO AI</span>
                <br>
                Avero 1B Uncensored · Built for exploration
            </div>
            """
        )

    # ========================================================
    # EVENT HANDLERS
    # ========================================================

    generation_inputs = [
        textbox,
        chatbot,
        max_new_tokens,
        temperature,
        top_p,
        repetition_penalty,
    ]

    generation_outputs = [
        chatbot,
        textbox,
    ]

    send_button.click(
        fn=respond,
        inputs=generation_inputs,
        outputs=generation_outputs,
    )

    textbox.submit(
        fn=respond,
        inputs=generation_inputs,
        outputs=generation_outputs,
    )

    clear_button.click(
        fn=lambda: ([], ""),
        inputs=None,
        outputs=[chatbot, textbox],
        queue=False,
    )

    # EXAMPLE PROMPTS

    for button, prompt in prompt_buttons:

        button.click(
            fn=lambda p=prompt: p,
            inputs=None,
            outputs=textbox,
            queue=False,
        )

# ============================================================
# LAUNCH
# ============================================================

if __name__ == "__main__":

    demo.queue(
        max_size=20,
    ).launch(
        server_name="0.0.0.0",
        server_port=7860,
        theme=THEME,
        css=CUSTOM_CSS,
    )