import json
import os
import re
import secrets

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

APP_ENV = os.getenv("APP_ENV", "dev").lower()
APP_PORT = int(os.getenv("PORT", "7860"))
MAX_UPLOAD_MB = int(os.getenv("MAX_UPLOAD_MB", "10"))
MAX_UPLOAD_BYTES = MAX_UPLOAD_MB * 1024 * 1024
API_KEY = os.getenv("MASKING_API_KEY", "")
EXPOSE_DOCS = os.getenv("EXPOSE_DOCS", "true" if APP_ENV != "prod" else "false").lower() == "true"
ALLOW_ORIGINS = [origin.strip() for origin in os.getenv("ALLOWED_ORIGINS", "").split(",") if origin.strip()]

app = FastAPI(
    title="Aadhaar Strict Masking API",
    docs_url="/docs" if EXPOSE_DOCS else None,
    redoc_url="/redoc" if EXPOSE_DOCS else None,
    openapi_url="/openapi.json" if EXPOSE_DOCS else None,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOW_ORIGINS or ["*"],
    allow_credentials=False,
    allow_methods=["POST", "GET"],
    allow_headers=["*"],
)

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


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


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)


def mask_numeric_value(value, visible_last=4, mask_char="x"):
    if not value or len(value) <= visible_last:
        return value
    return f'{mask_char * (len(value) - visible_last)}{value[-visible_last:]}'


def apply_mask(img, bbox, ratio):
    p1 = tuple(map(int, bbox[0]))
    p3 = tuple(map(int, bbox[2]))
    width = max(1, 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


def validate_api_key(x_api_key):
    if not API_KEY:
        return
    if not x_api_key or not secrets.compare_digest(x_api_key, API_KEY):
        raise HTTPException(status_code=401, detail="Unauthorized")


@app.get("/health")
async def health_check():
    return {"status": "ok", "env": APP_ENV}


@app.post("/v1/aadhaar/process")
async def process_document(
    file: UploadFile = File(...),
    x_api_key: str | None = Header(default=None, alias="x-api-key"),
):
    try:
        validate_api_key(x_api_key)

        contents = await file.read()
        if not contents:
            raise HTTPException(status_code=400, detail="Empty upload")
        if len(contents) > MAX_UPLOAD_BYTES:
            raise HTTPException(status_code=413, detail=f"File too large. Max allowed is {MAX_UPLOAD_MB} MB")

        img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
        if img is None:
            raise HTTPException(status_code=400, detail="Invalid image format")

        h, w = img.shape[:2]
        img = cv2.resize(img, (1200, int(h * (1200 / w))))
        processed_img = enhance_mobile_photo(img)
        results = reader.readtext(processed_img, detail=1)

        aadhaar_values = []
        vid_values = []
        masked_aadhaar_keys = set()
        masked_vid_keys = set()

        for i, (bbox, text, conf) in enumerate(results):
            if conf < 0.35:
                continue

            clean = get_clean_digits(text)
            text_upper = text.upper()

            if len(clean) == 12 and not clean.startswith(("0", "1")):
                key = (clean, tuple(map(int, bbox[0])), tuple(map(int, bbox[2])))
                if key not in masked_aadhaar_keys:
                    aadhaar_values.append(clean)
                    img = apply_mask(img, bbox, 0.68)
                    masked_aadhaar_keys.add(key)
                continue

            if "VID" in text_upper and len(clean) >= 12:
                key = (clean[:16], tuple(map(int, bbox[0])), tuple(map(int, bbox[2])))
                if key not in masked_vid_keys:
                    vid_values.append(clean[:16])
                    img = apply_mask(img, bbox, 0.79)
                    masked_vid_keys.add(key)

            elif "VID" in text_upper and i + 1 < len(results):
                next_bbox, next_text, next_conf = results[i + 1]
                if next_conf < 0.35:
                    continue

                next_clean = get_clean_digits(next_text)
                if len(next_clean) >= 12:
                    key = (next_clean[:16], tuple(map(int, next_bbox[0])), tuple(map(int, next_bbox[2])))
                    if key not in masked_vid_keys:
                        vid_values.append(next_clean[:16])
                        img = apply_mask(img, next_bbox, 0.79)
                        masked_vid_keys.add(key)

            elif len(clean) == 16:
                key = (clean, tuple(map(int, bbox[0])), tuple(map(int, bbox[2])))
                if key not in masked_vid_keys:
                    vid_values.append(clean)
                    img = apply_mask(img, bbox, 0.79)
                    masked_vid_keys.add(key)

        extracted = {
            "aadhaar": [mask_numeric_value(value, visible_last=4) for value in aadhaar_values] or None,
            "vid": [mask_numeric_value(value, visible_last=4) for value in vid_values] or None,
        }

        if isinstance(extracted["aadhaar"], list) and len(extracted["aadhaar"]) == 1:
            extracted["aadhaar"] = extracted["aadhaar"][0]
        if isinstance(extracted["vid"], list) and len(extracted["vid"]) == 1:
            extracted["vid"] = extracted["vid"][0]

        _, 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",
                "Cache-Control": "no-store",
                "Pragma": "no-cache",
            },
        )

    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        print("Error while processing document")
        raise HTTPException(status_code=500, detail="Internal server error")


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