import base64
import cv2
import numpy as np
import os
from fastapi import FastAPI
from pydantic import BaseModel
from collections import defaultdict

app = FastAPI()

@app.get("/")
def root():
    return {
        "status": "ok",
        "service": "iconCaptcha solver",
        "endpoint": "/solve"
    }

class Input(BaseModel):
    image_base64: str

def save_base64_image_cv(base64_str, output_path="final.png"):
    img_data = base64.b64decode(base64_str)
    nparr = np.frombuffer(img_data, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)

    if img.shape[2] == 4:
        alpha = img[:, :, 3] / 255.0
        rgb = img[:, :, :3]
        white_bg = np.ones_like(rgb, dtype=np.uint8) * 255
        img = (rgb * alpha[:, :, None] + white_bg * (1 - alpha[:, :, None])).astype(np.uint8)

    cv2.imwrite(output_path, img)

def extract_icon_positions(image_path):
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    icons, pos = [], []
    for c in contours:
        x, y, w, h = cv2.boundingRect(c)
        if w > 10 and h > 10:
            roi = cv2.resize(thresh[y:y+h, x:x+w], (50, 50))
            icons.append(roi)
            pos.append((x, y))
    return icons, pos

def img_hash(img):
    img = cv2.resize(img, (8, 8))
    return (img > img.mean()).astype(np.uint8).flatten()

def find_rarest(icon_features, positions):
    hashes = [img_hash(i) for i in icon_features]
    groups = defaultdict(list)

    for i, h in enumerate(hashes):
        for g in groups.values():
            if np.sum(h != hashes[g[0]]) < 3:
                g.append(i)
                break
        else:
            groups[len(groups)] = [i]

    idx = min(groups.values(), key=len)[0]
    return positions[idx]

@app.post("/solve")
def solve(data: Input):
    try:
        save_base64_image_cv(data.image_base64)
        icons, pos = extract_icon_positions("final.png")
        x, y = find_rarest(icons, pos)
        return {"x": x, "y": y}
    finally:
        if os.path.exists("final.png"):
            os.remove("final.png")