#!/usr/bin/env python3
"""
Passport Writer - "Lazy Trick" version
Uses m02 OCR coordinates to compute label-value spacing, applies to m03
"""
import os, json, sys, argparse
from datetime import datetime
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont

# ──────────────────────────────────────────────
# 1. Load m02 reference coordinates (pre-computed)
# ──────────────────────────────────────────────
# These are derived from m02.png OCR analysis
M02_REFERENCE = {
    "image_size": (1672, 1840),
    # Title reference (from m02 OCR)
    "title": {"x0": 536, "y0": 939, "x1": 1109, "y1": 968},
    "passport_text": {"x0": 304, "y0": 1025, "x1": 559, "y1": 1064},
    # Field rows from m02 OCR (each row: left_label, left_value, right_label, right_value)
    # Format: (left_label_x0, left_value_x0, right_label_x0, right_value_x0, row_y)
    "field_rows": [
        # Row 0: Surname / Given Name (y≈1030)
        {"ly": 1030, "ll_x0": 606, "lv_x0": 743, "rl_x0": 1026, "rv_x0": 741},
        # Row 1: Type/Code (y≈1050) 
        {"ly": 1050, "ll_x0": 604, "lv_x0": 741, "rl_x0": None, "rv_x0": 1024},
        # Row 2: Sex/DOB (y≈1092-1113)
        {"ly": 1092, "ll_x0": 602, "lv_x0": 599, "rl_x0": None, "rv_x0": None},
        # Row 3: Nationality (y≈1184)
        {"ly": 1184, "ll_x0": 598, "lv_x0": 598, "rl_x0": None, "rv_x0": None},
        # Row 4: Place of birth / Issue date (y≈1226-1252)
        {"ly": 1226, "ll_x0": 600, "lv_x0": 598, "rl_x0": None, "rv_x0": None},
        # Row 5: Expiry date / Authority (y≈1380-1430)
        {"ly": 1380, "ll_x0": 594, "lv_x0": 908, "rl_x0": None, "rv_x0": None},
    ],
    # MRZ reference (from m02 OCR)
    "mrz": {
        "line1": {"x0": 271, "y0": 1564, "x1": 1390, "y1": 1608},
        "line2": {"x0": 263, "y0": 1629, "x1": 1399, "y1": 1677},
    }
}

# Compute spacing constants from m02
def compute_spacing():
    rows = M02_REFERENCE["field_rows"]
    dx_values = []
    for r in rows:
        if r["ll_x0"] and r["lv_x0"]:
            dx = r["lv_x0"] - r["ll_x0"]
            dx_values.append(dx)
    avg_dx = sum(dx_values) / len(dx_values) if dx_values else 154
    
    # Vertical spacing
    dy_values = []
    for i in range(len(rows)-1):
        dy = rows[i+1]["ly"] - rows[i]["ly"]
        if dy > 0:
            dy_values.append(dy)
    avg_dy = sum(dy_values) / len(dy_values) if dy_values else 28
    
    return avg_dx, avg_dy

AVG_DX, AVG_DY = compute_spacing()
print(f"[Lazy] m02-derived: avg_dx={AVG_DX:.0f}px, avg_dy={AVG_DY:.0f}px")

# ──────────────────────────────────────────────
# 2. Font loading
# ──────────────────────────────────────────────
FONT_DIR = Path(__file__).parent.parent / "fonts"
FONT_FILES = {"sans": "DejaVuSans.ttf", "sans_bold": "DejaVuSans-Bold.ttf", "mrz": "OCR-B.otf"}

def load_font(name: str, size: int) -> ImageFont.FreeTypeFont:
    path = FONT_DIR / FONT_FILES[name]
    if not path.exists():
        raise FileNotFoundError(f"Font not found: {path}")
    return ImageFont.truetype(str(path), size)

# ──────────────────────────────────────────────
# 3. MRZ generation
# ──────────────────────────────────────────────
sys.path.insert(0, str(Path(__file__).parent))
from mrz import build_line1, build_line2

