import os

# Must be set before importing Paddle/PaddleOCR.
os.environ.setdefault("FLAGS_use_mkldnn", "1")
os.environ.setdefault("OMP_NUM_THREADS", "2")
os.environ.setdefault("MKL_NUM_THREADS", "2")

import contextlib
import hashlib
import io
import logging
import re
import tempfile
import threading
import time
from typing import Dict, List, Optional, Tuple

import cv2
import numpy as np
import pypdfium2 as pdfium
import torch
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import hf_hub_download
from PIL import Image, ImageOps
from ultralytics import YOLO

logging.getLogger("ppocr").setLevel(logging.ERROR)
logging.getLogger("paddleocr").setLevel(logging.ERROR)
from paddleocr import PaddleOCR

# ==========================================================
# CONFIG
# ==========================================================
APP_NAME = "2FAKYC PAN Validator"

# Fast endpoint settings
ENABLE_FALLBACK_OCR = True
ENABLE_FULL_IMAGE_FALLBACK = True
REJECT_TALL_SCREENSHOT_BAD_CROP = True

PAN_DETECTION_THRESHOLD = float(os.getenv("PAN_DETECTION_THRESHOLD", "0.50"))
YOLO_IMGSZ = int(os.getenv("YOLO_IMGSZ", "512"))
OCR_MIN_CONFIDENCE = float(os.getenv("OCR_MIN_CONFIDENCE", "0.30"))
MAX_OCR_CORRECTIONS = int(os.getenv("MAX_OCR_CORRECTIONS", "2"))
MAX_UPLOAD_MB = int(os.getenv("MAX_UPLOAD_MB", "12"))
MAX_PDF_PAGES = int(os.getenv("MAX_PDF_PAGES", "1"))

YOLO_DEVICE = "cpu"
torch.set_num_threads(2)
cv2.setNumThreads(2)

PAN_MODEL_REPO = "foduucom/pan-card-detection"
PAN_MODEL_REVISION = "5b6395bcfda0814d8817dc6a446fd70533f88a24"
PAN_MODEL_SHA256 = "a8721936f8585a53227445f997e1ebe10af5ba7faacd3602c01d65514c8dbbc8"

PAN_ENTITY_MAP = {
    "P": "Individual",
    "C": "Company",
    "F": "Firm / LLP",
    "H": "HUF",
    "T": "Trust",
    "A": "Association of Persons",
    "B": "Body of Individuals",
    "G": "Government Agency",
    "L": "Local Authority",
    "J": "Artificial Juridical Person",
}

LETTER_FIX = {
    "0": "O",
    "1": "I",
    "2": "Z",
    "5": "S",
    "6": "G",
    "8": "B",
}

DIGIT_FIX = {
    "O": "0",
    "Q": "0",
    "D": "0",
    "I": "1",
    "L": "1",
    "Z": "2",
    "S": "5",
    "G": "6",
    "B": "8",
}

STRICT_PAN_REGEX = re.compile(r"^[A-Z]{5}[0-9]{4}[A-Z]$")

PAN_DIRECT_PATTERN = re.compile(
    r"(?<![A-Z0-9])([A-Z0-9]{5}\s*-?\s*[0-9OQDILZSGBl]{4}\s*-?\s*[A-Z0-9])(?![A-Z0-9])",
    re.IGNORECASE,
)

BLOCKED_CONTEXT_WORDS = [
    "WHATSAPP",
    "PHONE",
    "MOBILE",
    "CONTACT",
    "EMAIL",
    "MAIL",
    "GMAIL",
    "YAHOO",
    "OUTLOOK",
    "CALL",
    "TEL",
]

PAN_KEYWORDS = [
    "INCOME TAX",
    "INCOME TAX DEPARTMENT",
    "GOVT OF INDIA",
    "GOVERNMENT OF INDIA",
    "PERMANENT ACCOUNT NUMBER",
    "PERMANENT ACCOUNT NUMBER CARD",
    "PERMANENT ACCOUNT",
    "ACCOUNT NUMBER",
    "DATE OF BIRTH",
    "DATE OF INCORPORATION",
    "DATE OF FORMATION",
    "INCORPORATION",
    "FORMATION",
    "SIGNATURE",
]

