import os
import re
import cv2
import json
import easyocr
import uvicorn
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI(title="Aadhaar Strict Masking API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

reader = easyocr.Reader(['en'], gpu=False)

# -------------------------------
# Enhance Image
# -------------------------------
def enhance_mobile_photo(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)).apply(gray)

# -------------------------------
# Clean Digits
# -------------------------------
def get_clean_digits(text):
    text = text.upper().replace('O','0').replace('I','1').replace('L','1').replace('B', '8')
    return re.sub(r'[^0-9]', '', text)

# -------------------------------
# Mask Function
# -------------------------------
def apply_mask(img, bbox, ratio):
    p1 = tuple(map(int, bbox[0]))
    p3 = tuple(map(int, bbox[2]))

    width = p3[0] - p1[0]
    mask_width = int(width * ratio)

    cv2.rectangle(img, p1, (p1[0] + mask_width, p3[1]), (0,0,0), -1)
    return img

# -------------------------------
# -------------------------------
# QR Mask
# -------------------------------
# def mask_qr_code(img):
#     detector = cv2.QRCodeDetector()
#     data, bbox, _ = detector.detectAndDecode(img)
#
#     if bbox is not None:
#         bbox = bbox[0].astype(int)
#         x1, y1 = bbox[0]
#         x2, y2 = bbox[2]
#         cv2.rectangle(img, (x1, y1), (x2, y2), (0,0,0), -1)
#
    #return img

# -------------------------------
# MAIN API
# -------------------------------
@app.post("/v1/aadhaar/process")
async def process_document(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)

        # ✅ QR Mask (before resize)
        #img = mask_qr_code(img)

         #Resize
        h, w = img.shape[:2]
        img = cv2.resize(img, (1200, int(h * (1200/w))))

        # ✅ QR Mask again (better detection)
        #img = mask_qr_code(img)

        processed_img = enhance_mobile_photo(img)
        
        results = reader.readtext(processed_img, detail=1)

        extracted = {"aadhaar": None, "vid": None}

        for i, (bbox, text, conf) in enumerate(results):
            clean = get_clean_digits(text)
            text_upper = text.upper()

            # Aadhaar (12 digit only)
            if len(clean) == 12 and not clean.startswith(('0','1')):
                extracted["aadhaar"] = clean
                img = apply_mask(img, bbox, 0.68)
                continue

            # VID: 16 digits or text context
            if "VID" in text_upper and len(clean) >= 12:
                extracted["vid"] = clean
                img = apply_mask(img, bbox, 0.79) # Mask first 12 digits
                
            # Scenario B: "VID" is in one box, numbers are in the NEXT box
            elif "VID" in text_upper:
                if i + 1 < len(results):
                    next_bbox, next_text, _ = results[i+1]
                    next_clean = get_clean_digits(next_text)
                    if len(next_clean) >= 12:
                        extracted["vid"] = next_clean
                        img = apply_mask(img, next_bbox, 0.79)

            # Scenario C: 16 digits found standalone
            elif len(clean) == 16:
                extracted["vid"] = clean
                img = apply_mask(img, bbox, 0.79)

        _, buffer = cv2.imencode('.jpg', img)
        return Response(
            content=buffer.tobytes(),
            media_type="image/jpeg",
            headers={
                "x-data": json.dumps(extracted),
                "Access-Control-Expose-Headers": "x-data"
            }
        )

    except Exception as e:
        print(f"Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

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