import os
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

import spaces

# Dummy function to satisfy Hugging Face ZeroGPU startup check
@spaces.GPU
def dummy_gpu_fn():
    pass


# 1. Download the GGUF model file from Hugging Face Hub
# Adjust repo_id and filename if you named them differently when pushing to Hub
REPO_ID = "iamfebin/wish-lawyer-gguf"
FILENAME = "Meta-Llama-3.1-8B-Instruct.Q4_K_M.gguf"

print(f"Downloading GGUF model '{FILENAME}' from repository '{REPO_ID}'...")
try:
    model_path = hf_hub_download(
        repo_id=REPO_ID,
        filename=FILENAME
    )
    print(f"Model downloaded successfully to: {model_path}")
except Exception as e:
    print(f"Error downloading model: {e}")
    print("If you pushed your model under a different filename or repository, please update REPO_ID and FILENAME in app.py.")
    model_path = None

# 2. Initialize the Llama model for CPU inference
llm = None
if model_path:
    print("Loading model into memory (using llama.cpp with 2 threads)...")
    try:
        llm = Llama(
            model_path=model_path,
            n_ctx=1024,      # Reduced from 2048 to lower memory footprint
            n_threads=2,     # Free Hugging Face Spaces provide 2 vCPUs
            use_mmap=False   # Read directly into RAM to prevent memory-mapped paging OOMs
        )
        print("Model loaded successfully!")
    except Exception as e:
        import traceback
        print("Failed to load model!")
        traceback.print_exc()

# 3. Prompt Template
wish_prompt_template = """### System:
You are an expert Wish Lawyer. Your job is to analyze dangerous human wishes, identify at least 3 hidden loopholes or catastrophic risks, and rewrite the wish into a single, legally ironclad sentence that protects the wisher completely.

### Context/Grantor:
{}

### Human Wish:
{}

### Risk Analysis (Internal Thought Process):
"""

def consult_wish_lawyer(wish, context):
    if not llm:
        return "Error: Model not loaded. Please check model download logs.", ""
        
    if not wish.strip():
        return "Please input a wish!", ""
        
    formatted_prompt = wish_prompt_template.format(context, wish)
    
    # Generate output using Llama.cpp CPU inference
    response = llm(
        formatted_prompt,
        max_tokens=512,
        temperature=0.5,
        top_p=0.9,
        stop=["### System:", "### Human Wish:", "### Context/Grantor:"]
    )
    
    generated_text = response["choices"][0]["text"].strip()
    
    # Split risk analysis and rewritten wish
    split_marker = "### Safer Rewritten Wish:"
    if split_marker not in generated_text:
        split_marker = "### Ironclad Rewritten Wish:"
        
    if split_marker in generated_text:
        parts = generated_text.split(split_marker)
        risk_analysis = parts[0].strip()
        rewritten_wish = parts[1].strip()
    else:
        risk_analysis = generated_text
        rewritten_wish = "No rewritten wish generated."
        
    return risk_analysis, rewritten_wish

# 4. Gradio UI Layout
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # ⚖️ Wish Lawyer: Supernatural Legal Counsel
        Have you ever been worried about the literal interpretations of genies, crossroads devils, or trickster monkey paws?
        This is an instruction-tuned **Llama 3.1 8B** model trained to audit supernatural wishes for loopholes and redraft them into ironclad contracts.
        """
    )
    
    with gr.Row():
        with gr.Column():
            wish_input = gr.Textbox(
                label="Your Human Wish", 
                placeholder="e.g., I want to never feel tired.",
                lines=3
            )
            context_input = gr.Dropdown(
                label="Grantor / Context",
                choices=[
                    "A trickster monkey paw",
                    "An erratic genie",
                    "A deal with a crossroads devil",
                    "A suspicious corporate contract",
                    "Standard Wish Rules Apply"
                ],
                value="Standard Wish Rules Apply"
            )
            submit_btn = gr.Button("Consult the Wish Lawyer", variant="primary")
            
        with gr.Column():
            risk_output = gr.Textbox(
                label="Risk Analysis (Loopholes Identified)", 
                lines=6
            )
            wish_output = gr.Textbox(
                label="Safer Rewritten Wish (Ironclad)", 
                lines=4
            )
            
    submit_btn.click(
        fn=consult_wish_lawyer,
        inputs=[wish_input, context_input],
        outputs=[risk_output, wish_output]
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="amber", secondary_hue="slate"))