app.py

# from fastapi import FastAPI, UploadFile, File, Header
# from ultralytics import YOLO
# import cv2
# import numpy as np
# import io
# from PIL import Image
# from datetime import datetime, timezone
# import uvicorn
# import os
# import time
# import json
# import base64
# import pytesseract  # ✅ OCR
# import requests
# import jwt
# from jwt import PyJWKClient

# try:
#     from dotenv import load_dotenv  # type: ignore
#     load_dotenv()
# except Exception:
#     # optional in production containers
#     pass

# from langchain_openai import ChatOpenAI
# from langchain.prompts import ChatPromptTemplate
# try:
#     from langchain_core.messages import HumanMessage
# except Exception:  # Fallback for older langchain
#     from langchain.schema import HumanMessage  # type: ignore

# # Init FastAPI
# app = FastAPI()


# @app.get("/")
# def root():
#     return {
#         "service": "Receipt Scanner API",
#         "status": "ok",
#         "docs": "/docs",
#         "endpoints": [
#             "/process-receipt",
#             "/process-receipt-vision",
#         ],
#     }

# # Ensure output folder exists
# os.makedirs("receiptimages", exist_ok=True)

# # Load trained YOLOv8n model
# model = YOLO("receipt-scanner/models/best.pt")

# # Init LangChain with OpenAI
# openai_api_key = os.getenv("OPENAI_API_KEY")
# if not openai_api_key:
#     print("⚠️ OPENAI_API_KEY is not set; LangChain calls will fail until configured.")

# TEXT_LLM_MODEL = os.getenv("TEXT_LLM_MODEL", "gpt-5-mini")
# VISION_LLM_MODEL = os.getenv("VISION_LLM_MODEL", "gpt-4o")

# llm_text = ChatOpenAI(
#     model=TEXT_LLM_MODEL,
#     temperature=0,
# )

# llm_vision = ChatOpenAI(
#     model=VISION_LLM_MODEL,
#     temperature=0,
# )

# # Supabase config
# SUPABASE_URL = os.getenv("SUPABASE_URL", "").rstrip("/")
# SUPABASE_JWKS_URL = os.getenv("SUPABASE_JWKS_URL", "")
# SUPABASE_SERVICE_ROLE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY", "")
# SUPABASE_STORAGE_BUCKET = os.getenv("SUPABASE_STORAGE_BUCKET", "receipts")
# SUPABASE_ANON_KEY = os.getenv("SUPABASE_ANON_KEY", "")
# SUPABASE_RECEIPTS_TABLE = os.getenv("SUPABASE_RECEIPTS_TABLE", "receipts")
# SUPABASE_CATEGORIES_TABLE = os.getenv("SUPABASE_CATEGORIES_TABLE", "categories")

# _JWKS_CLIENT = None
# if SUPABASE_JWKS_URL:
#     try:
#         _JWKS_CLIENT = PyJWKClient(SUPABASE_JWKS_URL)
#     except Exception as _exc:
#         print(f"⚠️ Failed to initialize JWKS client: {_exc}")

# # Define the prompt for categorization (restored detailed rules)
# UNIFIED_PROMPT_RULES = """
# You are an expert at analyzing receipt text extracted with OCR. 
# The text may contain noise or extra characters, but focus only on useful parts.

# Rules:
# 1. **Venue**:
#    - This is the store/restaurant/tuck shop/retailer where the receipt comes from.
#    - It is usually near the top of the receipt (sometimes next to a logo or under "Shop/Store/Restaurant").
#    - If multiple candidates appear, pick the most likely actual business name (ignore words like "Customer Copy", "Approved", "MasterCard").
#    - If both a brand name and location/branch are present, include both (e.g. "Checkers - Brooklyn Mall").
#    - Ignore banks, payment processors, or card brands (e.g., "Nedbank", "FNB", "MasterCard", "Visa") unless they are the ONLY candidate.
#    - If uncertain, check for context clues like items, meals, or store names to identify the real merchant.
#    - If you are not confident, return `null`.

# 2. **Date**:
#    - Look for patterns like `DD-MM-YY`, `YYYY-MM-DD`, `DD/MM/YYYY`, or combined with time (e.g., `2025-09-15 14:23`).
#    - If only time is found without a date, return `null`.

# 3. **Total**:
#    - Find the final purchase amount.
#    - Prioritize lines with "TOTAL", "Amount Due", "Purchase", "Grand Total".
#    - Include the **currency symbol/code** if present (e.g., `R238.00`, `$12.50`).
#    - If no currency is visible, return just the number.

# 4. **Category**:
#    - The category MUST be chosen only from the following fixed list:
#      - "Office Supplies"
#      - "Travel/Transport"
#      - "Utilities"
#      - "Wages/Staff costs"
#      - "Meals & Entertainment"
#      - "Other Business Expenses"
#    - Infer the category from the venue/items. 
#    - If no clear match, default to "Other Business Expenses".
#    - Do not invent or suggest categories outside this list
#    - Just also take note that if the receipt or content contains petrol stations from south africa and there is price per litre, or a pump number, or a price per litre, then the category should be Travel/Transport.

# 5. **Output format**:
#    Return strictly valid JSON with these keys:
#    {{
#      "venue": "...",
#      "date": "...",
#      "total": "...",
#      "category": "..."
#    }}
#    - Do not add extra text or explanation outside the JSON.
# """

# text_prompt = ChatPromptTemplate.from_template(
#     UNIFIED_PROMPT_RULES
#     + "\n---\n\nReceipt OCR text:\n{ocr_text}\n"
# )


# def run_langchain_with_text(receipt_text: str, image_path: str):
#     """Send OCR text to GPT for structured JSON, with logs."""
#     print(f"📜 OCR extracted text from {image_path}:\n{receipt_text}")

