import os
os.environ['TRANSFORMERS_CACHE'] = '/app/.cache'  # ✅ Set local writable cache

from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_sock import Sock
import uuid
import time
import requests
from transformers import pipeline

# Initialize Flask app and WebSocket
app = Flask(__name__)
CORS(app)
sock = Sock(app)

# AI classification pipeline using Facebook's zero-shot model (lightweight for free-tier)
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Active chat sessions: {session_id: [ {role, msg, ip, timestamp} ]}
SESSIONS = {}

# Labels we want AI to detect (can be expanded)
SENSITIVE_LABELS = ["terrorism", "blackmail", "national security threat"]

# Storage space API to log flagged chats
STORAGE_API = "https://mike23415-storage.hf.space/api/flag"

def flag_if_sensitive(text, ip, session_id, role):
    # AI checks if the message matches sensitive labels
    result = classifier(text, SENSITIVE_LABELS)
    scores = dict(zip(result["labels"], result["scores"]))

    # Threshold to consider a message flagged
    for label, score in scores.items():
        if score > 0.8:
            print(f"⚠️ FLAGGED: {label} with score {score}")
            payload = {
                "ip": ip,
                "session_id": session_id,
                "timestamp": time.time(),
                "role": role,
                "message": text,
                "label": label,
                "score": score
            }
            # Send to storage space for later review
            try:
                requests.post(STORAGE_API, json=payload, timeout=3)
            except Exception as e:
                print("Failed to send flag:", e)
            break

@sock.route('/ws/<session_id>')
def chat(ws, session_id):
    ip = request.remote_addr or "unknown"
    if session_id not in SESSIONS:
        SESSIONS[session_id] = []

    join_index = sum(1 for msg in SESSIONS[session_id] if msg["role"].startswith("Receiver")) + 1
    role = "Sender" if len(SESSIONS[session_id]) == 0 else f"Receiver {join_index}"

    try:
        while True:
            msg = ws.receive()
            if msg is None:
                break
            entry = {
                "role": role,
                "msg": msg,
                "ip": ip,
                "timestamp": time.time()
            }
            SESSIONS[session_id].append(entry)

            # AI flagging in background
            flag_if_sensitive(msg, ip, session_id, role)

            # Broadcast to all participants
            for conn in SESSIONS[session_id]:
                try:
                    ws.send(f"{role}: {msg}")
                except:
                    continue
    except:
        pass

@app.route("/")
def root():
    return "Real-time AI chat backend is running."

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)