backend/chart_reader.py

# ─────────────────────────────────────────────────────────────────────────────
# PakFit Chart Reader v2
# Uses Gemini Vision to read size chart images and extract measurements
# Updated to use google.genai (new library) with gemini-1.5-flash
# ─────────────────────────────────────────────────────────────────────────────

import os
import json
import base64
import requests
from pathlib import Path
from io import BytesIO
from PIL import Image
from dotenv import load_dotenv

# Load .env from project root
load_dotenv(Path(__file__).parent.parent / ".env")

# ── Configure Gemini ──────────────────────────────────────────────────────────
GEMINI_KEY = os.getenv("GEMINI_API_KEY")

try:
    from google import genai
    from google.genai import types
    if GEMINI_KEY:
        client = genai.Client(api_key=GEMINI_KEY)
        print("✓ Gemini Vision ready")
    else:
        client = None
        print("⚠ GEMINI_API_KEY not found in .env")
except ImportError:
    client = None
    print("⚠ google-genai not installed. Run: pip install google-genai")

GEMINI_MODEL = "gemini-2.0-flash"

# ── Headers for image download ────────────────────────────────────────────────
HEADERS = {
    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
                  "AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/120.0.0.0 Safari/537.36",
}

# ── Gemini prompt ─────────────────────────────────────────────────────────────
EXTRACTION_PROMPT = """
This is a Pakistani clothing size chart image.

Extract ALL size data from this chart and return ONLY valid JSON.
Do not include any explanation, markdown, or code blocks — just raw JSON.

Rules:
1. Convert any centimetre values to inches (divide by 2.54)
2. Use these exact field names: size_label, chest, shoulder, sleeve, collar, length, waist
3. If a measurement is not in the chart, omit that field
4. size_label should be: XS, S, M, L, XL, XXL
5. All numeric values should be floats rounded to 2 decimal places

Return format:
{"garment_type": "Kameez", "sizes": [{"size_label": "S", "chest": 23.0, "shoulder": 17.5, "sleeve": 23.5, "collar": 15.0, "length": 40.75}]}

Common field mappings:
- Ready Chest or Chest @ Arm Hole = chest
- Sleeves Length or Sleeve Length = sleeve
- Band Collar or Collar = collar
- Front Length or Length from HSP = length
- Trouser Waist or Bottom or Width = waist
"""


def download_image(image_url):
    """Download image from URL and return PIL Image."""
    try:
        r = requests.get(image_url, headers=HEADERS, timeout=15)
        r.raise_for_status()
        img = Image.open(BytesIO(r.content))
        if img.mode not in ("RGB", "L"):
            img = img.convert("RGB")
        print(f"  Image downloaded: {img.size[0]}x{img.size[1]} pixels")
        return img
    except Exception as e:
        print(f"  Image download failed: {e}")
        return None


def clean_gemini_response(text):
    """Remove markdown code blocks from Gemini response."""
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        lines = [l for l in lines if not l.startswith("```")]
        text = "\n".join(lines).strip()
    if text.startswith("json"):
        text = text[4:].strip()
    return text


def image_to_base64(img):
    """Convert PIL Image to base64 JPEG string."""
    buf = BytesIO()
    img.save(buf, format="JPEG", quality=85)
    return base64.b64encode(buf.getvalue()).decode()


def read_chart_from_image_url(image_url, brand=None, garment_type=None):
    """
    Download a size chart image and extract measurements using Gemini Vision.
    Returns dict with garment_type, sizes list, and source.
    """
    if not client:
        print("  Gemini client not available")
        return None

    print("  Reading chart image with Gemini Vision...")

    img = download_image(image_url)
    if not img:
        return None

    try:
        prompt = EXTRACTION_PROMPT
        if garment_type:
            prompt += f"\n\nNote: This is a {garment_type} size chart."
        if brand:
            prompt += f"\nBrand: {brand}"

        img_b64   = image_to_base64(img)
        img_bytes = base64.b64decode(img_b64)

        response = client.models.generate_content(
            model=GEMINI_MODEL,
            contents=[
                types.Part.from_bytes(data=img_bytes, mime_type="image/jpeg"),
                prompt,
            ]
        )

        raw_text   = response.text
        clean_text = clean_gemini_response(raw_text)
        data       = json.loads(clean_text)
        sizes      = data.get("sizes", [])
        detected   = data.get("garment_type", garment_type or "Kameez")

        print(f"  Gemini extracted {len(sizes)} sizes for {detected}")

        valid_sizes = []
        for size in sizes:
            if "size_label" not in size:
                continue
            cleaned = {"size_label": str(size["size_label"])}
            for field in ["chest", "shoulder", "sleeve", "collar", "length", "waist"]:
                if field in size and size[field]:
                    try:
                        val = float(size[field])
                        if field == "chest" and (val < 15 or val > 60):
                            continue
                        if field == "length" and (val < 20 or val > 60):
                            continue
                        cleaned[field] = round(val, 2)
                    except (ValueError, TypeError):
                        pass
            valid_sizes.append(cleaned)

        if not valid_sizes:
            print("  No valid sizes extracted")
            return None

        return {
            "garment_type": detected,
            "sizes":        valid_sizes,
            "source":       "gemini_vision",
        }

    except json.JSONDecodeError as e:
        print(f"  JSON parse error: {e}")
        return None
    except Exception as e:
        print(f"  Gemini Vision error: {e}")
        return None