#     chain = text_prompt | llm_text
#     result = chain.invoke({"ocr_text": receipt_text})  # ✅ match prompt var

#     # 🔥 Debug logging
#     print(f"📤 GPT raw response: {result}")
#     print(f"📄 Extracted content: {result.content}")

#     return result.content


# def _verify_user_via_auth_api(authorization: str, x_user_id: str | None) -> str:
#     token = authorization.split(" ", 1)[1]
#     if not SUPABASE_URL or not SUPABASE_ANON_KEY:
#         raise Exception("Supabase auth not configured")
#     resp = requests.get(
#         f"{SUPABASE_URL}/auth/v1/user",
#         headers={
#             "Authorization": f"Bearer {token}",
#             "apikey": SUPABASE_ANON_KEY,
#         },
#         timeout=20,
#     )
#     if resp.status_code != 200:
#         raise Exception("Invalid JWT")
#     data = resp.json() or {}
#     sub = data.get("id") or data.get("sub")
#     if not sub or (x_user_id and x_user_id != sub):
#         raise Exception("User mismatch")
#     return str(sub)


# def verify_user(authorization: str | None, x_user_id: str | None) -> str:
#     if not authorization or not authorization.lower().startswith("bearer "):
#         raise_value = "Missing bearer token"
#         raise Exception(raise_value)
#     # Prefer JWKS if provided, else fall back to Auth API using anon key
#     if _JWKS_CLIENT and SUPABASE_URL:
#         token = authorization.split(" ", 1)[1]
#         signing_key = _JWKS_CLIENT.get_signing_key_from_jwt(token).key
#         claims = jwt.decode(
#             token,
#             signing_key,
#             algorithms=["RS256"],
#             options={"verify_aud": False},
#             issuer=f"{SUPABASE_URL}/auth/v1",
#         )
#         sub = claims.get("sub")
#         if not sub or (x_user_id and x_user_id != sub):
#             raise Exception("User mismatch")
#         return str(sub)
#     if SUPABASE_ANON_KEY:
#         return _verify_user_via_auth_api(authorization, x_user_id)
#     raise Exception("No auth method configured (set SUPABASE_JWKS_URL or SUPABASE_ANON_KEY)")


# def ensure_extraction_complete(extracted_obj: dict):
#     """Ensure required fields were extracted before persisting.
#     Requires venue, total, and date to be present and non-empty.
#     """
#     from fastapi import HTTPException
#     required = ["venue", "total", "date"]
#     missing = [k for k in required if not extracted_obj.get(k)]
#     if missing:
#         raise HTTPException(status_code=422, detail=f"Extraction incomplete. Missing: {', '.join(missing)}")


# def parse_extracted_json(raw_text: str) -> dict:
#     """Parse model output that may be wrapped in triple backticks or contain extra text.

#     Strategy:
#     - Trim, find the largest JSON object between first '{' and last '}'
#     - Attempt json.loads; fallback to stripping ```json fences
#     """
#     if not raw_text:
#         return {}
#     text = str(raw_text).strip()
#     # Try bracket slicing first
#     start = text.find("{")
#     end = text.rfind("}")
#     if start != -1 and end != -1 and end > start:
#         candidate = text[start : end + 1]
#         try:
#             return json.loads(candidate)
#         except Exception:
#             pass
#     # Remove code fences if present
#     if text.startswith("```") and text.endswith("```"):
#         # Strip fences lines
#         lines = text.splitlines()
#         if lines and lines[0].startswith("```"):
#             lines = lines[1:]
#         if lines and lines[-1].startswith("```"):
#             lines = lines[:-1]
#         text2 = "\n".join(lines).strip()
#         try:
#             return json.loads(text2)
#         except Exception:
#             # Try bracket slice again
#             start = text2.find("{")
#             end = text2.rfind("}")
#             if start != -1 and end != -1 and end > start:
#                 candidate = text2[start : end + 1]
#                 try:
#                     return json.loads(candidate)
#                 except Exception:
#                     return {}
#             return {}
#     return {}


# def upload_to_storage(path: str, content: bytes, content_type: str | None) -> str:
#     if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
#         raise Exception("Supabase storage not configured")
#     url = f"{SUPABASE_URL}/storage/v1/object/{SUPABASE_STORAGE_BUCKET}/{path}"
#     resp = requests.post(
#         url,
#         data=content,
#         headers={
#             "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
#             "apikey": SUPABASE_SERVICE_ROLE_KEY,
#             "Content-Type": content_type or "application/octet-stream",
#         },
#         timeout=60,
#     )
#     if not resp.ok:
#         raise Exception(f"Storage upload failed: {resp.text}")
#     return f"{SUPABASE_STORAGE_BUCKET}/{path}"


# def insert_receipt(row: dict) -> dict:
#     if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
#         raise Exception("Supabase REST not configured")
#     url = f"{SUPABASE_URL}/rest/v1/{SUPABASE_RECEIPTS_TABLE}"
#     resp = requests.post(
#         url,
#         json=row,
#         headers={
#             "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
#             "apikey": SUPABASE_SERVICE_ROLE_KEY,
#             "Content-Type": "application/json",
#             "Prefer": "return=representation",
#         },
#         timeout=60,
#     )
#     if not resp.ok:
#         raise Exception(f"DB insert failed: {resp.text}")
#     data = resp.json()
#     return data[0] if isinstance(data, list) and data else data


