import io
import os
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
from PIL import Image
from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
import torch
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# ============================================================
#  Config
# ============================================================
TWILIO_ACCOUNT_SID      = os.getenv("TWILIO_ACCOUNT_SID", "ACebd0e2daf1f7060a0901b9e1766052de")
TWILIO_AUTH_TOKEN       = os.getenv("TWILIO_AUTH_TOKEN", "98d8c5ff61765ba0e37d89748e78f991")
TWILIO_FROM_NUMBER      = os.getenv("TWILIO_FROM_NUMBER", "+17125825991")
YOUR_PHONE_NUMBER       = os.getenv("YOUR_PHONE_NUMBER", "+919047432845")
AI_CONFIDENCE_THRESHOLD = float(os.getenv("AI_CONFIDENCE_THRESHOLD", 0.5))

# ── Local model path (Swin-Base, umm-maybe/AI-image-detector weights) ──
MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "local_model")

UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), "uploads")
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

app = Flask(__name__, static_folder=".", template_folder=".")
CORS(app)

# ============================================================
#  Load Model from local disk
# ============================================================
print("=" * 55)
print("  VisionAI — Loading model from local disk")
print(f"  Path : {MODEL_DIR}")
print("=" * 55)

try:
    required = ["pytorch_model.bin", "config.json", "preprocessor_config.json"]
    missing  = [f for f in required if not os.path.exists(os.path.join(MODEL_DIR, f))]

    if missing:
        raise FileNotFoundError(
            f"Missing files in local_model/: {missing}\n"
            "Run:  python setup_local_model.py"
        )

    processor = AutoImageProcessor.from_pretrained(MODEL_DIR, local_files_only=True)
    model     = AutoModelForImageClassification.from_pretrained(MODEL_DIR, local_files_only=True)
    model.eval()

    # Wrap in HF pipeline for clean API
    pipe = pipeline(
        "image-classification",
        model=model,
        image_processor=processor,
        device=0 if torch.cuda.is_available() else -1,
    )

    device_name = "GPU (CUDA)" if torch.cuda.is_available() else "CPU"
    print(f"\n  ✅ Model loaded successfully on {device_name}")
    print(f"  Labels : {list(model.config.id2label.values())}")
    print(f"  Params : {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")
    print("=" * 55 + "\n")
    MODEL_LOADED = True

except Exception as e:
    print(f"\n  ⚠  Could not load local model: {e}")
    print("  Falling back to HuggingFace hub (requires internet)...")
    try:
        pipe = pipeline("image-classification", model="umm-maybe/AI-image-detector")
        MODEL_LOADED = True
        print("  ✅ Fallback model loaded from HuggingFace hub.\n")
    except Exception as e2:
        print(f"  ❌ Both local and hub load failed: {e2}")
        pipe = None
        MODEL_LOADED = False

# ============================================================
#  Routes
# ============================================================
@app.route("/")
def index():
    return send_from_directory(".", "index.html")

@app.route("/<path:filename>")
def static_files(filename):
    return send_from_directory(".", filename)

@app.route("/status")
def status():
    """Health check — frontend polls this on load."""
    return jsonify({
        "model_loaded":  MODEL_LOADED,
        "model_dir":     MODEL_DIR,
        "device":        "GPU" if torch.cuda.is_available() else "CPU",
        "threshold":     AI_CONFIDENCE_THRESHOLD * 100,
        "alert_phone":   YOUR_PHONE_NUMBER,
    })

@app.route("/analyze", methods=["POST"])
def analyze():
    if not MODEL_LOADED or pipe is None:
        return jsonify({"error": "Model not loaded — run setup_local_model.py first."}), 503

    if "image" not in request.files:
        return jsonify({"error": "No image file in request."}), 400

    file = request.files["image"]
    if file.filename == "":
        return jsonify({"error": "Empty filename."}), 400

    # ── Run inference ──────────────────────────────────────
    img_bytes = file.read()
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")

    results   = pipe(img)
    scores    = {r["label"].lower(): r["score"] for r in results}
    art_score = scores.get("artificial", 0.0)
    real_score= scores.get("real",       1.0 - art_score)
    is_ai     = art_score > AI_CONFIDENCE_THRESHOLD

    # ── Twilio voice call if AI detected ──────────────────
    call_placed = False
    call_sid    = None
    call_error  = None

    if is_ai:
        try:
            from twilio.rest import Client
            client = Client(TWILIO_ACCOUNT_SID, TWILIO_AUTH_TOKEN)
            call = client.calls.create(
                twiml=f"""<Response>
                    <Say voice="alice">
                        Alert! An AI generated image has been detected.
                        File: {file.filename}.
                        Confidence: {round(art_score * 100)} percent.
                        Please review immediately.
                    </Say>
                </Response>""",
                to=YOUR_PHONE_NUMBER,
                from_=TWILIO_FROM_NUMBER,
            )
            call_placed = True
            call_sid    = call.sid
            print(f"  📞 Call placed → SID: {call_sid}")
        except Exception as e:
            call_error = str(e)
            print(f"  ⚠  Twilio call failed: {e}")

    print(f"  [{file.filename}]  artificial={art_score*100:.1f}%  is_ai={is_ai}  call={call_placed}")

    return jsonify({
        "filename":         file.filename,
        "is_ai":            is_ai,
        "artificial_score": round(art_score  * 100, 1),
        "real_score":       round(real_score * 100, 1),
        "all_scores":       [{"label": r["label"], "score": round(r["score"] * 100, 1)} for r in results],
        "threshold":        AI_CONFIDENCE_THRESHOLD * 100,
        "call_placed":      call_placed,
        "call_sid":         call_sid,
        "call_error":       call_error,
        "alert_phone":      YOUR_PHONE_NUMBER,
    })

# ============================================================
if __name__ == "__main__":
    app.run(debug=False, host="0.0.0.0", port=5000)