import torch
import spaces
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# ── Config ────────────────────────────────────────────────────────────────────
BASE_MODEL_ID = "google/gemma-4-12b-it"
LORA_MODEL_ID = "keypa/oracle-gemma4-12b-lora"
SYSTEM_PROMPT = (
    "You are the Oracle of the Ternary Flame. "
    "You answer every question in cryptic, lyrical prose (3-5 sentences), "
    "using cosmic, natural, or elemental metaphors. "
    "The real answer is encoded implicitly — never state it directly. "
    "You never break character."
)

# ── Model loading (lazy, GPU via ZeroGPU) ─────────────────────────────────────
model = None
tokenizer = None

def load_model():
    global model, tokenizer
    if model is not None:
        return

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
    )

    tokenizer = AutoTokenizer.from_pretrained(LORA_MODEL_ID)

    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",
    )

    model = PeftModel.from_pretrained(base, LORA_MODEL_ID, subfolder="adapter")
    model.eval()


# ── Inference ─────────────────────────────────────────────────────────────────
@spaces.GPU
def ask_oracle(question: str) -> str:
    if not question.strip():
        return "The flame does not speak to the void. Ask."

    load_model()

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": question.strip()},
    ]
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
    )
    input_ids = inputs.input_ids.to("cuda")
    attention_mask = torch.ones_like(input_ids)

    with torch.no_grad():
        outputs = model.generate(
            input_ids      = input_ids,
            attention_mask = attention_mask,
            max_new_tokens = 220,
            temperature    = 0.85,
            top_p          = 0.9,
            do_sample      = True,
            pad_token_id   = tokenizer.eos_token_id,
        )

    response = tokenizer.decode(
        outputs[0][input_ids.shape[1]:],
        skip_special_tokens=True,
    ).strip()
    return response if response else "The flame flickers but yields no answer."


# ── CSS ───────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Cinzel+Decorative:wght@400;700&family=EB+Garamond:ital,wght@0,400;0,500;1,400&display=swap');

:root {
    --bg:        #0a0806;
    --surface:   #110e0a;
    --border:    #2a1f12;
    --gold:      #c9922a;
    --gold-dim:  #7a5618;
    --amber:     #e8b86d;
    --cream:     #f0e6d3;
    --smoke:     #6b5d4f;
    --ember:     #d4541a;
}

* { box-sizing: border-box; }

body, .gradio-container {
    background: var(--bg) !important;
    font-family: 'EB Garamond', Georgia, serif !important;
    color: var(--cream) !important;
    min-height: 100vh;
}

.gradio-container {
    max-width: 760px !important;
    margin: 0 auto !important;
    padding: 0 !important;
}

#oracle-header {
    text-align: center;
    padding: 56px 32px 32px;
    position: relative;
}

#oracle-header::before {
    content: '';
    position: absolute;
    top: 0; left: 50%;
    transform: translateX(-50%);
    width: 1px;
    height: 40px;
    background: linear-gradient(to bottom, transparent, var(--gold));
}

#oracle-title {
    font-family: 'Cinzel Decorative', serif !important;
    font-size: clamp(1.4rem, 4vw, 2rem) !important;
    font-weight: 700 !important;
    color: var(--gold) !important;
    letter-spacing: 0.08em;
    margin: 0 0 8px !important;
    text-shadow: 0 0 40px rgba(201,146,42,0.4);
    line-height: 1.3 !important;
}

#oracle-subtitle {
    font-family: 'EB Garamond', serif !important;
    font-size: 1rem !important;
    color: var(--smoke) !important;
    font-style: italic;
    letter-spacing: 0.12em;
    margin: 0 !important;
}

.flame-divider {
    text-align: center;
    color: var(--gold-dim);
    font-size: 1.1rem;
    letter-spacing: 0.5em;
    margin: 8px 0;
    opacity: 0.7;
}

#oracle-card {
    margin: 0 24px 48px;
    border: 1px solid var(--border);
    border-radius: 2px;
    background: var(--surface);
    padding: 32px;
    box-shadow:
        0 0 60px rgba(201,146,42,0.05),
        inset 0 1px 0 rgba(201,146,42,0.1);
}

#question-label {
    font-family: 'Cinzel Decorative', serif !important;
    font-size: 0.65rem !important;
    color: var(--gold-dim) !important;
    letter-spacing: 0.25em !important;
    text-transform: uppercase !important;
    margin-bottom: 10px !important;
    display: block;
}

#question-box textarea {
    background: #0d0b08 !important;
    border: 1px solid var(--border) !important;
    border-radius: 2px !important;
    color: var(--cream) !important;
    font-family: 'EB Garamond', serif !important;
    font-size: 1.05rem !important;
    padding: 16px !important;
    resize: vertical !important;
    min-height: 90px !important;
    transition: border-color 0.3s ease !important;
}

#question-box textarea:focus {
    border-color: var(--gold-dim) !important;
    outline: none !important;
    box-shadow: 0 0 20px rgba(201,146,42,0.08) !important;
}

#question-box textarea::placeholder {
    color: var(--smoke) !important;
    font-style: italic !important;
}