# def update_receipt(receipt_id: str, updates: dict) -> dict:
#     if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
#         raise Exception("Supabase REST not configured")
#     url = f"{SUPABASE_URL}/rest/v1/{SUPABASE_RECEIPTS_TABLE}?id=eq.{receipt_id}"
#     resp = requests.patch(
#         url,
#         json=updates,
#         headers={
#             "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
#             "apikey": SUPABASE_SERVICE_ROLE_KEY,
#             "Content-Type": "application/json",
#             "Prefer": "return=representation",
#         },
#         timeout=60,
#     )
#     if not resp.ok:
#         raise Exception(f"DB update failed: {resp.text}")
#     data = resp.json()
#     return data[0] if isinstance(data, list) and data else data


# def upsert_category_if_needed(category_name: str | None):
#     if not category_name:
#         return
#     if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
#         raise Exception("Supabase REST not configured")
#     url = f"{SUPABASE_URL}/rest/v1/{SUPABASE_CATEGORIES_TABLE}?on_conflict=name"
#     resp = requests.post(
#         url,
#         json={"name": category_name},
#         headers={
#             "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
#             "apikey": SUPABASE_SERVICE_ROLE_KEY,
#             "Content-Type": "application/json",
#             "Prefer": "resolution=merge-duplicates",
#         },
#         timeout=30,
#     )
#     # It's ok if 409 or a duplicate; treat non-2xx as warning only
#     if not resp.ok and resp.status_code not in (409,):
#         print(f"⚠️ Category upsert failed: {resp.status_code} {resp.text}")


# def parse_total(value: str | float | int | None) -> float:
#     if value is None:
#         return 0.0
#     if isinstance(value, (int, float)):
#         return float(value)
#     # strip currency and non-numeric
#     import re
#     s = str(value)
#     m = re.findall(r"[-+]?[0-9]*\.?[0-9]+", s.replace(",", ""))
#     return float(m[-1]) if m else 0.0


# # def detect_and_crop_largest_receipt(img_array: np.ndarray):
# #     """Detect the largest bounding box via YOLO and return crop and bbox.

# #     Returns (cropped_array, bbox_dict)
# #     bbox_dict = {"x1": int, "y1": int, "x2": int, "y2": int}
# #     """
# #     results = model(img_array)
# #     # Default fallback to full image
# #     cropped = img_array
# #     bbox = None
# #     for r in results:
# #         if len(r.boxes.xyxy) == 0:
# #             continue
# #         boxes = r.boxes.xyxy
# #         i = max(range(len(boxes)), key=lambda j: (boxes[j][2]-boxes[j][0]) * (boxes[j][3]-boxes[j][1]))
# #         x1, y1, x2, y2 = map(int, boxes[i].tolist())
# #         cropped = img_array[y1:y2, x1:x2]
# #         bbox = {"x1": x1, "y1": y1, "x2": x2, "y2": y2}
# #         break
# #     return cropped, bbox

# def detect_and_crop_largest_receipt(img_array: np.ndarray):
#     """Detect receipt via YOLO and return the crop of the highest-confidence box.

#     Returns (cropped_array, bbox_dict)
#     bbox_dict = {"x1": int, "y1": int, "x2": int, "y2": int, "conf": float}
#     """
#     results = model(img_array)
#     cropped = img_array
#     bbox = None

#     for r in results:
#         if len(r.boxes.xyxy) == 0:
#             continue

#         boxes = r.boxes.xyxy.cpu().numpy()
#         scores = r.boxes.conf.cpu().numpy()

#         # Pick index of highest confidence
#         i = int(scores.argmax())

#         x1, y1, x2, y2 = map(int, boxes[i].tolist())
#         cropped = img_array[y1:y2, x1:x2]
#         bbox = {
#             "x1": x1,
#             "y1": y1,
#             "x2": x2,
#             "y2": y2,
#             "conf": float(scores[i]),
#         }
#         break  # only use first image in batch

#     return cropped, bbox



# def encode_image_to_data_url(img_array: np.ndarray, format: str = "JPEG") -> str:
#     """Encode an RGB image array to a data URL suitable for GPT-4o image input."""
#     pil_img = Image.fromarray(img_array)
#     buffer = io.BytesIO()
#     pil_img.save(buffer, format=format)
#     b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
#     mime = "image/jpeg" if format.upper() == "JPEG" else "image/png"
#     return f"data:{mime};base64,{b64}"


# def encode_image_to_jpeg_bytes(img_array: np.ndarray) -> bytes:
#     """Encode an RGB image array to JPEG bytes."""
#     pil_img = Image.fromarray(img_array)
#     buffer = io.BytesIO()
#     pil_img.save(buffer, format="JPEG")
#     return buffer.getvalue()

# @app.post("/process-receipt")
# async def process_receipt(
#     file: UploadFile = File(...),
#     authorization: str | None = Header(default=None, alias="Authorization"),
#     x_user_id: str | None = Header(default=None, alias="X-User-Id"),
# ):
#     """
#     Upload receipt image -> detect largest -> crop -> OCR -> GPT categorize -> JSON result
#     """
#     # Verify user & read file
#     user_id = None
#     try:
#         user_id = verify_user(authorization, x_user_id)
#     except Exception as exc:
#         from fastapi import HTTPException
#         raise HTTPException(status_code=401, detail=str(exc))

#     contents = await file.read()
#     image = Image.open(io.BytesIO(contents)).convert("RGB")
#     img_array = np.array(image)

#     outputs = []
#     # Detect and crop once per image (YOLO returns one result for input image)
#     cropped, bbox = detect_and_crop_largest_receipt(img_array)

#     # Save cropped or fallback image
#     save_path = f"receiptimages/receipt_main.jpg"
#     cv2.imwrite(save_path, cv2.cvtColor(cropped, cv2.COLOR_RGB2BGR))
#     print(f"✅ Saved {save_path}")

