from flask import Flask, render_template, request, Response, jsonify
import cv2
import os
import numpy as np
import pickle
from datetime import datetime

# --- Flask App ---
app = Flask(__name__)

FACE_DATA_DIR = 'face_data'
FACE_CASCADE_PATH = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'

if not os.path.exists(FACE_DATA_DIR):
    os.makedirs(FACE_DATA_DIR)

face_cascade = cv2.CascadeClassifier(FACE_CASCADE_PATH)

camera = None
face_recognizer = cv2.face.LBPHFaceRecognizer_create()
is_trained = False
has_webcam = os.path.exists("/dev/video0")  # deteksi webcam di server

def load_face_data():
    global is_trained
    faces, labels, names = [], [], []

    if os.path.exists(os.path.join(FACE_DATA_DIR, 'names.pkl')):
        with open(os.path.join(FACE_DATA_DIR, 'names.pkl'), 'rb') as f:
            names = pickle.load(f)

    for idx, name in enumerate(names):
        face_dir = os.path.join(FACE_DATA_DIR, name)
        if os.path.exists(face_dir):
            for filename in os.listdir(face_dir):
                if filename.endswith('.jpg'):
                    img_path = os.path.join(face_dir, filename)
                    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                    faces.append(img)
                    labels.append(idx)

    if faces:
        face_recognizer.train(faces, np.array(labels))
        is_trained = True
        return names
    return []

def get_camera():
    global camera
    if has_webcam:
        if camera is None:
            camera = cv2.VideoCapture(0)
        return camera
    return None

@app.route('/')
def index():
    names = load_face_data()
    return render_template('index.html', registered_faces=names, has_webcam=has_webcam)

@app.route('/register')
def register():
    return render_template('register.html', has_webcam=has_webcam)

@app.route('/recognize')
def recognize():
    return render_template('recognize.html', has_webcam=has_webcam)

@app.route('/video_feed')
def video_feed():
    if not has_webcam:
        return "Webcam tidak tersedia di server ini", 404

    def generate():
        cam = get_camera()
        while True:
            success, frame = cam.read()
            if not success:
                break
            ret, buffer = cv2.imencode('.jpg', frame)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')

    return Response(generate(), mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/recognition_feed')
def recognition_feed():
    if not has_webcam:
        return "Webcam tidak tersedia di server ini", 404

    def generate():
        cam = get_camera()
        names = load_face_data()

        while True:
            success, frame = cam.read()
            if not success:
                break

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = face_cascade.detectMultiScale(gray, 1.3, 5)

            for (x, y, w, h) in faces:
                cv2.rectangle(frame, (x, y), (x+w, y+h), (255, 0, 0), 2)

                if is_trained and names:
                    roi_gray = gray[y:y+h, x:x+w]
                    roi_gray = cv2.resize(roi_gray, (100, 100))

                    id_, confidence = face_recognizer.predict(roi_gray)
                    if confidence < 100:
                        name = names[id_]
                        confidence_text = f"{name} ({round(100-confidence)}%)"
                    else:
                        confidence_text = "Unknown"

                    cv2.putText(frame, confidence_text, (x, y-10),
                              cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)

            ret, buffer = cv2.imencode('.jpg', frame)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')

    return Response(generate(), mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/capture_face', methods=['POST'])
def capture_face():
    if has_webcam:
        name = request.json.get('name', '').strip()
        if not name:
            return jsonify({'error': 'Nama tidak boleh kosong'})

        cam = get_camera()
        success, frame = cam.read()
        if not success:
            return jsonify({'error': 'Gagal mengambil gambar dari kamera'})

        return save_face(name, frame)
    else:
        return jsonify({'error': 'Webcam tidak tersedia, gunakan /upload_face'})

@app.route('/upload_face', methods=['POST'])
def upload_face():
    name = request.form.get('name', '').strip()
    file = request.files.get('file')

    if not name:
        return jsonify({'error': 'Nama tidak boleh kosong'})
    if not file:
        return jsonify({'error': 'File tidak ditemukan'})

    np_img = np.frombuffer(file.read(), np.uint8)
    frame = cv2.imdecode(np_img, cv2.IMREAD_COLOR)
    return save_face(name, frame)

def save_face(name, frame):
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)

    if len(faces) == 0:
        return jsonify({'error': 'Tidak ada wajah yang terdeteksi'})
    if len(faces) > 1:
        return jsonify({'error': 'Terdeteksi lebih dari satu wajah'})

    (x, y, w, h) = faces[0]
    face_roi = gray[y:y+h, x:x+w]
    face_roi = cv2.resize(face_roi, (100, 100))

    person_dir = os.path.join(FACE_DATA_DIR, name)
    if not os.path.exists(person_dir):
        os.makedirs(person_dir)

    filename = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
    cv2.imwrite(os.path.join(person_dir, filename), face_roi)

    names_file = os.path.join(FACE_DATA_DIR, 'names.pkl')
    if os.path.exists(names_file):
        with open(names_file, 'rb') as f:
            names = pickle.load(f)
    else:
        names = []

    if name not in names:
        names.append(name)
        with open(names_file, 'wb') as f:
            pickle.dump(names, f)

    load_face_data()
    return jsonify({'success': f'Wajah {name} berhasil didaftarkan'})

# --- Wrapper untuk Hugging Face ---
from fastapi import FastAPI
from starlette.middleware.wsgi import WSGIMiddleware

flask_app = app
asgi_app = FastAPI()
asgi_app.mount("/", WSGIMiddleware(flask_app))

if __name__ == '__main__':
    flask_app.run(debug=True, host='0.0.0.0', port=5000)