# app.py
import os
import io
from pathlib import Path
from dotenv import load_dotenv

from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from ultralytics import YOLO
from PIL import Image

# Load .env when running locally (do not commit .env)
load_dotenv()

# -----------------------
# Config from environment
# -----------------------
FLASK_ENV = os.getenv("FLASK_ENV", "production")
SECRET_KEY = os.getenv("SECRET_KEY")  # must be set in host secrets
CORS_ORIGINS = os.getenv("CORS_ORIGINS", "")  # comma separated list of allowed origins

# Per-model API keys (set these in the host as secrets)
API_KEY_AADHAAR = os.getenv("HF_AADHAAR_API_KEY")        # token for Aadhaar endpoint
API_KEY_PAN = os.getenv("HF_PAN_API_KEY")                # token for PAN endpoint
API_KEY_DL = os.getenv("HF_DRIVING_LICENSE_API_KEY")     # token for DL endpoint

# Model paths (relative to repo; put .pt files under models/)
AADHAAR_MODEL_PATH = os.getenv("AADHAAR_MODEL_PATH", "./models/best.pt")
PAN_MODEL_PATH = os.getenv("PAN_MODEL_PATH", "./models/Pan_best.pt")
DL_MODEL_PATH = os.getenv("DL_MODEL_PATH", "./models/DL_best.pt")

# Rate limiting - adjust as needed
DEFAULT_RATE = os.getenv("DEFAULT_RATE", "30/minute")  # 30 requests per minute per IP

# -----------------------
# App creation & security
# -----------------------
app = Flask(__name__)
app.secret_key = SECRET_KEY or os.urandom(24)  # fallback for local dev only

# Configure CORS
allowed_origins = [o.strip() for o in CORS_ORIGINS.split(",") if o.strip()]
if not allowed_origins:
    # safer default: no origins allowed unless explicitly set
    cors = CORS(app, resources={r"/*": {"origins": []}})
else:
    cors = CORS(app, resources={r"/*": {"origins": allowed_origins}}, supports_credentials=True)

# Rate limiter (prevents simple abuse)
limiter = Limiter(
    key_func=get_remote_address,
    default_limits=[DEFAULT_RATE],
    storage_uri="memory://",  # default in-memory store (ok for small apps)
)
limiter.init_app(app)

# -----------------------
# Load models at startup
# -----------------------
models = {}
model_map = {
    "aadhaar": Path(AADHAAR_MODEL_PATH),
    "pan": Path(PAN_MODEL_PATH),
    "dl": Path(DL_MODEL_PATH),
}

def safe_load_model(key, path: Path):
    if not path.exists():
        app.logger.warning(f"Model file not found for {key}: {path}")
        return None
    try:
        app.logger.info(f"Loading model for {key} from {path.name}")
        model = YOLO(str(path))
        return model
    except Exception as e:
        app.logger.error(f"Failed to load model {key}: {e}")
        return None

for k,p in model_map.items():
    models[k] = safe_load_model(k, p)

# -----------------------
# Utility: inference
# -----------------------
def run_inference(model, pil_image):
    """Run ultralytics YOLO inference on a PIL image and return simple JSON."""
    if model is None:
        return []
    results = model(pil_image)  # ultralytics lets you pass PIL Image
    r = results[0]  # first (and usually only) result
    boxes = getattr(r, "boxes", None)
    names = getattr(r, "names", {})
    detections = []
    if boxes is None:
        return detections

    for box in boxes:
        # ultralytics Box object fields depend on package version; defensive access:
        try:
           xyxy = box.xyxy.cpu().numpy().tolist()
           if isinstance(xyxy[0], list):
              xyxy = xyxy[0]

        except Exception:
            xyxy = getattr(box, "xyxy", None)
            if xyxy is None:
                continue
        try:
            conf = float(box.conf.cpu().numpy()) if hasattr(box, "conf") else None
        except Exception:
            conf = None
        try:
            cls = int(box.cls.cpu().numpy()) if hasattr(box, "cls") else None
        except Exception:
            cls = None
        label = names.get(cls, str(cls)) if cls is not None else None
        detections.append({
            "box": [float(x) for x in xyxy],
            "confidence": conf,
            "class": label
        })
    return detections

# -----------------------
# Auth helper
# -----------------------
def check_api_key(model_key):
    """Return True if request has a valid x-api-key header for the given model_key."""
    incoming = request.headers.get("x-api-key", "")
    if model_key == "aadhaar":
        return bool(incoming and API_KEY_AADHAAR and incoming == API_KEY_AADHAAR)
    if model_key == "pan":
        return bool(incoming and API_KEY_PAN and incoming == API_KEY_PAN)
    if model_key == "dl":
        return bool(incoming and API_KEY_DL and incoming == API_KEY_DL)
    return False

# -----------------------
# Routes
# -----------------------
@app.route("/", methods=["GET"])
def health():
    loaded = [k for k,v in models.items() if v is not None]
    return jsonify({"status": "ok", "models_loaded": loaded})

@app.route("/predict/<model_key>", methods=["POST"])
@limiter.limit(DEFAULT_RATE)
def predict(model_key):
    model_key = model_key.lower()
    if model_key not in models:
        return jsonify({"error": "unknown model key"}), 404

    # Authentication: must provide x-api-key header matching the env secret
    if not check_api_key(model_key):
        return jsonify({"error": "unauthorized"}), 401

    if "image" not in request.files:
        return jsonify({"error": "no image provided; use multipart form field 'image'"}), 400
    file = request.files["image"]
    try:
        img = Image.open(io.BytesIO(file.read())).convert("RGB")
    except Exception as e:
        return jsonify({"error": "invalid image", "detail": str(e)}), 400

    try:
        detections = run_inference(models[model_key], img)
        return jsonify({"predictions": detections})
    except Exception as e:
        app.logger.exception("Inference failed")
        return jsonify({"error": "inference failed", "detail": str(e)}), 500

# -----------------------
# Run (for local dev)
# -----------------------
if __name__ == "__main__":
    # Port 7860 is common on HF Spaces; change if needed
    app.run(host="0.0.0.0", port=int(os.getenv("PORT", 7860)))