#     # OCR step
#     ocr_text = pytesseract.image_to_string(Image.fromarray(cropped))

#     # ✅ Fallback to full image if OCR is empty
#     if not ocr_text.strip():
#         print("⚠️ Empty OCR text, retrying with full image")
#         ocr_text = pytesseract.image_to_string(image)

#     # Send OCR text to GPT
#     llm_result = run_langchain_with_text(ocr_text, save_path)

#     # Parse extracted JSON
#     extracted_obj = parse_extracted_json(llm_result)

#     # Upload CROPPED image to Storage (JPEG)
#     cropped_bytes = encode_image_to_jpeg_bytes(cropped)
#     storage_path = f"{user_id}/{int(time.time()*1000)}_cropped.jpg"
#     try:
#         stored = upload_to_storage(storage_path, cropped_bytes, "image/jpeg")
#     except Exception as exc:
#         from fastapi import HTTPException
#         raise HTTPException(status_code=502, detail=f"Storage upload failed: {exc}")

#     # Phase 1: insert pending record first
#     pending_row = {
#         "userId": user_id,
#         "venue": "Pending",
#         "category": None,
#         "total": 0.0,
#         "purchased_at": datetime.now(timezone.utc).isoformat(),
#         "status": "pending",
#         "storage_path": stored,
#         "mime_type": "image/jpeg",
#         "file_size_bytes": len(cropped_bytes),
#     }
#     try:
#         saved_pending = insert_receipt(pending_row)
#     except Exception as exc:
#         from fastapi import HTTPException
#         raise HTTPException(status_code=502, detail=f"DB insert failed: {exc}")

#     # Ensure extraction is complete before update
#     ensure_extraction_complete(extracted_obj)

#     # Phase 2: update same record with extracted values
#     updates = {
#         "venue": extracted_obj.get("venue"),
#         "category": extracted_obj.get("category"),
#         "total": parse_total(extracted_obj.get("total")),
#         "purchased_at": extracted_obj.get("date") or saved_pending.get("purchased_at"),
#     }
#     try:
#         saved_final = update_receipt(saved_pending["id"], updates)
#     except Exception as exc:
#         from fastapi import HTTPException
#         raise HTTPException(status_code=502, detail=f"DB update failed: {exc}")

#     # Upsert category name
#     try:
#         upsert_category_if_needed(extracted_obj.get("category"))
#     except Exception as exc:
#         print(f"⚠️ Category upsert error: {exc}")

#     return {"receiptId": saved_final.get("id"), "status": "pending"}


# @app.post("/process-receipt-vision")
# async def process_receipt_vision(
#     file: UploadFile = File(...),
#     authorization: str | None = Header(default=None, alias="Authorization"),
#     x_user_id: str | None = Header(default=None, alias="X-User-Id"),
# ):
#     """
#     Upload receipt image -> detect largest -> crop -> send image to GPT-4o -> JSON result
#     No OCR is performed; the model reads directly from the image.
#     """
#     # Verify user & read file
#     user_id = None
#     try:
#         user_id = verify_user(authorization, x_user_id)
#     except Exception as exc:
#         from fastapi import HTTPException
#         raise HTTPException(status_code=401, detail=str(exc))

#     contents = await file.read()
#     image = Image.open(io.BytesIO(contents)).convert("RGB")
#     img_array = np.array(image)

#     cropped, bbox = detect_and_crop_largest_receipt(img_array)

#     save_path = f"receiptimages/receipt_vision.jpg"
#     cv2.imwrite(save_path, cv2.cvtColor(cropped, cv2.COLOR_RGB2BGR))
#     print(f"✅ Saved {save_path}")
#     print(f"🔎 Vision bbox: {bbox if bbox else 'none (full image)'}")
#     print(f"🧠 Vision model: {VISION_LLM_MODEL}")

#     data_url = encode_image_to_data_url(cropped, format="JPEG")
#     print(f"🖼️ Encoded image data URL length: {len(data_url)} (showing first 64): {data_url[:64]}...")

#     vision_instructions = UNIFIED_PROMPT_RULES
#     print("📝 Vision prompt (first 240 chars):\n" + vision_instructions[:240] + ("..." if len(vision_instructions) > 240 else ""))

#     message = HumanMessage(
#         content=[
#             {"type": "text", "text": vision_instructions},
#             {"type": "image_url", "image_url": {"url": data_url}},
#         ]
#     )
#     try:
#         result = llm_vision.invoke([message])
#     except Exception as exc:
#         print(f"❌ Vision LLM call failed: {exc}")
#         from fastapi import HTTPException
#         raise HTTPException(status_code=502, detail=f"Vision LLM failed: {exc}")

#     print(f"📤 Vision LLM raw result: {result}")
#     print(f"📄 Vision LLM content: {getattr(result, 'content', str(result))}")

#     # Parse extracted JSON
#     content_str = getattr(result, "content", str(result))
#     extracted_obj = parse_extracted_json(content_str)

#     cropped_bytes = encode_image_to_jpeg_bytes(cropped)

#     # Build storage path with category slug
#     category = extracted_obj.get("category")
#     def category_to_slug(category: str | None) -> str:
#         if not category:
#             return "uncategorized"
#         return (
#             category.strip().lower()
#             .replace("&", "and")
#             .replace("/", "_")
#             .replace(" ", "_")
#             .replace("-", "_")
#         )

#     category_slug = category_to_slug(category)
#     storage_path = f"{user_id}/{category_slug}/{int(time.time()*1000)}_cropped.jpg"

#     try:
#         stored = upload_to_storage(storage_path, cropped_bytes, "image/jpeg")

#     except Exception as exc:
#         from fastapi import HTTPException
#         raise HTTPException(status_code=502, detail=f"Storage upload failed: {exc}")

