# ocr_server.py - Clean natural language messages (no Reference number text)
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import base64
import binascii
from PIL import Image
import io
import torch
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
import re

app = FastAPI()

# Enable CORS for mobile app
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Request model for Base64
class OCRRequest(BaseModel):
    image_base64: str
    filename: str = "image.jpg"

# Load model
print("Loading OCR model...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_id = "prithivMLmods/coreOCR-7B-050325-preview"

processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
model = Qwen2VLForConditionalGeneration.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
).to(device).eval()

print(f"Model loaded on {device}")

def extract_numbers(text):
    numbers = re.findall(r'\d+', text)
    return ''.join(numbers) if numbers else ""

def base64_to_image(base64_string):
    if ',' in base64_string and base64_string.startswith('data:'):
        base64_string = base64_string.split(',', 1)[1]
    image_bytes = base64.b64decode(base64_string)
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    return image

@app.post("/ocr")
async def ocr_image(request: OCRRequest):
    try:
        # Convert base64 to image
        image = base64_to_image(request.image_base64)
        
        # Run OCR
        messages = [{
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": "Extract all numbers from this meter reading. Return only the numbers."},
            ]
        }]
        
        prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(
            text=[prompt_full],
            images=[image],
            return_tensors="pt",
        ).to(device)
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs, 
                max_new_tokens=200,
                temperature=0.1,
                do_sample=False
            )
        
        result = processor.decode(outputs[0], skip_special_tokens=True)
        numbers_only = extract_numbers(result)
        
        # Clean natural language messages
        if numbers_only:
            message = "Meter reading successfully extracted"
            ref_no = numbers_only[-14:] if len(numbers_only) >= 14 else numbers_only
        else:
            message = "No numbers found in the image. Please provide a clear meter reading photo"
            ref_no = ""
        
        return JSONResponse({
            "success": True,
            "message": message,
            "ref_no": ref_no,
            "numbers": numbers_only
        })
        
    except binascii.Error:
        return JSONResponse({
            "success": False,
            "message": "Invalid image format. Please send a valid Base64 encoded image"
        }, status_code=400)
        
    except Exception as e:
        error_message = str(e)
        if "image" in error_message.lower():
            message = "Could not process the image. Please ensure it's a valid photo of a meter reading"
        elif "timeout" in error_message.lower():
            message = "OCR processing timed out. Please try with a smaller or clearer image"
        else:
            message = "Failed to process image. Please try again"
        
        return JSONResponse({
            "success": False,
            "message": message
        }, status_code=500)

@app.get("/health")
async def health_check():
    return {
        "status": "ok",
        "model": "coreOCR-7B",
        "message": "OCR server is running normally"
    }

@app.post("/ocr-file")
async def ocr_image_file(file: bytes = None):
    """Legacy file upload endpoint for testing"""
    try:
        image = Image.open(io.BytesIO(file)).convert("RGB")
        
        messages = [{
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": "Extract all numbers from this meter reading. Return only the numbers."},
            ]
        }]
        
        prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(
            text=[prompt_full],
            images=[image],
            return_tensors="pt",
        ).to(device)
        
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.1, do_sample=False)
        
        result = processor.decode(outputs[0], skip_special_tokens=True)
        numbers_only = extract_numbers(result)
        
        if numbers_only:
            message = "Meter reading successfully extracted"
            ref_no = numbers_only[-14:] if len(numbers_only) >= 14 else numbers_only
        else:
            message = "No numbers found in the image. Please provide a clear meter reading photo"
            ref_no = ""
        
        return JSONResponse({
            "success": True,
            "message": message,
            "ref_no": ref_no,
            "numbers": numbers_only
        })
    except Exception as e:
        return JSONResponse({
            "success": False,
            "message": "Failed to process image. Please try again"
        }, status_code=500)

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