import os, uuid, shutil
from flask import Flask, request, jsonify
from PIL import Image
import face_recognition
import requests
from io import BytesIO
from bs4 import BeautifulSoup

TMP_DIR = "temp"
os.makedirs(TMP_DIR, exist_ok=True)

app = Flask(__name__)

def scrape_yandex_reverse(image_bytes):
    session = requests.Session()
    headers = {"User-Agent": "Mozilla/5.0"}

    temp_path = os.path.join(TMP_DIR, f"{uuid.uuid4().hex}.jpg")
    with open(temp_path, "wb") as f:
        f.write(image_bytes)

    files = {"upfile": open(temp_path, "rb")}
    upload_url = "https://yandex.com/images/search"
    r = session.post(upload_url, files=files, headers=headers, allow_redirects=True)
    search_url = r.url

    r2 = session.get(search_url, headers=headers)
    soup = BeautifulSoup(r2.text, "html.parser")
    image_elements = soup.select("img.serp-item__thumb")
    image_links = [img.get("src") for img in image_elements if img.get("src")]

    return image_links[:10]

@app.route('/upload', methods=['POST'])
def upload():
    if 'image' not in request.files:
        return jsonify({"error": "No image uploaded"}), 400

    uploaded = request.files['image']
    image = face_recognition.load_image_file(uploaded)
    try:
        original_enc = face_recognition.face_encodings(image)[0]
    except:
        return jsonify({"error": "No face found"}), 400

    matches = []
    image_bytes = uploaded.read()
    results = scrape_yandex_reverse(image_bytes)

    for link in results:
        try:
            response = requests.get(link, timeout=6)
            img = Image.open(BytesIO(response.content)).convert("RGB")
            path = os.path.join(TMP_DIR, f"{uuid.uuid4().hex}.jpg")
            img.save(path)

            comp = face_recognition.load_image_file(path)
            encs = face_recognition.face_encodings(comp)
            for enc in encs:
                if face_recognition.compare_faces([original_enc], enc, tolerance=0.5)[0]:
                    matches.append({"image_url": link, "source_url": link})
                    break
        except:
            continue

    shutil.rmtree(TMP_DIR)
    os.makedirs(TMP_DIR, exist_ok=True)
    return jsonify({"matches": matches})

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