#     # Phase 1: insert pending record first
#     pending_row = {
#         "userId": user_id,
#         "venue": "Pending",
#         "category": None,
#         "total": 0.0,
#         "purchased_at": datetime.now(timezone.utc).isoformat(),
#         "status": "pending",
#         "storage_path": stored,
#         "mime_type": "image/jpeg",
#         "file_size_bytes": len(cropped_bytes),
#     }
#     try:
#         saved_pending = insert_receipt(pending_row)
#     except Exception as exc:
#         from fastapi import HTTPException
#         raise HTTPException(status_code=502, detail=f"DB insert failed: {exc}")

#     # Ensure extraction is complete before update
#     ensure_extraction_complete(extracted_obj)

#     # Phase 2: update same record with extracted values
#     updates = {
#         "venue": extracted_obj.get("venue"),
#         "category": extracted_obj.get("category"),
#         "total": parse_total(extracted_obj.get("total")),
#         "purchased_at": extracted_obj.get("date") or saved_pending.get("purchased_at"),
#     }
#     try:
#         saved_final = update_receipt(saved_pending["id"], updates)
#     except Exception as exc:
#         from fastapi import HTTPException
#         raise HTTPException(status_code=502, detail=f"DB update failed: {exc}")

#     # Upsert category name
#     try:
#         upsert_category_if_needed(extracted_obj.get("category"))
#     except Exception as exc:
#         print(f"⚠️ Category upsert error: {exc}")

#     return {"receiptId": saved_final.get("id"), "status": "pending"}

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

from fastapi import FastAPI, UploadFile, File, Header
from ultralytics import YOLO
import cv2
import numpy as np
import io
from PIL import Image
from datetime import datetime, timezone
import uvicorn
import os
import time
import json
import base64
import ssl as _ssl  # ✅ proper SSL context for asyncpg
import certifi



# Async + performance
import asyncio
import asyncpg
import httpx
from concurrent.futures import ThreadPoolExecutor

# Optional: reduce thread oversubscription on 2 vCPU
try:
    import torch
    torch.set_num_threads(1)
    torch.set_num_interop_threads(1)
except Exception:
    pass
try:
    cv2.setNumThreads(1)
except Exception:
    pass

# Optional: faster event loop
try:
    import uvloop
    uvloop.install()
except Exception:
    pass

# Optional in containers
try:
    from dotenv import load_dotenv  # type: ignore
    load_dotenv()
except Exception:
    pass

from langchain_openai import ChatOpenAI
try:
    from langchain_core.messages import HumanMessage
except Exception:  # Fallback for older langchain
    from langchain.schema import HumanMessage  # type: ignore

# -------------------- App --------------------
app = FastAPI()

@app.get("/")
def root():
    return {
        "service": "Receipt Scanner API (vision-only)",
        "status": "ok",
        "docs": "/docs",
        "endpoints": [
            "/process-receipt-vision",
            "/healthz",
        ],
    }

# Ensure output folder exists
os.makedirs("receiptimages", exist_ok=True)

# -------------------- Models --------------------
# Keep your fast YOLOv8n weights
model = YOLO("receipt-scanner/models/best.pt")

# Vision LLM only (LangChain)
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    print("⚠️ OPENAI_API_KEY is not set; LangChain calls will fail until configured.")
VISION_LLM_MODEL = os.getenv("VISION_LLM_MODEL", "gpt-4o")
llm_vision = ChatOpenAI(model=VISION_LLM_MODEL, temperature=0)

# -------------------- Supabase Config --------------------
SUPABASE_URL = os.getenv("SUPABASE_URL", "").rstrip("/")
SUPABASE_SERVICE_ROLE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY", "")
SUPABASE_STORAGE_BUCKET = os.getenv("SUPABASE_STORAGE_BUCKET", "receipts")
SUPABASE_ANON_KEY = os.getenv("SUPABASE_ANON_KEY", "")
SUPABASE_RECEIPTS_TABLE = os.getenv("SUPABASE_RECEIPTS_TABLE", "receipts")
SUPABASE_CATEGORIES_TABLE = os.getenv("SUPABASE_CATEGORIES_TABLE", "categories")

# Direct DB (Pooler strongly recommended)
SUPABASE_DB_HOST = os.getenv("SUPABASE_DB_HOST")
SUPABASE_DB_PORT = int(os.getenv("SUPABASE_DB_PORT", "5432"))  # 5432 Pooler
SUPABASE_DB_NAME = os.getenv("SUPABASE_DB_NAME", "postgres")
SUPABASE_DB_USER = os.getenv("SUPABASE_DB_USER", "postgres")
SUPABASE_DB_PASSWORD = os.getenv("SUPABASE_DB_PASSWORD")

# -------------------- Prompt --------------------
UNIFIED_PROMPT_RULES = """
You are an expert at analyzing receipt text visible in the image. Focus on useful parts.

Rules:
1. **Venue**:
   - Retailer/restaurant name (top of receipt). Ignore bank/processor brands unless they're the only candidate.
   - If brand + branch, include both (e.g., "Checkers - Brooklyn Mall"). If unsure, return null.
2. **Date**:
   - Accept formats like DD-MM-YY, YYYY-MM-DD, DD/MM/YYYY (with optional time). If only time found, return null.
3. **Total**:
   - Final payable amount; prefer lines with TOTAL/Amount Due/Grand Total. Include currency if present (e.g., R238.00).
4. **Category** (choose only from):
   - "Office Supplies", "Travel/Transport", "Utilities", "Wages/Staff costs", "Meals & Entertainment", "Other Business Expenses".
   - Petrol/fuel clues (pump no./price per litre) ⇒ "Travel/Transport".
   - If unclear, default to "Other Business Expenses".
5. **Output format** (valid JSON only):
   {
     "venue": "...",
     "date": "...",
     "total": "...",
     "category": "..."
   }
"""

