import base64
import cv2
import numpy as np
import uvicorn
from fastapi import FastAPI, HTTPException
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 preprocess_image_memory(base64_str):
    try:
        img_data = base64.b64decode(base64_str)
        nparr = np.frombuffer(img_data, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
        
        if img is None:
            raise ValueError("Invalid/corrupt image")

        if len(img.shape) == 3 and 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)
        
        return img
    except Exception as e:
        raise ValueError(f"Error decoding image: {str(e)}")

def extract_icon_positions(img):
    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):
    if not icon_features:
        return None, None
        
    hashes = [img_hash(i) for i in icon_features]
    groups = defaultdict(list)
    
    for i, h in enumerate(hashes):
        found = False
        for label, group in groups.items():
            if np.sum(h != hashes[group[0]]) < 3:
                group.append(i)
                found = True
                break
        if not found:
            groups[len(groups)] = [i]
            
    if not groups:
        return None, None

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

@app.post("/solve")
def solve(data: Input):
    try:
        img = preprocess_image_memory(data.image_base64)

        icons, pos = extract_icon_positions(img)
        
        if not icons:
            return {"error": "No icons found", "x": 0, "y": 0}

        x, y = find_rarest(icons, pos)

        return {"x": int(x), "y": int(y)}
        
    except Exception as e:
        return {"error": str(e)}

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