NEGATIVE_DOCUMENT_KEYWORDS = [
    "AADHAAR",
    "AADHAR",
    "UNIQUE IDENTIFICATION",
    "DRIVING LICENCE",
    "DRIVER LICENSE",
    "PASSPORT",
    "VOTER",
    "ELECTION COMMISSION",
]

MODEL_LOCK = threading.Lock()

# ==========================================================
# FASTAPI
# ==========================================================
app = FastAPI(title=APP_NAME)

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


# ==========================================================
# RESPONSE HELPERS
# ==========================================================
def accepted_response(pan_number: str) -> Dict:
    classification_code = pan_number[3]
    classification_name = PAN_ENTITY_MAP[classification_code]

    return {
        "status": "accepted",
        "message": "Valid PAN card.",
        "data": {
            "pan_number": pan_number,
            "classification_code": classification_code,
            "classification_name": classification_name,
            "masked_pan": mask_pan_for_user_schema(pan_number),
            "kyc_route": "standard",
        },
    }


def rejected_response(message: str) -> Dict:
    return {
        "status": "rejected",
        "message": str(message),
        "data": {},
    }


# ==========================================================
# MODEL LOADING
# ==========================================================
def sha256_file(path: str, chunk_size: int = 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as file:
        while chunk := file.read(chunk_size):
            digest.update(chunk)
    return digest.hexdigest()


@contextlib.contextmanager
def allow_legacy_checkpoint_load():
    original_load = torch.load

    def patched_load(*args, **kwargs):
        kwargs["weights_only"] = False
        return original_load(*args, **kwargs)

    torch.load = patched_load
    try:
        yield
    finally:
        torch.load = original_load


print("Downloading PAN YOLO model...")
pan_model_path = hf_hub_download(
    repo_id=PAN_MODEL_REPO,
    filename="best.pt",
    revision=PAN_MODEL_REVISION,
)

actual_hash = sha256_file(pan_model_path)
if actual_hash != PAN_MODEL_SHA256:
    raise RuntimeError(
        "PAN model hash verification failed. "
        f"Expected {PAN_MODEL_SHA256}, got {actual_hash}"
    )

print("Loading PAN YOLO detector...")
with allow_legacy_checkpoint_load():
    pan_detector = YOLO(pan_model_path)

print("Loading PaddleOCR...")
ocr_reader = PaddleOCR(
    lang="en",
    use_doc_orientation_classify=False,
    use_doc_unwarping=False,
    use_textline_orientation=False,
    engine="paddle",
    device="cpu",
    enable_mkldnn=True,
    cpu_threads=2,
    text_rec_score_thresh=OCR_MIN_CONFIDENCE,
)

print("Warming up models...")
dummy = np.ones((420, 680, 3), dtype=np.uint8) * 255
try:
    _ = pan_detector.predict(
        dummy,
        imgsz=YOLO_IMGSZ,
        conf=0.25,
        device=YOLO_DEVICE,
        verbose=False,
    )
except Exception as error:
    print("YOLO warmup warning:", error)

try:
    _ = ocr_reader.predict(dummy)
except Exception as error:
    print("OCR warmup warning:", error)

print("Models loaded successfully.")


# ==========================================================
# PAN RULES
# ==========================================================
def compact_alnum(text: str) -> str:
    return re.sub(r"[^A-Z0-9]", "", str(text).upper())


def normalize_pan_candidate(raw_candidate: str) -> Optional[str]:
    cleaned = compact_alnum(raw_candidate)

    if len(cleaned) != 10:
        return None

    chars = list(cleaned)
    corrections = 0

    letter_positions = {0, 1, 2, 3, 4, 9}
    digit_positions = {5, 6, 7, 8}

    for index in letter_positions:
        char = chars[index]
        if "A" <= char <= "Z":
            continue
        replacement = LETTER_FIX.get(char)
        if replacement is None:
            return None
        chars[index] = replacement
        corrections += 1

    for index in digit_positions:
        char = chars[index]
        if char.isdigit():
            continue
        replacement = DIGIT_FIX.get(char)
        if replacement is None:
            return None
        chars[index] = replacement
        corrections += 1

    candidate = "".join(chars)

    if corrections > MAX_OCR_CORRECTIONS:
        return None
    if not STRICT_PAN_REGEX.fullmatch(candidate):
        return None
    if candidate[3] not in PAN_ENTITY_MAP:
        return None

    return candidate


def has_blocked_context(text: str) -> bool:
    compact = compact_alnum(text)
    return any(word in compact for word in BLOCKED_CONTEXT_WORDS)


def find_pan_number(ocr_tokens: List[str]) -> Optional[str]:
    """
    Safe PAN extraction.
    Does not slide through long random text like Phone/Whatsapp:93448 35708.
    """
    for token in ocr_tokens:
        token = str(token).strip()
        if not token:
            continue
        if has_blocked_context(token):
            continue

        for match in PAN_DIRECT_PATTERN.finditer(token):
            candidate = normalize_pan_candidate(match.group(1))
            if candidate:
                return candidate

        compact = compact_alnum(token)
        if len(compact) == 10:
            candidate = normalize_pan_candidate(compact)
            if candidate:
                return candidate

    # Limited adjacent-token recovery only for short OCR fragments.
    short_tokens = []
    for token in ocr_tokens:
        token = str(token).strip()
        if has_blocked_context(token):
            continue
        compact = compact_alnum(token)
        if 1 <= len(compact) <= 10:
            short_tokens.append(compact)

    for start in range(len(short_tokens) - 1):
        combined = short_tokens[start] + short_tokens[start + 1]
        if len(combined) == 10:
            candidate = normalize_pan_candidate(combined)
            if candidate:
                return candidate

    return None


def mask_pan_for_user_schema(pan: str) -> str:
    # User required: first 5 masked, next 4 visible, last masked. Example: XXXXX1234X
    return "XXXXX" + pan[5:9] + "X"


# ==========================================================
# PAN DOCUMENT KEYWORDS
# ==========================================================
def normalize_keyword_text(text: str) -> str:
    text = str(text).upper()
    text = re.sub(r"[^A-Z0-9\s]", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def compact_keyword_text(text: str) -> str:
    return re.sub(r"[^A-Z0-9]", "", str(text).upper())


def keyword_matches(full_text: str, keyword: str) -> bool:
    normal_full = normalize_keyword_text(full_text)
    compact_full = compact_keyword_text(full_text)
    normal_keyword = normalize_keyword_text(keyword)
    compact_keyword = compact_keyword_text(keyword)
    return normal_keyword in normal_full or compact_keyword in compact_full


def calculate_pan_keyword_info(ocr_tokens: List[str]) -> Dict:
    full_text = "\n".join([str(t) for t in ocr_tokens])

    matched_pan_keywords = []
    for keyword in PAN_KEYWORDS:
        if keyword_matches(full_text, keyword):
            matched_pan_keywords.append(keyword)

    negative_matches = []
    for keyword in NEGATIVE_DOCUMENT_KEYWORDS:
        if keyword_matches(full_text, keyword):
            negative_matches.append(keyword)

    return {
        "matched_pan_keywords": matched_pan_keywords,
        "negative_document_keywords": negative_matches,
        "pan_keyword_score": len(matched_pan_keywords),
        "negative_keyword_score": len(negative_matches),
    }


# ==========================================================
# OCR + IMAGE HELPERS
# ==========================================================
def extract_ocr_tokens(image_bgr: np.ndarray) -> List[str]:
    tokens = []
    results = ocr_reader.predict(image_bgr)

    for result in results:
        payload = result.json
        if callable(payload):
            payload = payload()

        data = payload.get("res", payload)
        texts = data.get("rec_texts", [])
        scores = data.get("rec_scores", [])

        if len(scores) != len(texts):
            scores = [1.0] * len(texts)

        for text, score in zip(texts, scores):
            text = str(text).strip()
            if text and float(score) >= OCR_MIN_CONFIDENCE:
                tokens.append(text)

    return tokens


def image_bytes_to_bgr(file_bytes: bytes, filename: str) -> List[np.ndarray]:
    ext = os.path.splitext(filename or "")[1].lower()

    if ext == ".pdf":
        images = []
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=True) as tmp:
            tmp.write(file_bytes)
            tmp.flush()
            pdf = pdfium.PdfDocument(tmp.name)
            pages_to_process = min(len(pdf), MAX_PDF_PAGES)
            for page_index in range(pages_to_process):
                page = pdf[page_index]
                bitmap = page.render(scale=2.0)
                pil_image = bitmap.to_pil().convert("RGB")
                img_rgb = np.asarray(pil_image)
                images.append(cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR))
        return images

    try:
        pil_image = Image.open(io.BytesIO(file_bytes))
        pil_image = ImageOps.exif_transpose(pil_image).convert("RGB")
    except Exception as error:
        raise ValueError(f"Could not read uploaded file as image/PDF: {error}")

    img_rgb = np.asarray(pil_image)
    return [cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)]