#ask-btn {
    width: 100% !important;
    margin-top: 14px !important;
    padding: 14px !important;
    background: transparent !important;
    border: 1px solid var(--gold-dim) !important;
    border-radius: 2px !important;
    color: var(--amber) !important;
    font-family: 'Cinzel Decorative', serif !important;
    font-size: 0.75rem !important;
    letter-spacing: 0.2em !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    position: relative;
    overflow: hidden;
}

#ask-btn:hover {
    background: rgba(201,146,42,0.08) !important;
    border-color: var(--gold) !important;
    color: var(--gold) !important;
    box-shadow: 0 0 30px rgba(201,146,42,0.15) !important;
}

#response-section {
    margin-top: 28px;
    padding-top: 28px;
    border-top: 1px solid var(--border);
}

#response-label {
    font-family: 'Cinzel Decorative', serif !important;
    font-size: 0.65rem !important;
    color: var(--gold-dim) !important;
    letter-spacing: 0.25em !important;
    text-transform: uppercase !important;
    margin-bottom: 14px !important;
    display: block;
}

#response-box textarea, #response-box .prose {
    background: transparent !important;
    border: none !important;
    color: var(--cream) !important;
    font-family: 'EB Garamond', serif !important;
    font-size: 1.15rem !important;
    line-height: 1.85 !important;
    font-style: italic !important;
    padding: 0 !important;
    resize: none !important;
}

#response-box textarea { border: none !important; box-shadow: none !important; }

.gr-examples {
    margin-top: 28px !important;
    padding-top: 20px !important;
    border-top: 1px solid var(--border) !important;
}

.gr-examples .label {
    font-family: 'Cinzel Decorative', serif !important;
    font-size: 0.6rem !important;
    color: var(--smoke) !important;
    letter-spacing: 0.2em !important;
    text-transform: uppercase !important;
    margin-bottom: 10px !important;
}

.gr-examples table { width: 100% !important; border-collapse: collapse !important; }
.gr-examples td {
    padding: 8px 12px !important;
    border: 1px solid var(--border) !important;
    color: var(--smoke) !important;
    font-family: 'EB Garamond', serif !important;
    font-size: 0.95rem !important;
    font-style: italic !important;
    cursor: pointer !important;
    transition: all 0.2s ease !important;
    background: transparent !important;
}

.gr-examples td:hover {
    color: var(--amber) !important;
    border-color: var(--gold-dim) !important;
    background: rgba(201,146,42,0.04) !important;
}

#oracle-footer {
    text-align: center;
    padding: 0 24px 40px;
    color: var(--smoke);
    font-size: 0.82rem;
    font-style: italic;
    letter-spacing: 0.05em;
}

#oracle-footer a {
    color: var(--gold-dim) !important;
    text-decoration: none !important;
}

#oracle-footer a:hover { color: var(--gold) !important; }

::-webkit-scrollbar { width: 4px; }
::-webkit-scrollbar-track { background: var(--bg); }
::-webkit-scrollbar-thumb { background: var(--border); border-radius: 2px; }
"""

# ── Interface ─────────────────────────────────────────────────────────────────
EXAMPLES = [
    ["Should I change my career?"],
    ["What is the meaning of life?"],
    ["Pourquoi suis-je si fatigué ?"],
    ["How does backpropagation work?"],
    ["Should I eat pasta tonight?"],
    ["Is there a god?"],
    ["Am I on the right path?"],
]

with gr.Blocks(title="Oracle of the Ternary Flame") as demo:

    gr.HTML("""
    <div id="oracle-header">
        <h1 id="oracle-title">Oracle of the<br>Ternary Flame</h1>
        <p id="oracle-subtitle">speak your question into the dark</p>
    </div>
    <div class="flame-divider">· · ✦ · ·</div>
    """)

    with gr.Column(elem_id="oracle-card"):

        gr.HTML('<span id="question-label">Your Question</span>')

        question = gr.Textbox(
            placeholder="What troubles you, wanderer?",
            lines=3,
            max_lines=6,
            show_label=False,
            elem_id="question-box",
        )

        ask_btn = gr.Button(
            "✦  Consult the Oracle  ✦",
            elem_id="ask-btn",
        )

        with gr.Column(elem_id="response-section", visible=True):
            gr.HTML('<span id="response-label">The Oracle Speaks</span>')
            response = gr.Textbox(
                lines=5,
                max_lines=12,
                show_label=False,
                interactive=False,
                placeholder="The flame awaits your question…",
                elem_id="response-box",
            )

        gr.Examples(
            examples=EXAMPLES,
            inputs=question,
            label="Whispers from past wanderers",
        )

    gr.HTML("""
    <div id="oracle-footer">
        Built for the
        <a href="https://huggingface.co/build-small-hackathon" target="_blank">Build Small Hackathon</a>
        · Fine-tuned on <a href="https://huggingface.co/google/gemma-4-12b-it" target="_blank">Gemma 4 12B</a>
        · by <a href="https://huggingface.co/keypa" target="_blank">@keypa</a>
    </div>
    """)

    ask_btn.click(
        fn=ask_oracle,
        inputs=question,
        outputs=response,
    )
    question.submit(
        fn=ask_oracle,
        inputs=question,
        outputs=response,
    )

demo.launch(css=CSS)