from fastapi import FastAPI, File, UploadFile, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import easyocr
import io
from PIL import Image
import numpy as np
import base64
import os
from pathlib import Path
from business_card import llm_calling

app = FastAPI()

# Initialize EasyOCR with fallback directories
def initialize_reader():
    possible_dirs = [
        "/tmp/.EasyOCR",  # Writable temp directory
        os.path.join(os.getcwd(), ".EasyOCR"),  # Current working directory
    ]

    for model_dir in possible_dirs:
        try:
            Path(model_dir).mkdir(parents=True, exist_ok=True)
            reader = easyocr.Reader(
                ['en', 'th'],
                gpu=False,
                model_storage_directory=model_dir,
                user_network_directory=model_dir,
                download_enabled=True
            )
            print(f"EasyOCR initialized successfully in {model_dir}")
            return reader
        except Exception as e:
            print(f"Failed to initialize in {model_dir}: {str(e)}")
            continue

    raise RuntimeError("Could not find a writable directory for EasyOCR")

reader = initialize_reader()

# Optional: for future use or documentation
class ImageBase64Payload(BaseModel):
    image: str  # base64 string of image

@app.get("/")
def root():
    return {"message": "Hello from Hugging Face FastAPI!"}

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

@app.post("/extract")
async def extract_business_card(request: Request, file: UploadFile = File(None)):
    try:
        # Case 1: Image uploaded as a file
        if file:
            image_data = await file.read()
        else:
            # Case 2: Base64 string in JSON body
            json_data = await request.json()
            base64_str = json_data.get("image", "")

            if not base64_str:
                return JSONResponse(
                    status_code=400,
                    content={"status": "error", "message": "'image' field is required in JSON payload"},
                )

            if "," in base64_str:
                base64_str = base64_str.split(",")[1]

            try:
                image_data = base64.b64decode(base64_str)
            except Exception:
                return JSONResponse(
                    status_code=400,
                    content={"status": "error", "message": "Invalid base64 string"},
                )

        # Decode image into NumPy format
        image = Image.open(io.BytesIO(image_data)).convert("RGB")
        image_np = np.array(image)

        # Perform OCR
        ocr_result = reader.readtext(image_np, detail=0)
        full_text = " ".join(ocr_result)

        # Call LLM to extract structured info
        extracted_data = llm_calling(full_text)

        return {
            "status": "success",
            "data": extracted_data,
            "ocr_text": full_text  # Optional for debugging
        }

    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"status": "error", "message": str(e), "type": type(e).__name__},
        )