def resize_input_for_speed(image_bgr: np.ndarray, max_side: int = 1280) -> np.ndarray:
    height, width = image_bgr.shape[:2]
    largest = max(height, width)

    if largest <= max_side:
        return image_bgr

    scale = max_side / largest
    new_width = int(width * scale)
    new_height = int(height * scale)

    return cv2.resize(image_bgr, (new_width, new_height), interpolation=cv2.INTER_AREA)


def crop_with_padding(image_bgr: np.ndarray, xyxy, padding_ratio: float = 0.04) -> np.ndarray:
    height, width = image_bgr.shape[:2]
    x1, y1, x2, y2 = [float(value) for value in xyxy]

    pad_x = (x2 - x1) * padding_ratio
    pad_y = (y2 - y1) * padding_ratio

    x1 = max(0, int(x1 - pad_x))
    y1 = max(0, int(y1 - pad_y))
    x2 = min(width, int(x2 + pad_x))
    y2 = min(height, int(y2 + pad_y))

    crop = image_bgr[y1:y2, x1:x2]
    return crop if crop.size else image_bgr


def resize_for_ocr_fast(
    image_bgr: np.ndarray,
    min_width: int = 750,
    max_width: int = 1100,
    max_side: int = 1100,
    allow_upscale: bool = True,
) -> np.ndarray:
    height, width = image_bgr.shape[:2]
    if width <= 0 or height <= 0:
        return image_bgr

    largest = max(height, width)

    # Reduce very large/tall screenshots first.
    if largest > max_side:
        scale = max_side / largest
        new_width = int(width * scale)
        new_height = int(height * scale)
        return cv2.resize(image_bgr, (new_width, new_height), interpolation=cv2.INTER_AREA)

    # For real PAN crops, small upscaling helps OCR.
    if allow_upscale and width < min_width:
        scale = min_width / width
    elif width > max_width:
        scale = max_width / width
    else:
        return image_bgr

    new_width = int(width * scale)
    new_height = int(height * scale)
    return cv2.resize(image_bgr, (new_width, new_height), interpolation=cv2.INTER_LINEAR)


