from flask import Flask, request, jsonify
from flask_cors import CORS
import os
import logging
import torch
import threading
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    AutoModelForCausalLM,
    AutoConfig,
)

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)

# Set environment variables
os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'

# Explainer model globals
explanation_tokenizer = None
explanation_model = None
explanation_is_encoder_decoder = None
model_loaded = False  # Tracks readiness

# Configuration via environment
EXPLAINER_MODEL_NAME = os.environ.get('EXPLAINER_MODEL_NAME', 'HuggingFaceH4/zephyr-7b-beta').strip()
EXPLAINER_TRUST_REMOTE_CODE = os.environ.get('EXPLAINER_TRUST_REMOTE_CODE', '0').strip() in ('1', 'true', 'True')


def load_explainer():
    """Load local explanation model in background."""
    global explanation_tokenizer, explanation_model, explanation_is_encoder_decoder, model_loaded
    try:
        logger.info(f"Loading explainer model: {EXPLAINER_MODEL_NAME} ...")
        config = AutoConfig.from_pretrained(EXPLAINER_MODEL_NAME, trust_remote_code=EXPLAINER_TRUST_REMOTE_CODE)
        explanation_is_encoder_decoder = getattr(config, 'is_encoder_decoder', False)
        explanation_tokenizer = AutoTokenizer.from_pretrained(EXPLAINER_MODEL_NAME, trust_remote_code=EXPLAINER_TRUST_REMOTE_CODE)
        if explanation_is_encoder_decoder:
            explanation_model = AutoModelForSeq2SeqLM.from_pretrained(EXPLAINER_MODEL_NAME, trust_remote_code=EXPLAINER_TRUST_REMOTE_CODE)
        else:
            explanation_model = AutoModelForCausalLM.from_pretrained(EXPLAINER_MODEL_NAME, trust_remote_code=EXPLAINER_TRUST_REMOTE_CODE)
        explanation_model.eval()
        model_loaded = True
        arch = 'seq2seq' if explanation_is_encoder_decoder else 'causal'
        logger.info(f"Explainer model loaded successfully! Detected architecture: {arch}")
    except Exception as e:
        logger.error(f"Error loading explainer model: {e}")
        explanation_tokenizer = None
        explanation_model = None
        explanation_is_encoder_decoder = None


def generate_explanation_local(message: str, label: str) -> str | None:
    """Generate explanation using local model if available."""
    if not model_loaded:
        return None
    try:
        system_instruction = (
            f"Explain in a short and concise why the the text was classified as '{label}'. "
            "If phishing Highlight signals such as urgency, credential requests, suspicious links, sender authenticity, grammar, tone, or context."
        )
        prompt = f"{system_instruction}\n\nMessage:\n{message}\n\nExplanation:"

        if explanation_is_encoder_decoder:
            inputs = explanation_tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512)
            with torch.no_grad():
                output_ids = explanation_model.generate(
                    **inputs,
                    max_new_tokens=180,
                    temperature=0.7,
                    top_p=0.95,
                    do_sample=True,
                    num_return_sequences=1,
                )
            return explanation_tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
        else:
            inputs = explanation_tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=1024)
            with torch.no_grad():
                output_ids = explanation_model.generate(
                    **inputs,
                    max_new_tokens=200,
                    temperature=0.7,
                    top_p=0.95,
                    do_sample=True,
                    pad_token_id=explanation_tokenizer.eos_token_id or explanation_tokenizer.pad_token_id,
                )
            generated = explanation_tokenizer.decode(output_ids[0], skip_special_tokens=True)
            if generated.startswith(prompt):
                generated = generated[len(prompt):]
            return generated.strip()
    except Exception as e:
        logger.error(f"Error generating local explanation: {e}")
        return None


@app.route("/", methods=["GET"])
def home():
    return jsonify({
        "status": "healthy",
        "message": "Anti-Phishing Explanation API",
        "endpoints": {
            "/explain": "POST - Generate explanation for a classification",
            "/health": "GET - Health check"
        }
    })


@app.route("/health", methods=["GET"])
def health():
    return jsonify({
        "status": "healthy",
        "model_ready": model_loaded,
        "architecture": 'seq2seq' if explanation_is_encoder_decoder else ('causal' if model_loaded else None)
    })


@app.route("/explain", methods=["POST"])
def explain():
    if not model_loaded:
        return jsonify({"error": "Model not yet loaded"}), 503
    try:
        data = request.get_json()
        if not data or "message" not in data or "label" not in data:
            return jsonify({"error": "Missing 'message' or 'label' field"}), 400

        message = data["message"]
        label = data["label"]
        if not message.strip() or label not in ("Safe", "Phishing"):
            return jsonify({"error": "Invalid 'message' or 'label'"}), 400

        explanation = generate_explanation_local(message.strip(), label)
        return jsonify({
            "label": label,
            "message": message,
            "explanation": explanation or ""
        })
    except Exception as e:
        logger.error(f"Error in explain endpoint: {e}")
        return jsonify({"error": "Internal server error"}), 500


if __name__ == "__main__":
    threading.Thread(target=load_explainer, daemon=True).start()  # Load model after server starts
    app.run(debug=False, host="0.0.0.0", port=7860)