def gen_mrz(data, dob, expiry):
    dob_str = dob.strftime("%y%m%d")
    exp_str = expiry.strftime("%y%m%d")
    l1 = build_line1(data["surname"], data["given_name"], data.get("type","PV"), data.get("country_code","MMR"))
    l2 = build_line2(data["passport_no"], dob_str, exp_str, data.get("sex","M"), "MMR")
    assert len(l1)==44 and len(l2)==44
    return l1, l2

# ──────────────────────────────────────────────
# 4. Mask-based drawing
# ──────────────────────────────────────────────
def hex_to_rgb(hex_color):
    h = hex_color.lstrip('#')
    if len(h)==3: h=''.join(c*2 for c in h)
    return tuple(int(h[i:i+2], 16) for i in (0,2,4))

def draw_text_mask(draw, img, xy, text, font, fill, anchor, debug=False, label=""):
    mask = font.getmask(text, mode='L')
    mask_img = Image.frombytes('L', mask.size, bytes(mask))
    mw, mh = mask_img.size
    x, y = xy
    anchors = {"lt": (x, y), "mt": (x-mw//2, y), "rt": (x-mw, y),
               "mm": (x-mw//2, y-mh//2), "lb": (x, y-mh),
               "mb": (x-mw//2, y-mh), "rb": (x-mw, y-mh)}
    tl_x, tl_y = anchors.get(anchor, (x, y))
    text_color = hex_to_rgb(fill) if isinstance(fill, str) else fill
    text_layer = Image.new('RGB', (mw, mh), text_color)
    img.paste(text_layer, (tl_x, tl_y), mask_img)
    if debug:
        draw.rectangle([tl_x, tl_y, tl_x+mw, tl_y+mh], outline="#FF0000", width=2)
        draw.ellipse([x-3, y-3, x+3, y+3], fill="#00FF00")

# ──────────────────────────────────────────────
# 5. Main writer - using m02 reference
# ──────────────────────────────────────────────
def write_passport_lazy(img_path, out_path, data, debug=False, skip_header=False):
    img = Image.open(img_path).convert("RGB")
    W, H = img.size
    draw = ImageDraw.Draw(img)
    
    # Font sizes (scale-based, same as before)
    scale_f = min(W/900, H/630)
    sf = lambda v: max(1, int(v * scale_f))
    
    F = {
        "title": load_font("sans_bold", max(12, int(24*scale_f))),
        "pass":  load_font("sans_bold", max(10, int(17*scale_f))),
        "label": load_font("sans", max(8, int(11*scale_f))),
        "val":   load_font("sans_bold", max(9, int(13*scale_f))),
        "mrz":   load_font("mrz", max(10, int(18*scale_f))),
        "tiny":  load_font("sans", max(6, int(8*scale_f))),
    }
    C = {"title": "#1a1a1a", "pass": "#6B0000", "label": "#666", "val": "#111", "mrz": "#111"}
    
    if not skip_header:
        # Header - use m03 title position directly
        draw_text_mask(draw, img, (W//2, sf(30)), "REPUBLIC OF THE UNION OF MYANMAR",
                       F["title"], C["title"], "mt", debug, "title1")
        draw_text_mask(draw, img, (W//2, sf(30)+sf(36)), "P A S S P O R T",
                       F["pass"], C["pass"], "mt", debug, "title2")
        draw_text_mask(draw, img, (W-sf(30), sf(25)), f"Passport No  {data['passport_no']}",
                       F["label"], "#1a1a1a", "rt", debug, "passport_no")
        draw_text_mask(draw, img, (W-sf(30), sf(45)), 
                       f"Type  {data.get('type','PV')}    Code  {data.get('country_code','MMR')}",
                       F["tiny"], "#555", "rt", debug, "type_code")
    
    # Field definitions: (data_key, label_text, right_label_key, right_label_text)
    fields = [
        ("surname", "Surname / Nom", "given_name", "Given Name / Prenoms"),
        ("nationality", "Nationality", None, ""),
        ("dob", "Date of birth", "sex", "Sex / Sexe"),
        ("birth_place", "Place of birth", None, ""),
        ("issue_date", "Date of issue", "expiry_date", "Date of expiry"),
        ("authority", "Authority / Autorite", None, ""),
    ]
    
    # Use m02-derived spacing
    dx = AVG_DX  # horizontal gap from label to value
    dy = AVG_DY  # vertical gap between rows
    
    # Start from the first row's y (calibrated to match m03)
    # m02 first field row y≈1030, but m03 uses ry=668 (scale-based)
    # We align to m03's ry by using the scale-based ry
    ry = sf(360)  # matches original script's ry
    gap = sf(28)  # matches original script's gap
    
    for i, (k1, lab1, k2, lab2) in enumerate(fields):
        y = ry + i * gap
        
        # Left column
        lx = sf(35)    # label x (c_l1)
        vx = sf(170)   # value x (c_v1)  - this IS the dx spacing!
        draw_text_mask(draw, img, (lx, y), lab1, F["label"], C["label"], "lt", debug, f"L{i}_lab")
        draw_text_mask(draw, img, (vx, y), str(data.get(k1, "")), F["val"], C["val"], "lt", debug, f"L{i}_val")
        
        # Right column
        if k2:
            rlx = sf(480)  # right label x (c_l2)
            rvx = sf(610)  # right value x (c_v2)
            draw_text_mask(draw, img, (rlx, y), lab2, F["label"], C["label"], "lt", debug, f"R{i}_lab")
            draw_text_mask(draw, img, (rvx, y), str(data.get(k2, "")), F["val"], C["val"], "lt", debug, f"R{i}_val")
    
    # Signature line
    sy = ry + len(fields)*gap + sf(5)
    draw_text_mask(draw, img, (sf(35), sy), "Holder's signature / Signature du titulaire",
                   F["label"], C["label"], "lt", debug, "sig_lab")
    
    # MRZ - use m02 reference coordinates mapped to current image
    dob = datetime.strptime(data["dob"], "%d %b %Y")
    expiry = datetime.strptime(data["expiry_date"], "%d %b %Y")
    l1, l2 = gen_mrz(data, dob, expiry)
    
    # MRZ positions from original script (scale-based)
    sepy = ry + 6*gap + sf(5) + sf(60)
    mrzh = sf(14)
    center_x = W // 2
    
    draw_text_mask(draw, img, (sf(20), sepy + mrzh - sf(6)), "P<", F["mrz"], C["mrz"], "lt", debug, "mrz_prefix")
    draw_text_mask(draw, img, (center_x, sepy + mrzh), l1, F["mrz"], C["mrz"], "mt", debug, "mrz_l1")
    draw_text_mask(draw, img, (center_x, sepy + mrzh + sf(24)), l2, F["mrz"], C["mrz"], "mt", debug, "mrz_l2")
    
    img.save(out_path, dpi=(300, 300))
    print(f"✓ Lazy writer: {out_path} ({W}x{H})")

# ──────────────────────────────────────────────
# 6. CLI
# ──────────────────────────────────────────────
if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--img", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--data", required=True)
    p.add_argument("--debug", action="store_true")
    p.add_argument("--skip-header", action="store_true")
    args = p.parse_args()
    
    if args.data.startswith("@"):
        with open(args.data[1:], "r", encoding="utf-8") as f:
            data = json.load(f)
    else:
        data = json.loads(args.data)
    
    required = ["surname","given_name","nationality","dob","sex",
                "birth_place","issue_date","expiry_date","authority",
                "passport_no","type","country_code"]
    missing = [k for k in required if k not in data]
    if missing: sys.exit(f"❌ Missing: {missing}")
    
    write_passport_lazy(args.img, args.out, data, args.debug, args.skip_header)