def enhance_for_ocr_fast(image_bgr: np.ndarray, allow_upscale: bool = True) -> np.ndarray:
    image_bgr = resize_for_ocr_fast(image_bgr, allow_upscale=allow_upscale)
    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    clahe = cv2.createCLAHE(clipLimit=1.8, tileGridSize=(8, 8))
    gray = clahe.apply(gray)
    return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)


def is_tall_mobile_screenshot(image_bgr: np.ndarray) -> bool:
    height, width = image_bgr.shape[:2]
    if width <= 0:
        return False
    return (height / width) >= 1.65


# ==========================================================
# YOLO CROP QUALITY + OCR FALLBACK
# ==========================================================
def get_box_metrics(image_bgr: np.ndarray, box) -> Dict:
    height, width = image_bgr.shape[:2]
    image_area = max(1, height * width)

    x1, y1, x2, y2 = [float(v) for v in box]
    box_w = max(1.0, x2 - x1)
    box_h = max(1.0, y2 - y1)
    box_area = box_w * box_h

    return {
        "area_ratio": round(box_area / image_area, 4),
        "aspect_ratio": round(box_w / box_h, 4),
        "box_width": round(box_w, 2),
        "box_height": round(box_h, 2),
    }


def is_bad_yolo_crop(image_bgr: np.ndarray, box) -> bool:
    metrics = get_box_metrics(image_bgr, box)
    area_ratio = metrics["area_ratio"]
    aspect_ratio = metrics["aspect_ratio"]

    # PAN card should not be a tiny text strip.
    if area_ratio < 0.08:
        return True

    # PAN card is landscape-ish.
    if aspect_ratio < 1.15 or aspect_ratio > 3.20:
        return True

    return False


