from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
import easyocr
import numpy as np
import cv2
import time
from google.cloud import vision
import dotenv
import hashlib
import time

dotenv.load_dotenv()

client = vision.ImageAnnotatorClient()

app = FastAPI()

# Cache storage: {hash: {'result': response, 'timestamp': time}}
cache = {}
CACHE_TTL = 600  # 10 minutes in seconds

def get_image_hash(image_bytes: bytes) -> str:
    """Generate SHA-256 hash of image bytes."""
    return hashlib.sha256(image_bytes).hexdigest()

def is_cache_valid(timestamp: float) -> bool:
    """Check if cache entry is still valid (within TTL)."""
    return (time.time() - timestamp) < CACHE_TTL

def clear_expired_cache():
    """Remove expired cache entries."""
    current_time = time.time()
    expired_keys = [key for key, value in cache.items() if not is_cache_valid(value['timestamp'])]
    for key in expired_keys:
        cache.pop(key, None)

@app.post('/ocr')
async def ocr(image: UploadFile = File(...)):
    print("ocr called", time.time())
    try:
        # Check if image file is provided
        if not image:
            raise HTTPException(status_code=400, detail="No image file provided")
        
        # Read image file into memory
        image_bytes = await image.read()
        
        # Generate hash of image bytes
        image_hash = get_image_hash(image_bytes)
        
        # Check cache
        clear_expired_cache()  # Clean up expired entries
        if image_hash in cache and is_cache_valid(cache[image_hash]['timestamp']):
            print("Serving from cache")
            return cache[image_hash]['result']
        
        # Process OCR if not in cache
        image = vision.Image(content=image_bytes)
        response = client.text_detection(image=image)
        texts = response.text_annotations

        text_results = []            
        for text in texts:
            text_results.append({
                'text': text.description
            })

        response_data = JSONResponse({
            'success': True,
            'results': text_results
        })

        # Store in cache
        cache[image_hash] = {
            'result': response_data,
            'timestamp': time.time()
        }

        return response_data

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host='0.0.0.0', port=6001)