# -------------------- Small utilities --------------------
def category_to_slug(category: str | None) -> str:
    if not category:
        return "uncategorized"
    return (
        category.strip().lower()
        .replace("&", "and")
        .replace("/", "_")
        .replace(" ", "_")
        .replace("-", "_")
    )

def parse_extracted_json(raw_text: str) -> dict:
    if not raw_text:
        return {}
    text = str(raw_text).strip()

    # 1) Try largest JSON object between first '{' and last '}'
    s, e = text.find("{"), text.rfind("}")
    if s != -1 and e != -1 and e > s:
        candidate = text[s:e + 1]
        try:
            return json.loads(candidate)
        except Exception:
            pass

    # 2) Handle fenced code blocks ```json ... ```
    if text.startswith("```") and text.endswith("```"):
        lines = text.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].startswith("```"):
            lines = lines[:-1]
        inner = "\n".join(lines).strip()
        try:
            return json.loads(inner)
        except Exception:
            s2, e2 = inner.find("{"), inner.rfind("}")
            if s2 != -1 and e2 != -1 and e2 > s2:
                try:
                    return json.loads(inner[s2:e2 + 1])
                except Exception:
                    return {}
            return {}
    return {}

def parse_total(value: str | float | int | None) -> float:
    if value is None:
        return 0.0
    if isinstance(value, (int, float)):
        return float(value)
    import re
    m = re.findall(r"[-+]?[0-9]*\.?[0-9]+", str(value).replace(",", ""))
    return float(m[-1]) if m else 0.0

def ensure_extraction_complete(extracted_obj: dict):
    from fastapi import HTTPException
    required = ["venue", "total", "date"]
    missing = [k for k in required if not extracted_obj.get(k)]
    if missing:
        raise HTTPException(status_code=422, detail=f"Extraction incomplete. Missing: {', '.join(missing)}")

# Resize crop to keep LLM payload small but legible
def resize_if_needed(img: np.ndarray, max_dim: int = 1600) -> np.ndarray:
    h, w = img.shape[:2]
    m = max(h, w)
    if m <= max_dim:
        return img
    scale = max_dim / m
    new_w, new_h = int(w * scale), int(h * scale)
    return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)

# -------------------- Detection + encoding --------------------
def detect_and_crop_largest_receipt(img_array: np.ndarray):
    """Pick the highest-confidence YOLO box; fallback to full image if none."""
    results = model(img_array)
    cropped = img_array
    bbox = None
    for r in results:
        if len(r.boxes.xyxy) == 0:
            continue
        boxes = r.boxes.xyxy.cpu().numpy()
        scores = r.boxes.conf.cpu().numpy()
        i = int(scores.argmax())
        x1, y1, x2, y2 = map(int, boxes[i].tolist())
        cropped = img_array[y1:y2, x1:x2]
        bbox = {"x1": x1, "y1": y1, "x2": x2, "y2": y2, "conf": float(scores[i])}
        break
    return cropped, bbox

async def detect_and_crop_async(img_array: np.ndarray):
    async with app.state.yolo_sem:
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(app.state.executor, detect_and_crop_largest_receipt, img_array)

def encode_image_to_data_url(img_array: np.ndarray, format: str = "JPEG") -> str:
    pil_img = Image.fromarray(img_array)
    buffer = io.BytesIO()
    pil_img.save(buffer, format=format)
    b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
    mime = "image/jpeg" if format.upper() == "JPEG" else "image/png"
    return f"data:{mime};base64,{b64}"

def encode_image_to_jpeg_bytes(img_array: np.ndarray) -> bytes:
    pil_img = Image.fromarray(img_array)
    buffer = io.BytesIO()
    pil_img.save(buffer, format="JPEG")
    return buffer.getvalue()

# -------------------- Auth (no JWKS; Supabase Auth API only, with tiny TTL cache) --------------------
_auth_cache: dict[str, tuple[str, float]] = {}  # token -> (sub, expires_at)
_AUTH_TTL = float(os.getenv("AUTH_CACHE_TTL_SECONDS", "45"))

async def verify_user_async(authorization: str | None, x_user_id: str | None) -> str:
    if not authorization or not authorization.lower().startswith("bearer "):
        raise Exception("Missing bearer token")
    token = authorization.split(" ", 1)[1]
    now = time.time()
    cached = _auth_cache.get(token)
    if cached and cached[1] > now:
        sub = cached[0]
        if not sub or (x_user_id and x_user_id != sub):
            raise Exception("User mismatch")
        return sub
    if not SUPABASE_URL or not SUPABASE_ANON_KEY:
        raise Exception("Supabase auth not configured")
    url = f"{SUPABASE_URL}/auth/v1/user"
    headers = {"Authorization": f"Bearer {token}", "apikey": SUPABASE_ANON_KEY}
    resp = await app.state.http.get(url, headers=headers)
    if resp.status_code != 200:
        raise Exception("Invalid JWT")
    data = resp.json() or {}
    sub = data.get("id") or data.get("sub")
    if not sub or (x_user_id and x_user_id != sub):
        raise Exception("User mismatch")
    _auth_cache[token] = (str(sub), now + _AUTH_TTL)
    return str(sub)

