main.py

import os
import re
import json
import tempfile
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from pdf2image import convert_from_path
from fastapi import FastAPI, UploadFile, Form, HTTPException, Header, Depends
from fastapi.responses import JSONResponse
import pytesseract
from PIL import Image
import requests

# -----------------------------
# CONFIG
# -----------------------------
IMG_SIZE = (224, 224)
MODEL_PATH = "./final_model.keras"
class_names = ['Other', 'Aadhaar Card', 'Pan Card', 'Voter Id']

type_mapping = {
    "aadhaar_card": "Aadhaar Card",
    "pan_card": "Pan Card",
    "voter_id": "Voter Id"
}

EXTRACTABLE_TYPES = {"Aadhaar Card", "Pan Card", "Voter Id"}

# -----------------------------
# AUTH CONFIG (HF SECRETS)
# -----------------------------
MASTER_SECRET_KEY = os.getenv("MASTER_SECRET_KEY")
SARVAM_API_KEY = os.getenv("SARVAM_API_KEY")

PROJECT_KEYS_JSON = os.getenv("PROJECT_KEYS_JSON", "{}")
try:
    PROJECT_KEYS = json.loads(PROJECT_KEYS_JSON)
except Exception:
    PROJECT_KEYS = {}

# -----------------------------
# AUTH VALIDATION
# -----------------------------
def verify_keys(
    x_secret_key: str = Header(None),
    x_project_id: str = Header(None),
    x_project_key: str = Header(None),
):
    if not MASTER_SECRET_KEY:
        raise HTTPException(status_code=500, detail="MASTER_SECRET_KEY not configured in Space secrets")

    if not x_secret_key or x_secret_key != MASTER_SECRET_KEY:
        raise HTTPException(status_code=401, detail="Invalid Secret Key")

    if not x_project_id or not x_project_key:
        raise HTTPException(status_code=401, detail="Project Id and Project Key are required")

    if x_project_id not in PROJECT_KEYS:
        raise HTTPException(status_code=401, detail=f"Unknown project: {x_project_id}")

    if PROJECT_KEYS.get(x_project_id) != x_project_key:
        raise HTTPException(status_code=401, detail="Invalid Project Key")

    return True

# -----------------------------
# LOAD MODEL
# -----------------------------
model = tf.keras.models.load_model(MODEL_PATH)

# -----------------------------
# Predict a single image array (H,W,C)
# -----------------------------
def predict_array(img_array):
    img_array = tf.cast(img_array, tf.float32)
    img_array = tf.image.resize(img_array, IMG_SIZE)
    img_array = tf.expand_dims(img_array, axis=0)

    pred = model.predict(img_array, verbose=0)[0]
    class_id = int(np.argmax(pred))
    conf = float(np.max(pred))

    return class_names[class_id], conf, pred

# -----------------------------
# Predict from image path
# -----------------------------
def predict_image(img_path):
    img = image.load_img(img_path, target_size=IMG_SIZE, color_mode="rgb")
    img_array = image.img_to_array(img)
    label, conf, _ = predict_array(img_array)
    return {
        "type": label,
        "confidence": round(conf, 6)
    }

# -----------------------------
# Predict from PDF (max 2 pages only)
# -----------------------------
def predict_pdf(pdf_path, dpi=200, max_pages=2):
    pages = convert_from_path(pdf_path, dpi=dpi)

    if len(pages) > max_pages:
        return {
            "error": "Maximum page limit reached",
            "max_pages": max_pages,
            "found_pages": len(pages)
        }

    results = []
    labels = []

    for page in pages:
        page_np = np.array(page)

        if page_np.shape[-1] == 4:
            page_np = page_np[..., :3]

        label, conf, _ = predict_array(page_np)
        results.append((label, conf))
        labels.append(label)

    if len(labels) == 2:
        if labels[0] != labels[1]:
            return {
                "error": "Invalid input",
                "reason": "Pages belong to different document types",
                "page1_type": labels[0],
                "page2_type": labels[1]
            }

        final_label = labels[0]
        final_conf = float(max(results[0][1], results[1][1]))

        return {
            "type": final_label,
            "confidence": round(final_conf, 6)
        }

    return {
        "type": labels[0],
        "confidence": round(float(results[0][1]), 6)
    }

# -----------------------------
# Main predict dispatcher
# -----------------------------
def predict_file(path):
    if not os.path.exists(path):
        return {"error": "File not found", "path": path}

    ext = os.path.splitext(path)[1].lower()

    if ext in [".png", ".jpg", ".jpeg", ".bmp", ".webp"]:
        return predict_image(path)

    if ext == ".pdf":
        return predict_pdf(path)

    return {"error": "Unsupported file type", "ext": ext}

