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="Prohorizon Advanced Masking")

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

# Initialize OCR once
reader = easyocr.Reader(['en', 'hi'], gpu=False)

def enhance_mobile_photo(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    enhanced = clahe.apply(gray)
    return enhanced

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

def apply_mask(img, bbox, mask_ratio):
    """Applies a black rectangle over a percentage of the detected text box."""
    p1 = tuple(map(int, bbox[0])) # Top-left
    p3 = tuple(map(int, bbox[2])) # Bottom-right
    
    width = p3[0] - p1[0]
    mask_width = int(width * mask_ratio)
    
    # Draw the mask (0,0,0) is black
    cv2.rectangle(img, p1, (p1[0] + mask_width, p3[1]), (0, 0, 0), -1)
    return img

@app.post("/v1/aadhaar/process")
async def process_document(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        nparr = np.frombuffer(contents, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        # 1. Resize for Speed & Accuracy
        h, w = img.shape[:2]
        img = cv2.resize(img, (1200, int(h * (1200/w))))
        
        # 2. Enhance image for OCR pass
        processed_img = enhance_mobile_photo(img)

        # 3. OCR Pass
        results = reader.readtext(processed_img, detail=1)
        
        extracted = {"aadhaar": None, "vid": None, "name": None}

        # 4. Pattern Detection & Masking
        for i, (bbox, text, conf) in enumerate(results):
            clean = get_clean_digits(text)
            text_upper = text.upper()
            
            # --- 1. Aadhaar Masking (12 Digits) ---
            if 11 <= len(clean) <= 12:
                extracted["aadhaar"] = clean
                img = apply_mask(img, bbox, 0.68) # Masks 8 digits, leaves 4
                continue # Move to next box once masked

            # --- 2. VID Masking (16 Digits) ---
            # Scenario A: VID label and numbers together or standalone 16 digits
            if len(clean) >= 15 or ( "VID" in text_upper and len(clean) >= 8):
                extracted["vid"] = clean
                img = apply_mask(img, bbox, 0.75) # Masks 12 digits, leaves 4
                continue

            # Scenario B: "VID" label in current box, numbers in the NEXT box
            if "VID" in text_upper and i + 1 < len(results):
                next_bbox, next_text, _ = results[i+1]
                next_clean = get_clean_digits(next_text)
                if len(next_clean) >= 8:
                    extracted["vid"] = next_clean
                    img = apply_mask(img, next_bbox, 0.75)
                    continue

            # --- 3. Name Extraction ---
            if not extracted["name"] and conf > 0.7:
                if not re.search(r'\d', text) and len(text.split()) >= 2:
                    if not any(k in text_upper for k in ["GOVT", "INDIA", "MALE", "FEMALE", "DOB", "YEAR"]):
                        extracted["name"] = text.strip()

        # 5. Final Image Response
        _, buffer = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
        
        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:
        raise HTTPException(status_code=500, detail=str(e))

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