def read_chart_from_table(table_data, brand=None, garment_type=None):
    """Parse size chart from HTML table data using Gemini."""
    if not client:
        return None

    try:
        headers    = table_data.get("headers", [])
        rows       = table_data.get("rows", [])
        if not rows:
            return None

        table_text = "Headers: " + " | ".join(headers) + "\n"
        for row in rows:
            table_text += " | ".join(str(v) for v in row.values()) + "\n"

        prompt = f"""
This is a Pakistani clothing size chart as text.

{table_text}

Extract size data and return ONLY valid JSON (no markdown):
{{"garment_type": "Kameez", "sizes": [{{"size_label": "S", "chest": 23.0, "shoulder": 17.5}}]}}

Convert cm to inches if needed. Use: size_label, chest, shoulder, sleeve, collar, length, waist.
"""
        if brand:
            prompt += f"\nBrand: {brand}"
        if garment_type:
            prompt += f"\nGarment: {garment_type}"

        response   = client.models.generate_content(
            model=GEMINI_MODEL,
            contents=[prompt]
        )
        clean_text = clean_gemini_response(response.text)
        data       = json.loads(clean_text)
        sizes      = data.get("sizes", [])
        print(f"  Gemini extracted {len(sizes)} sizes from table")

        return {
            "garment_type": data.get("garment_type", garment_type or "Kameez"),
            "sizes":        sizes,
            "source":       "gemini_table",
        }

    except Exception as e:
        print(f"  Table reading error: {e}")
        return None


def get_bilingual_explanation(brand, garment_type, recommended_size,
                               fitscore, all_sizes, buyer_measurements):
    """Generate bilingual Urdu + English explanation using Gemini."""
    if not client:
        return {
            "english": f"Recommended size {recommended_size} at {brand} with {fitscore:.0f}% confidence.",
            "urdu":    f"{brand} میں تجویز کردہ سائز {recommended_size} ہے۔"
        }

    prompt = f"""
A Pakistani man is buying a {garment_type} from {brand}.
PakFit recommended size {recommended_size} with {fitscore:.0f}% confidence.
His measurements: {buyer_measurements}

Write a SHORT friendly explanation in English and Urdu (1-2 sentences each).
Return ONLY this JSON (no markdown):
{{"english": "English here", "urdu": "اردو یہاں"}}
"""

    try:
        response   = client.models.generate_content(
            model=GEMINI_MODEL,
            contents=[prompt]
        )
        clean_text = clean_gemini_response(response.text)
        return json.loads(clean_text)
    except Exception:
        return {
            "english": f"Based on your measurements, size {recommended_size} at {brand} is your best fit with {fitscore:.0f}% confidence.",
            "urdu":    f"آپ کی پیمائش کے مطابق، {brand} میں سائز {recommended_size} آپ کے لیے بہترین ہے۔"
        }


# ─────────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    print("Testing PakFit Chart Reader v2...")
    print("=" * 50)

    explanation = get_bilingual_explanation(
        brand="J.",
        garment_type="Kameez",
        recommended_size="M",
        fitscore=97.5,
        all_sizes=[{"size": "M", "score": 97.5}],
        buyer_measurements={"chest": 24, "shoulder": 18.5}
    )
    print(f"English: {explanation['english']}")
    print(f"Urdu: {explanation['urdu']}")
    print("\nChart reader ready.")