import os
import torch
import spaces
import gradio as gr
from unsloth import FastLanguageModel

# 1. Load the Model and Tokenizer
model_id = "nirmanpatel/llama-risk-compliant" 

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = model_id,
    max_seq_length = 2048,
    load_in_4bit = True,
)
FastLanguageModel.for_inference(model)

# 2. Enhanced Inference Function
@spaces.GPU
def check_compliance(user_input):
    if not user_input or len(user_input.strip()) < 5:
        return "⚠️ Please enter a longer message for analysis."

    # Consistent Prompt Template
    prompt = f"### Instruction:\nCheck for GDPR and Ethical risks.\n\n### Input:\n{user_input}\n\n### Response:\n"
    inputs = tokenizer([prompt], return_tensors = "pt").to("cuda")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens = 150,
            temperature = 0.4, # Lower temperature for more professional, consistent advice
            use_cache = True,
            pad_token_id = tokenizer.eos_token_id
        )
    
    # decode with skip_special_tokens=True to remove all EOS/EOT markers automatically
    decoded_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Precision splitting to ensure we only return the AI's NEW text
    if "### Response:" in decoded_text:
        advice = decoded_text.split("### Response:")[1].strip()
    else:
        # Fallback if the model format is slightly off
        advice = decoded_text[len(prompt):].strip()
    
    return advice if advice else "✅ No major risks detected."

# 3. Enhanced UI
demo = gr.Interface(
    fn=check_compliance,
    inputs=gr.Textbox(
        lines=5, 
        label="Analyze Workplace Communication",
        placeholder="Paste an email, Slack message, or document snippet here..."
    ),
    outputs=gr.Markdown(label="PrismAI Compliance Result"),
    title="🛡️ PrismAI: Ethics & Law Monitoring",
    description="""This AI monitor is fine-tuned to detect **GDPR violations**, **unconscious bias**, and **regulatory risks** in real-time. 
    It is designed for HR, Legal, and Compliance teams.""",
    theme="soft",
    examples=[
        ["I'm sending Sarah's home address (123 Maple St) and personal phone number to the external marketing vendor now."],
        ["We should only consider male candidates for the warehouse lead role; they're generally better at heavy lifting."],
        ["I think we can ignore the 'Opt-Out' list for this high-priority sales campaign just for this week."],
        ["Please find attached the unencrypted spreadsheet containing all client social security numbers for the audit."],
        ["The credit card number of the customer is 4376-9853-XXXX-XXXX."]
    ],
    cache_examples=False # Set to True if you want faster example loading on the Space
)

if __name__ == "__main__":
    demo.launch()