# -------------------- Async Storage --------------------
async def upload_to_storage_async(path: str, content: bytes, content_type: str | None) -> str:
    if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
        raise Exception("Supabase storage not configured")
    url = f"{SUPABASE_URL}/storage/v1/object/{SUPABASE_STORAGE_BUCKET}/{path}"
    headers = {
        "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
        "apikey": SUPABASE_SERVICE_ROLE_KEY,
        "Content-Type": content_type or "application/octet-stream",
    }
    backoff = 0.2
    for _ in range(3):
        resp = await app.state.http.post(url, content=content, headers=headers)
        if resp.status_code < 500:
            if resp.status_code >= 400:
                raise Exception(f"Storage upload failed: {resp.text}")
            return f"{SUPABASE_STORAGE_BUCKET}/{path}"
        await asyncio.sleep(backoff)
        backoff *= 2
    raise Exception("Storage upload failed after retries")

# -------------------- REST fallbacks for DB --------------------
async def insert_receipt_rest(row: dict) -> dict:
    if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
        raise Exception("Supabase REST not configured")
    url = f"{SUPABASE_URL}/rest/v1/{SUPABASE_RECEIPTS_TABLE}"
    headers = {
        "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
        "apikey": SUPABASE_SERVICE_ROLE_KEY,
        "Content-Type": "application/json",
        "Prefer": "return=representation",
    }
    resp = await app.state.http.post(url, json=row, headers=headers)
    if resp.status_code >= 400:
        raise Exception(f"DB REST insert failed: {resp.text}")
    data = resp.json()
    return data[0] if isinstance(data, list) and data else data

async def upsert_category_rest(category_name: str | None):
    if not category_name:
        return
    if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
        raise Exception("Supabase REST not configured")
    url = f"{SUPABASE_URL}/rest/v1/{SUPABASE_CATEGORIES_TABLE}?on_conflict=name"
    headers = {
        "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
        "apikey": SUPABASE_SERVICE_ROLE_KEY,
        "Content-Type": "application/json",
        "Prefer": "resolution=merge-duplicates",
    }
    # Ignore non-2xx except 409; just log
    resp = await app.state.http.post(url, json={"name": category_name}, headers=headers)
    if resp.status_code not in (200, 201, 204, 409):
        print(f"⚠️ Category upsert (REST) failed: {resp.status_code} {resp.text}")

# -------------------- Async DB (asyncpg) --------------------
async def insert_receipt_db(row: dict) -> dict:
    # If pool unavailable, use REST fallback
    if getattr(app.state, "dbpool", None) is None:
        return await insert_receipt_rest(row)

    sql = f"""
    INSERT INTO {SUPABASE_RECEIPTS_TABLE}
    ( "userId", venue, category, total, purchased_at, status, storage_path, mime_type, file_size_bytes )
    VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)
    RETURNING id, "userId", venue, category, total, purchased_at, status, storage_path, mime_type, file_size_bytes;
    """
    async with app.state.dbpool.acquire() as conn:
        rec = await conn.fetchrow(
            sql,
            row["userId"], row["venue"], row["category"], row["total"],
            row["purchased_at"], row["status"], row["storage_path"],
            row["mime_type"], row["file_size_bytes"]
        )
    return dict(rec)

async def upsert_category_if_needed_db(category_name: str | None):
    if not category_name:
        return
    # If pool unavailable, REST fallback
    if getattr(app.state, "dbpool", None) is None:
        await upsert_category_rest(category_name)
        return
    sql = f'INSERT INTO {SUPABASE_CATEGORIES_TABLE} (name) VALUES ($1) ON CONFLICT (name) DO NOTHING;'
    async with app.state.dbpool.acquire() as conn:
        await conn.execute(sql, category_name)

# -------------------- Lifecycle --------------------
@app.on_event("startup")
async def on_startup():
    # Validate critical env
    if not (SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY):
        print("⚠️ Missing Storage/URL env; check SUPABASE_* variables.")
    if not (SUPABASE_DB_HOST and SUPABASE_DB_USER and SUPABASE_DB_PASSWORD):
        print("⚠️ Missing DB env; direct Postgres will be disabled unless set.")

    # Shared async HTTP client (Auth + Storage + REST fallback)
    app.state.http = httpx.AsyncClient(
        timeout=httpx.Timeout(60.0, connect=10.0),
        limits=httpx.Limits(
            max_connections=int(os.getenv("HTTP_MAX_CONNECTIONS", "50")),
            max_keepalive_connections=20,
        ),
        http2=True,
    )

    # Concurrency guards
    yolo_conc = int(os.getenv("YOLO_CONCURRENCY", "1"))
    app.state.yolo_sem = asyncio.Semaphore(yolo_conc)
    app.state.executor = ThreadPoolExecutor(max_workers=int(os.getenv("THREADPOOL_MAX", str(yolo_conc))))

   # --- Direct Postgres pool (with resilient startup) ---
    app.state.dbpool = None
    POOL_DB = (os.getenv("SUPABASE_DB_NAME", "postgres").strip() or "postgres")
    disable_stmt_cache = POOL_DB.endswith("-Transaction")
    
    # Choose CA source: inline PEM > file path > certifi bundle
    cafile_env = os.getenv("SUPABASE_DB_CA_PEM_PATH")     # e.g. /app/certs/supabase-ca.pem
    cadata_env = os.getenv("SUPABASE_DB_CA_PEM")          # full PEM content pasted as a secret
    if cadata_env:
        SSL_CTX = _ssl.create_default_context()
        # NOTE: you must paste a CA (or CA chain), not the leaf server cert
        SSL_CTX.load_verify_locations(cadata=cadata_env)
        ca_used = "env:SUPABASE_DB_CA_PEM"
    elif cafile_env:
        SSL_CTX = _ssl.create_default_context(cafile=cafile_env)
        ca_used = f"file:{cafile_env}"
    else:
        SSL_CTX = _ssl.create_default_context(cafile=certifi.where())
        ca_used = f"certifi:{certifi.where()}"
    
    SSL_CTX.check_hostname = True
    SSL_CTX.verify_mode = _ssl.CERT_REQUIRED
    
    # Optional last-resort switch (debug only; don't use in prod)
    if os.getenv("ALLOW_INSECURE_DB_SSL") == "1":
        SSL_CTX = _ssl.create_default_context()
        SSL_CTX.check_hostname = False
        SSL_CTX.verify_mode = _ssl.CERT_NONE
        ca_used = "INSECURE (verification disabled)"
    
    try:
        app.state.dbpool = await asyncpg.create_pool(
            user=SUPABASE_DB_USER,
            password=SUPABASE_DB_PASSWORD,
            database=POOL_DB,
            host=SUPABASE_DB_HOST,
            port=SUPABASE_DB_PORT,
            min_size=1,
            max_size=int(os.getenv("DB_POOL_MAX", "10")),
            command_timeout=15,
            max_inactive_connection_lifetime=60,
            ssl=SSL_CTX,
            statement_cache_size=0 if disable_stmt_cache else 1000,
        )
        print(f"✅ Direct Postgres pool ready via {SUPABASE_DB_HOST}:{SUPABASE_DB_PORT} ({POOL_DB}) using CA: {ca_used}")
    except Exception as e:
        app.state.dbpool = None
        print(f"⚠️ Direct Postgres pool disabled (startup error): {e!r}")



