from fastapi import FastAPI, File, UploadFile, HTTPException

import pytesseract
import cv2
import os
from PIL import Image
import json
import unicodedata
from pdf2image import convert_from_bytes
from pypdf import PdfReader
import numpy as np
from typing import List
import io
import logging
import time
import asyncio
import psutil
import cachetools
import hashlib
import google.generativeai as genai
from dotenv import load_dotenv

app = FastAPI()

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %
(message)s')
logger = logging.getLogger(__name__)

# Load environment variables

load_dotenv()

# Configure Gemini API

api_key = os.getenv("GOOGLE_API_KEY")
if not api_key:
logger.error("GOOGLE_API_KEY not set")
raise HTTPException(status_code=500, detail="GOOGLE_API_KEY not set")
genai.configure(api_key=api_key)
model = genai.GenerativeModel("gemini-2.0-flash")

# Set Tesseract path

pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"

# In-memory caches (1-hour TTL)

raw_text_cache = cachetools.TTLCache(maxsize=100, ttl=3600)
structured_data_cache = cachetools.TTLCache(maxsize=100, ttl=3600)

def log_memory_usage():
"""Log current memory usage."""
process = psutil.Process()
mem_info = process.memory_info()
return f"Memory usage: {mem_info.rss / 1024 / 1024:.2f} MB"

def get_file_hash(file_bytes):
"""Generate MD5 hash of file content."""
return hashlib.md5(file_bytes).hexdigest()

def get_text_hash(raw_text):
"""Generate MD5 hash of raw text."""
return hashlib.md5(raw_text.encode('utf-8')).hexdigest()
def get_poppler_path():
"""Determine the correct poppler path based on the system."""
import platform
import shutil

# Check if poppler utilities are in PATH

if shutil.which('pdftoppm'):
return None # Use system PATH

# Common poppler paths for different systems

common_paths = [
"/usr/bin", # Linux
"/usr/local/bin", # macOS with Homebrew
"/opt/homebrew/bin", # macOS with Apple Silicon Homebrew
"/usr/share/poppler/bin", # Some Linux distributions
"C:\\poppler\\Library\\bin", # Windows
"C:\\Program Files\\poppler\\bin", # Windows alternative
]

for path in common_paths:

if os.path.exists(os.path.join(path, "pdftoppm")) or
os.path.exists(os.path.join(path, "pdftoppm.exe")):
return path

return None

async def process_image(img_bytes, filename, idx):

"""Process a single image (JPG/JPEG/PNG) with OCR."""
start_time = time.time()
logger.info(f"Starting OCR for {filename} image {idx}, {log_memory_usage()}")
try:
img = Image.open(io.BytesIO(img_bytes))
img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
img_pil = Image.fromarray(cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB))
custom_config = r'--oem 1 --psm 6 -l eng+ara' # Reduced for performance
page_text = pytesseract.image_to_string(img_pil, config=custom_config)
logger.info(f"Completed OCR for {filename} image {idx}, took {time.time() -
start_time:.2f} seconds, {log_memory_usage()}")
return page_text + "\n"
except Exception as e:
logger.error(f"OCR failed for {filename} image {idx}: {str(e)},
{log_memory_usage()}")
return ""

async def process_pdf_page(img, page_idx):

"""Process a single PDF page with OCR."""
start_time = time.time()
logger.info(f"Starting OCR for PDF page {page_idx}, {log_memory_usage()}")
try:
img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
img_pil = Image.fromarray(cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB))
custom_config = r'--oem 1 --psm 6 -l eng+ara' # Reduced for performance
page_text = pytesseract.image_to_string(img_pil, config=custom_config)
logger.info(f"Completed OCR for PDF page {page_idx}, took {time.time() -
start_time:.2f} seconds, {log_memory_usage()}")
return page_text + "\n"
 except Exception as e:
logger.error(f"OCR failed for PDF page {page_idx}: {str(e)},
{log_memory_usage()}")
return ""

async def process_with_gemini(filename: str, raw_text: str):

"""Process raw text with Gemini to extract structured data."""
start_time = time.time()
logger.info(f"Starting Gemini processing for {filename}, {log_memory_usage()}")

# Check structured data cache

text_hash = get_text_hash(raw_text)
if text_hash in structured_data_cache:
logger.info(f"Structured data cache hit for {filename},
{log_memory_usage()}")
return structured_data_cache[text_hash]

# Truncate text for Gemini

if len(raw_text) > 10000:
raw_text = raw_text[:10000]
logger.info(f"Truncated raw text for {filename} to 10000 characters,
{log_memory_usage()}")

try:
prompt = f"""
You are an intelligent invoice data extractor. Given raw text from an
invoice in any language and extract key business fields in the specified JSON
format. Support English. Handle synonyms (e.g., 'total' = 'net', 'tax' =
'GST'/'TDS'). The 'Products' field is dynamic and may contain multiple items, each
with 'qty', 'description', 'unit_price', and 'amount'. Detect the currency (e.g.,
USD, INR, EUR) from symbols ($, ₹, €) or text; default to USD if unclear. If a
field is missing, include it with an empty string ("") or appropriate default
(e.g., 0 for numbers).

Raw text:
{raw_text}

Output JSON:
{{
"invoice": {{
"invoice_number": "",
"invoice_date": "YYYY-MM-DD",
"due_date": "YYYY-MM-DD",
"purchase_order_number": "",
"vendor": {{
"vendor_id": "",
"name": "",
"address": {{
"line1": "",
"line2": "",
"city": "",
"state": "",
"postal_code": "",
"country": ""
}},
"contact": {{
"email": "",
"phone": ""
}},
 "tax_id": ""
}},
"buyer": {{
"buyer_id": "",
"name": "",
"address": {{
"line1": "",
"line2": "",
"city": "",
"state": "",
"postal_code": "",
"country": ""
}},
"contact": {{
"email": "",
"phone": ""
}},
"tax_id": ""
}},
"items": [
{{
"item_id": "",
"description": "",
"quantity": 0,
"unit_of_measure": "",
"unit_price": 0,
"total_price": 0,
"tax_rate": 0,
"tax_amount": 0,
"discount": 0,
"net_amount": 0
}}
],
"sub_total": 0,
"tax_total": 0,
"discount_total": 0,
"total_amount": 0,
"currency": ""
}}
}}
"""
response = model.generate_content(prompt)
llm_output = response.text
json_start = llm_output.find("{")
json_end = llm_output.rfind("}") + 1
json_str = llm_output[json_start:json_end]
structured_data = json.loads(json_str)
structured_data_cache[text_hash] = structured_data
logger.info(f"Gemini processing for {filename}, took {time.time() -
start_time:.2f} seconds, {log_memory_usage()}")
return structured_data
except Exception as e:
logger.error(f"Gemini processing failed for {filename}: {str(e)},
{log_memory_usage()}")
return {"error": f"Gemini processing failed: {str(e)}"}

@app.post("/ocr")
async def extract_and_structure(files: List[UploadFile] = File(...)):
output_json = {
 "success": True,
"message": "",
"data": []
}
success_count = 0
fail_count = 0

logger.info(f"Starting processing for {len(files)} files,

{log_memory_usage()}")

for file in files:

total_start_time = time.time()
logger.info(f"Processing file: {file.filename}, {log_memory_usage()}")

# Validate file format

valid_extensions = {'.pdf', '.jpg', '.jpeg', '.png'}
file_ext = os.path.splitext(file.filename.lower())[1]
if file_ext not in valid_extensions:
fail_count += 1
output_json["data"].append({
"filename": file.filename,
"structured_data": {"error": f"Unsupported file format:
{file_ext}"},
"error": f"Unsupported file format: {file_ext}"
})
logger.error(f"Unsupported file format for {file.filename}:
{file_ext}")
continue

# Read file into memory

try:
file_start_time = time.time()
file_bytes = await file.read()
file_stream = io.BytesIO(file_bytes)
file_hash = get_file_hash(file_bytes)
logger.info(f"Read file {file.filename}, took {time.time() -
file_start_time:.2f} seconds, size: {len(file_bytes)/1024:.2f} KB,
{log_memory_usage()}")
except Exception as e:
fail_count += 1
output_json["data"].append({
"filename": file.filename,
"structured_data": {"error": f"Failed to read file: {str(e)}"},
"error": f"Failed to read file: {str(e)}"
})
logger.error(f"Failed to read file {file.filename}: {str(e)},
{log_memory_usage()}")
continue

# Check raw text cache

raw_text = ""
if file_hash in raw_text_cache:
raw_text = raw_text_cache[file_hash]
logger.info(f"Raw text cache hit for {file.filename},
{log_memory_usage()}")
else:
if file_ext == '.pdf':
# Try extracting embedded text
try:
 extract_start_time = time.time()
reader = PdfReader(file_stream)
for page in reader.pages:
text = page.extract_text()
if text:
raw_text += text + "\n"
logger.info(f"Embedded text extraction for {file.filename},
took {time.time() - extract_start_time:.2f} seconds, text length: {len(raw_text)},
{log_memory_usage()}")
except Exception as e:
logger.warning(f"Embedded text extraction failed for
{file.filename}: {str(e)}, {log_memory_usage()}")

# If no embedded text, perform OCR

if not raw_text.strip():
try:
convert_start_time = time.time()

# Get the correct poppler path

poppler_path = get_poppler_path()

# Convert PDF to images with proper poppler path handling

if poppler_path:
images = convert_from_bytes(file_bytes,
poppler_path=poppler_path, dpi=100)
logger.info(f"Using poppler path: {poppler_path}")
else:
# Try without specifying poppler_path (use system PATH)
images = convert_from_bytes(file_bytes, dpi=100)
logger.info("Using poppler from system PATH")

logger.info(f"PDF to images conversion for {file.filename},

{len(images)} pages, took {time.time() - convert_start_time:.2f} seconds,
{log_memory_usage()}")

ocr_start_time = time.time()
page_texts = []
for i, img in enumerate(images):
page_text = await process_pdf_page(img, i)
page_texts.append(page_text)
raw_text = "".join(page_texts)
logger.info(f"Total OCR for {file.filename}, took
{time.time() - ocr_start_time:.2f} seconds, text length: {len(raw_text)},
{log_memory_usage()}")
except Exception as e:
fail_count += 1
error_msg = f"OCR failed: {str(e)}"
if "poppler" in str(e).lower():
error_msg += ". Please ensure Poppler is installed and
accessible in PATH."
output_json["data"].append({
"filename": file.filename,
"structured_data": {"error": error_msg},
"error": error_msg
})
logger.error(f"OCR failed for {file.filename}: {str(e)},
{log_memory_usage()}")
continue
else: # JPG/JPEG/PNG
 try:
ocr_start_time = time.time()
raw_text = await process_image(file_bytes, file.filename, 0)
logger.info(f"Image OCR for {file.filename}, took {time.time()
- ocr_start_time:.2f} seconds, text length: {len(raw_text)}, {log_memory_usage()}")
except Exception as e:
fail_count += 1
output_json["data"].append({
"filename": file.filename,
"structured_data": {"error": f"Image OCR failed:
{str(e)}"},
"error": f"Image OCR failed: {str(e)}"
})
logger.error(f"Image OCR failed for {file.filename}: {str(e)},
{log_memory_usage()}")
continue

# Normalize text
try:
normalize_start_time = time.time()
raw_text = unicodedata.normalize('NFKC', raw_text)
raw_text = raw_text.encode().decode('utf-8')
raw_text_cache[file_hash] = raw_text
logger.info(f"Text normalization for {file.filename}, took
{time.time() - normalize_start_time:.2f} seconds, text length: {len(raw_text)},
{log_memory_usage()}")
except Exception as e:
logger.warning(f"Text normalization failed for {file.filename}:
{str(e)}, {log_memory_usage()}")

# Process with Gemini

structured_data = await process_with_gemini(file.filename, raw_text)
success_count += 1
output_json["data"].append({
"filename": file.filename,
"structured_data": structured_data,
"error": ""
})

logger.info(f"Total processing for {file.filename}, took {time.time() -

total_start_time:.2f} seconds, {log_memory_usage()}")

output_json["message"] = f"Processed {len(files)} files. {success_count}

succeeded, {fail_count} failed."
if fail_count > 0 and success_count == 0:
output_json["success"] = False

logger.info(f"Completed processing for {len(files)} files, {success_count}

succeeded, {fail_count} failed, {log_memory_usage()}")
return output_json