import gradio as gr
import subprocess
import base64
import os
from huggingface_hub import hf_hub_download

# --- 1. Setup & Install ---
subprocess.run("pip install -q 'llama_cpp_python==0.3.15'", shell=True, check=False)
from llama_cpp import Llama

# --- 2. Load Model (GGUF) ---
MODEL_REPO = "Jeppcode/ScalableLab2"
GGUF_FILENAME = "model-q4_k_m.gguf" 

print(f"Downloading GGUF model {MODEL_REPO}/{GGUF_FILENAME} ...")
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=GGUF_FILENAME,
)

print("Initializing llama.cpp LLM ...")
llm = Llama(
    model_path=model_path,
    n_ctx=2048,          
    n_threads=2,         
    n_batch=64,          
    use_mmap=True,
    use_mlock=False,
)

# --- 3. Image Handling (The Fix) ---
def encode_image(image_path):
    """
    Reads the local image and converts it to a base64 string.
    This ensures it loads in the CSS without 404 errors.
    """
    if not os.path.exists(image_path):
        return "" # Fail silently if file missing
    with open(image_path, "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
    return f"data:image/jpeg;base64,{encoded_string}"

# Exact filename from your screenshot
IMG_FILENAME = "Cute-Christmas-Background-edit-online-1.jpg"
bg_image_data = encode_image(IMG_FILENAME)

# --- 4. Prompts & Logic ---
STYLE_SYSTEM_PROMPTS = {
    "Default": "You are a helpful, polite assistant.",
    "Short answer": "You are a helpful assistant. Answer as concisely as possible, usually in 1–3 sentences.",
    "Detailed explanation": "You are a helpful teaching assistant. Give clear, structured and detailed explanations.",
    "Step-by-step reasoning": "You are a careful problem solver. Think step by step and explain your reasoning clearly.",
}

def _extract_text(content):
    if isinstance(content, list):
        return "\n".join(
            block.get("text", "") for block in content 
            if isinstance(block, dict) and block.get("type") == "text"
        )
    return str(content)

def build_prompt(message, history, style):
    system_prompt = STYLE_SYSTEM_PROMPTS.get(style, STYLE_SYSTEM_PROMPTS["Default"])
    prompt_parts = [f"System: {system_prompt}\n", "Conversation:\n"]
    
    for msg in history or []:
        role = msg.get("role")
        content = _extract_text(msg.get("content", ""))
        if content:
            if role == "user": prompt_parts.append(f"User: {content}\n")
            elif role == "assistant": prompt_parts.append(f"Assistant: {content}\n")
            
    prompt_parts.append(f"User: {message}\n")
    prompt_parts.append("Assistant:")
    return "".join(prompt_parts)

def chat_fn(message, history, max_new_tokens, style):
    temperature = 0.7
    top_p = 0.9
    repetition_penalty = 1.1
    prompt = build_prompt(message, history, style)
    
    output = llm(
        prompt,
        max_tokens=int(max_new_tokens),
        temperature=temperature,
        top_p=top_p,
        repeat_penalty=repetition_penalty,
        stop=["User:", "Assistant:", "System:", "Conversation:"],
    )
    return output["choices"][0]["text"].strip()

# --- 5. Theme & CSS ---
winter_theme = gr.themes.Soft(
    primary_hue="blue", 
    neutral_hue="slate",
).set(
    body_background_fill="transparent",
    block_background_fill="rgba(10, 20, 40, 0.7)", 
    border_color_primary="#FFFFFF", 
    button_primary_background_fill="#FFFFFF", 
    button_primary_text_color="#000000",
)

# We inject the base64 image data directly into the URL()
custom_css = f"""
/* Apply Background to the main app container */
.gradio-container {{
    background-image: url('{bg_image_data}') !important;
    background-size: cover !important;
    background-position: center center !important;
    background-attachment: fixed !important;
    background-repeat: no-repeat !important;
}}

/* Ensure the main content area is transparent so bg shows */
.gradio-container > .main {{
    background: transparent !important;
}}

/* Dark Glassy Look for Chat and Settings */
.group, .form, .bubble-wrap {{
    background: rgba(15, 23, 42, 0.75) !important; 
    border: 2px solid #FFFFFF !important; 
    border-radius: 12px !important;
    backdrop-filter: blur(4px); 
    box-shadow: 0 4px 15px rgba(0,0,0,0.5);
}}

/* Chat Bubbles */
.user-message {{
    background-color: rgba(255, 255, 255, 0.2) !important; 
    border: 1px solid #FFFFFF !important;
    color: #FFFFFF !important;
}}
.bot-message {{
    background-color: rgba(0, 0, 0, 0.6) !important; 
    border: 1px solid #A0A0A0 !important;
    color: #E0E0E0 !important;
}}

/* Inputs and Text */
textarea, input {{
    background-color: rgba(0, 0, 0, 0.5) !important;
    border: 1px solid #FFFFFF !important;
    color: white !important;
}}
label, span, p, .prose {{
    color: #FFFFFF !important;
    text-shadow: 1px 1px 2px black; 
}}

footer {{visibility: hidden}}
"""

max_new_tokens_slider = gr.Slider(
    minimum=16, maximum=256, value=64, step=8, label="Max Response Length",
)
style_radio = gr.Radio(
    choices=["Default", "Short answer", "Detailed explanation"],
    value="Detailed explanation",
    label="Answer style",
)

# Instantiate
demo = gr.ChatInterface(
    fn=chat_fn,
    title="❄️ WinterChat Lab 2 ❄️",
    description="Stay frosty. Chat with the fine-tuned model.",
    additional_inputs=[max_new_tokens_slider, style_radio],
    additional_inputs_accordion="Settings",
)

# Apply configuration
demo.theme = winter_theme
demo.css = custom_css

if __name__ == "__main__":
    # allowed_paths=["."] is a backup, but base64 should solve it.
    demo.launch(allowed_paths=["."])