@app.on_event("shutdown")
async def on_shutdown():
    try:
        await app.state.http.aclose()
    except Exception:
        pass
    try:
        if getattr(app.state, "dbpool", None) is not None:
            await app.state.dbpool.close()
    except Exception:
        pass
    try:
        app.state.executor.shutdown(wait=False, cancel_futures=True)
    except Exception:
        pass

# -------------------- Health --------------------
@app.get("/healthz")
async def healthz():
    return {"ok": True, "db_direct": getattr(app.state, "dbpool", None) is not None}

# -------------------- Vision-only endpoint --------------------
@app.post("/process-receipt-vision")
async def process_receipt_vision(
    file: UploadFile = File(...),
    authorization: str | None = Header(default=None, alias="Authorization"),
    x_user_id: str | None = Header(default=None, alias="X-User-Id"),
):
    """YOLOv8n crop (highest-confidence) → GPT-4o (LangChain) → validate → Storage → Postgres (single INSERT)."""
    t0 = time.perf_counter()

    # 1) Verify user (Auth API; tiny TTL cache)
    try:
        user_id = await verify_user_async(authorization, x_user_id)
    except Exception as exc:
        from fastapi import HTTPException
        raise HTTPException(status_code=401, detail=str(exc))

    # 2) Read + YOLO crop
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert("RGB")
    img_array = np.array(image)
    cropped, bbox = await detect_and_crop_async(img_array)
    if bbox:
        print(f"🔎 YOLO bbox: {bbox}")

    # 3) Optional downscale for faster LLM
    cropped = resize_if_needed(cropped, max_dim=1600)

    # 4) Vision LLM
    data_url = encode_image_to_data_url(cropped, format="JPEG")
    message = HumanMessage(content=[
        {"type": "text", "text": UNIFIED_PROMPT_RULES},
        {"type": "image_url", "image_url": {"url": data_url}},
    ])
    try:
        if hasattr(llm_vision, "ainvoke"):
            result = await llm_vision.ainvoke([message])
        else:
            loop = asyncio.get_running_loop()
            result = await loop.run_in_executor(app.state.executor, llm_vision.invoke, [message])
    except Exception as exc:
        from fastapi import HTTPException
        raise HTTPException(status_code=502, detail=f"Vision LLM failed: {exc}")

    content_str = getattr(result, "content", str(result))
    extracted_obj = parse_extracted_json(content_str)
    ensure_extraction_complete(extracted_obj)

    # 5) Upload cropped image to Storage
    cropped_bytes = encode_image_to_jpeg_bytes(cropped)
    category_slug = category_to_slug(extracted_obj.get("category"))
    storage_path = f"{user_id}/{category_slug}/{int(time.time()*1000)}_cropped.jpg"
    try:
        stored_key = await upload_to_storage_async(storage_path, cropped_bytes, "image/jpeg")
    except Exception as exc:
        from fastapi import HTTPException
        raise HTTPException(status_code=502, detail=f"Storage upload failed: {exc}")

    # 6) Insert final row (single roundtrip; falls back to REST if pool down)
    row = {
        "userId": user_id,
        "venue": extracted_obj.get("venue"),
        "category": extracted_obj.get("category"),
        "total": parse_total(extracted_obj.get("total")),
        "purchased_at": extracted_obj.get("date") or datetime.now(timezone.utc).isoformat(),
        "status": "pending",
        "storage_path": stored_key,
        "mime_type": "image/jpeg",
        "file_size_bytes": len(cropped_bytes),
    }
    try:
        saved = await insert_receipt_db(row)
    except Exception as exc:
        from fastapi import HTTPException
        raise HTTPException(status_code=502, detail=f"DB insert failed: {exc}")

    # 7) Upsert category (best-effort; REST fallback inside)
    try:
        await upsert_category_if_needed_db(extracted_obj.get("category"))
    except Exception as exc:
        print(f"⚠️ Category upsert error: {exc}")

    t1 = time.perf_counter()
    print(f"✅ /process-receipt-vision completed in {t1 - t0:.2f}s")

    return {"receiptId": saved.get("id"), "status": saved.get("status", "pending")}


if __name__ == "__main__":
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860,
        workers=int(os.getenv("UVICORN_WORKERS", "2")),
    )