from huggingface_hub import InferenceClient
import gradio as gr
import random

from prompts import GAME_MASTER, COMPRESS_HISTORY, THREE_D_BASE

temperature = 0.92
top_p = 0.95
repetition_penalty = 1.05
MAX_HISTORY = 80

def format_prompt(message, history):
    prompt = "<s>"
    for user, bot in history or []:
        prompt += f"[INST] {user} [/INST] {bot} </s> "
    prompt += f"[INST] {message} [/INST]"
    return prompt

def compress_history(history_str):
    client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
    prompt = COMPRESS_HISTORY.format(history=history_str)
    output = ""
    for chunk in client.text_generation(prompt, temperature=0.7, max_new_tokens=1024, top_p=0.9, stream=True):
        output += getattr(chunk, 'token', chunk).text
    return output

def generate(msg, history, max_tokens=1024, stats=None, world_style="Epic Unreal Engine"):
    client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
    
    stats = stats or "Health: 100\nCreativity: 95\nVision: 88"
    dice = random.randint(1, 10)
    
    gm_text = GAME_MASTER.format(history=history or "New world creation begins...", stats=stats, dice=dice, prompt=msg)
    full_prompt = format_prompt(gm_text, history)
    
    kwargs = {
        "temperature": temperature,
        "max_new_tokens": max_tokens,
        "top_p": top_p,
        "repetition_penalty": repetition_penalty,
        "do_sample": True,
        "seed": random.randint(1, 999999999)
    }
    
    output = ""
    for chunk in client.text_generation(full_prompt, **kwargs, stream=True, details=True, return_full_text=False):
        output += chunk.token.text
        yield [(msg, output)], stats, None
    
    # Append 3D prompts
    three_d = f"\n\n=== 3D GENERATION PROMPTS ===\n{THREE_D_BASE.format(subject='main scene', style=world_style, lighting='volumetric god rays', mood='epic cinematic', camera='wide establishing shot')}\n"
    three_d += f"1. Main World: {output[:250].replace(chr(10), ' ')} ...\n"
    three_d += f"2. Key Asset: Character or object from the scene...\n"
    
    final = output + three_d
    
    if history is None:
        history = []
    history.append((msg, final))
    
    yield history, stats, gr.Dropdown(choices=["Continue", "New Location", "Generate Asset"], value="Continue")

def clear_fn():
    return None, None, None

with gr.Blocks(title="3D World Architect RPG") as demo:
    gr.HTML("<center><h1>🌌 3D World Architect RPG</h1><h3>Mixtral 8x7B • Narrative + Production 3D Prompts</h3></center>")
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=650, label="World Narrative + 3D Prompts", show_copy_button=True)
            msg = gr.Textbox(placeholder="Enter your action or describe a new world...", label="Your Command", lines=3)
            with gr.Row():
                send = gr.Button("Send", variant="primary")
                clear = gr.Button("New World")
        
        with gr.Column(scale=1):
            stats = gr.Textbox(value="Health: 100\nCreativity: 95\nVision: 88", label="Architect Stats", lines=12)
            style = gr.Dropdown(["Epic Unreal Engine", "Photorealistic", "Stylized Fantasy", "Cyberpunk Neon", "Low Poly"], value="Epic Unreal Engine", label="3D Style")
            tokens = gr.Slider(512, 2048, 1024, label="Max Tokens")

    send.click(generate, [msg, chatbot, tokens, stats, style], [chatbot, stats, msg])
    clear.click(clear_fn, outputs=[chatbot, msg, stats])

demo.launch(show_api=False)