def run_pan_yolo(image_bgr: np.ndarray) -> Dict:
    start = time.time()
    results = pan_detector.predict(
        image_bgr,
        imgsz=YOLO_IMGSZ,
        conf=0.25,
        device=YOLO_DEVICE,
        verbose=False,
    )
    elapsed = round(time.time() - start, 3)

    boxes = results[0].boxes

    if boxes is None or len(boxes) == 0:
        return {
            "detected": False,
            "confidence": 0.0,
            "box": None,
            "crop": None,
            "bad_crop": True,
            "metrics": None,
            "elapsed_seconds": elapsed,
        }

    best_index = int(torch.argmax(boxes.conf).item())
    best_confidence = float(boxes.conf[best_index].item())
    best_box = boxes.xyxy[best_index].tolist()
    metrics = get_box_metrics(image_bgr, best_box)

    if best_confidence < PAN_DETECTION_THRESHOLD:
        return {
            "detected": False,
            "confidence": round(best_confidence, 4),
            "box": [round(float(x), 2) for x in best_box],
            "crop": None,
            "bad_crop": True,
            "metrics": metrics,
            "elapsed_seconds": elapsed,
        }

    bad_crop = is_bad_yolo_crop(image_bgr, best_box)
    crop = image_bgr if bad_crop else crop_with_padding(image_bgr, best_box)

    return {
        "detected": True,
        "confidence": round(best_confidence, 4),
        "box": [round(float(x), 2) for x in best_box],
        "crop": crop,
        "bad_crop": bad_crop,
        "metrics": metrics,
        "elapsed_seconds": elapsed,
    }