# -----------------------------
# OCR: Extract raw text from file
# -----------------------------
def extract_text_from_file(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    all_text = []

    try:
        if ext in [".png", ".jpg", ".jpeg", ".bmp", ".webp"]:
            pil_img = Image.open(path).convert("RGB")
            text = pytesseract.image_to_string(pil_img, lang="eng+hin")
            all_text.append(text)

        elif ext == ".pdf":
            pages = convert_from_path(path, dpi=250)
            for page in pages:
                text = pytesseract.image_to_string(page, lang="eng+hin")
                all_text.append(text)

    except Exception as e:
        return f"OCR_ERROR: {str(e)}"

    return "\n".join(all_text).strip()

# -----------------------------
# Sarvam AI: Extract structured fields from raw OCR text
# -----------------------------
def extract_fields_with_sarvam(raw_text: str, doc_type: str) -> dict:
    empty_fields = {
        "name": None,
        "address": None,
        "date_of_birth": None,
        "gender": None,
        "id_number": None,
        "mobile_number": None,
    }

    if not SARVAM_API_KEY:
        return {**empty_fields, "ocr_error": "SARVAM_API_KEY not configured in Space secrets"}

    system_prompt = (
        "You are a document field extractor for Indian identity documents. "
        "Given raw OCR text, extract specific fields and return the result wrapped in <output> tags like this:\n\n"
        "<output>\n"
        "{\n"
        '  "name": "...",\n'
        '  "address": "...",\n'
        '  "date_of_birth": "...",\n'
        '  "gender": "...",\n'
        '  "id_number": "...",\n'
        '  "mobile_number": "..."\n'
        "}\n"
        "</output>\n\n"
        "Fields to extract:\n"
        "  - name: Full name of the document holder\n"
        "  - address: Full address (combine multi-line if needed)\n"
        "  - date_of_birth: In DD/MM/YYYY or YYYY format\n"
        "  - gender: Male / Female / Transgender\n"
        "  - id_number: Aadhaar=12-digit number, PAN=10-char alphanumeric, Voter ID=alphanumeric card number\n"
        "  - mobile_number: 10-digit mobile number if present\n\n"
        "Use null for any field not found. Only output the <output> block, nothing else outside it."
    )

    user_prompt = (
        f"Document Type: {doc_type}\n\n"
        f"Raw OCR Text:\n{raw_text}\n\n"
        "Extract the fields and return inside <output> tags."
    )

    try:
        response = requests.post(
            "https://api.sarvam.ai/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {SARVAM_API_KEY}",
                "Content-Type": "application/json"
            },
            json={
                "model": "sarvam-m",
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                "temperature": 0.0,
                "max_tokens": 512
            },
            timeout=30
        )

        if response.status_code != 200:
            return {
                **empty_fields,
                "sarvam_error": f"API returned {response.status_code}: {response.text[:300]}"
            }

        content = response.json()["choices"][0]["message"]["content"].strip()

        # Extract JSON from <output>...</output> tag
        output_match = re.search(r"<output>(.*?)</output>", content, flags=re.DOTALL)
        if not output_match:
            return {
                **empty_fields,
                "sarvam_error": f"No <output> tag found in response: {content[:300]}"
            }

        json_str = output_match.group(1).strip()
        extracted = json.loads(json_str)

        # Normalize keys
        fields = ["name", "address", "date_of_birth", "gender", "id_number", "mobile_number"]
        normalized = {f: extracted.get(f) for f in fields}

        # Handle year_of_birth fallback
        if normalized["date_of_birth"] is None and "year_of_birth" in extracted:
            normalized["date_of_birth"] = str(extracted["year_of_birth"])

        return normalized

    except json.JSONDecodeError as e:
        return {**empty_fields, "sarvam_parse_error": f"JSON parse failed: {str(e)}"}
    except requests.exceptions.Timeout:
        return {**empty_fields, "sarvam_error": "Sarvam API request timed out"}
    except Exception as e:
        return {**empty_fields, "sarvam_error": str(e)}

# -----------------------------
# FastAPI App
# -----------------------------
app = FastAPI(
    title="ID Document Validator",
    description="Classify document type, authenticate, and extract structured fields via OCR + Sarvam AI."
)

@app.post("/predict")
async def predict(
    file: UploadFile,
    expected_type: str = Form(None),
    auth: bool = Depends(verify_keys)
):
    if not file.filename:
        raise HTTPException(status_code=400, detail="No file provided")

    suffix = os.path.splitext(file.filename)[1]

    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
        temp_path = temp_file.name
        temp_file.write(await file.read())

    try:
        # Step 1: Classify document type
        result = predict_file(temp_path)

        # Step 2: Authenticate against expected_type
        if expected_type and expected_type in type_mapping and "type" in result:
            if result["type"] == type_mapping[expected_type]:
                result["Authentication"] = "Valid"
            else:
                result["Authentication"] = "Not valid"

        # Step 3: OCR + Sarvam field extraction
        doc_type = result.get("type")
        if doc_type in EXTRACTABLE_TYPES:
            raw_text = extract_text_from_file(temp_path)

            if raw_text.startswith("OCR_ERROR"):
                result["extracted_fields"] = {
                    "name": None,
                    "address": None,
                    "date_of_birth": None,
                    "gender": None,
                    "id_number": None,
                    "mobile_number": None,
                    "ocr_error": raw_text
                }
            else:
                result["extracted_fields"] = extract_fields_with_sarvam(raw_text, doc_type)

        return JSONResponse(content=result)

    finally:
        if os.path.exists(temp_path):
            os.unlink(temp_path)


@app.get("/")
def read_root():
    return {"message": "ID Validator API is running"}