def run_fast_ocr_with_fallback(
    card_bgr: np.ndarray,
    full_image_bgr: Optional[np.ndarray] = None,
    yolo_bad_crop: bool = False,
) -> Dict:
    combined_tokens = []
    seen = set()

    def add_tokens(tokens: List[str]) -> None:
        for token in tokens:
            key = re.sub(r"\s+", " ", str(token).strip().upper())
            if key and key not in seen:
                seen.add(key)
                combined_tokens.append(token)

    # Main OCR image.
    first_image = enhance_for_ocr_fast(card_bgr, allow_upscale=not yolo_bad_crop)
    tokens = extract_ocr_tokens(first_image)
    add_tokens(tokens)
    detected_pan = find_pan_number(combined_tokens)

    # Fallback 1: OCR middle/lower area only if PAN missing and crop was not already full image.
    if ENABLE_FALLBACK_OCR and not detected_pan and not yolo_bad_crop:
        h, w = first_image.shape[:2]
        pan_region = first_image[int(h * 0.25):int(h * 0.85), 0:int(w * 0.95)]
        if pan_region.size:
            tokens = extract_ocr_tokens(pan_region)
            add_tokens(tokens)
            detected_pan = find_pan_number(combined_tokens)

    # Fallback 2: full-image OCR only when crop was normal but OCR still failed.
    if (
        ENABLE_FULL_IMAGE_FALLBACK
        and full_image_bgr is not None
        and not detected_pan
        and not yolo_bad_crop
    ):
        full_image = enhance_for_ocr_fast(full_image_bgr, allow_upscale=False)
        tokens = extract_ocr_tokens(full_image)
        add_tokens(tokens)
        detected_pan = find_pan_number(combined_tokens)

    return {
        "tokens": combined_tokens,
        "detected_pan": detected_pan,
    }


# ==========================================================
# FINAL ANALYZER
# ==========================================================
def analyze_pan_image(image_bgr: np.ndarray) -> Dict:
    image_bgr = resize_input_for_speed(image_bgr, max_side=1280)
    tall_screenshot = is_tall_mobile_screenshot(image_bgr)

    # Gate 1: YOLO PAN detection.
    yolo_result = run_pan_yolo(image_bgr)

    if not yolo_result["detected"]:
        return rejected_response("Not a PAN card")

    # Fast false-positive protection for WhatsApp/mobile screenshots.
    if REJECT_TALL_SCREENSHOT_BAD_CROP and tall_screenshot and yolo_result["bad_crop"]:
        return rejected_response("Not a PAN card")

    # Gate 2: OCR.
    ocr_result = run_fast_ocr_with_fallback(
        card_bgr=yolo_result["crop"],
        full_image_bgr=image_bgr,
        yolo_bad_crop=yolo_result["bad_crop"],
    )

    ocr_tokens = ocr_result["tokens"]
    detected_pan = ocr_result["detected_pan"]
    keyword_info = calculate_pan_keyword_info(ocr_tokens)

    if not ocr_tokens:
        return rejected_response("Image is blurry or text is not readable")

    if not detected_pan:
        return rejected_response("PAN number not readable")

    # Reject wrong document types.
    if keyword_info["negative_keyword_score"] > 0 and keyword_info["pan_keyword_score"] == 0:
        return rejected_response("Not a PAN card")

    # A PAN-format string alone is not enough.
    if keyword_info["pan_keyword_score"] < 1:
        return rejected_response("Not a PAN card")

    return accepted_response(detected_pan)


# ==========================================================
# ROUTES
# ==========================================================
@app.get("/")
def root():
    return {"ok": True, "service": APP_NAME, "endpoint": "/predict"}


@app.get("/health")
def health():
    return {
        "ok": True,
        "models_loaded": True,
        "yolo_imgsz": YOLO_IMGSZ,
        "pan_detection_threshold": PAN_DETECTION_THRESHOLD,
    }


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        file_bytes = await file.read()

        if not file_bytes:
            return rejected_response("No file uploaded")

        if len(file_bytes) > MAX_UPLOAD_MB * 1024 * 1024:
            return rejected_response(f"File too large. Max allowed size is {MAX_UPLOAD_MB} MB")

        images_bgr = image_bytes_to_bgr(file_bytes, file.filename or "upload")

        if not images_bgr:
            return rejected_response("Could not read image")

        # Keep serialized inference on HF free CPU to avoid CPU overload.
        with MODEL_LOCK:
            # PDF support: return the first accepted page, otherwise final rejection.
            last_rejection = rejected_response("Not a PAN card")
            for image_bgr in images_bgr:
                result = analyze_pan_image(image_bgr)
                if result.get("status") == "accepted":
                    return result
                last_rejection = result

            return last_rejection

    except Exception as error:
        return rejected_response(f"Processing error: {type(error).__name__}: {error}")