app/lens_text_core.py

import os, time, asyncio, base64, re, threading, hashlib, logging
from io import BytesIO
from typing import Any, Dict, List, Union
from urllib.parse import urlparse

import httpx
from PIL import Image

from seleniumbase import Driver

CHROME_BINARY_PATH = os.getenv("CHROME_BINARY_PATH", "").strip() 

LOGGER = logging.getLogger("lens_text_core")
if not LOGGER.handlers:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    )

COOKIE_JSON_URL = os.getenv("COOKIE_JSON_URL", "")
UA = "Mozilla/5.0 (Lens Text OCR)"

CHROME_EXTRA_ARGS = os.getenv(
    "CHROME_EXTRA_ARGS",
    "--disable-gpu --no-sandbox --disable-dev-shm-usage "
    "--window-size=1920,1080 --headless=new",
).split()

_CACHE_TTL    = 600  
_BROWSER_TTL  = 900  
_IDLE_TIMEOUT = int(os.getenv("CHROME_IDLE_SECONDS", "60"))

def _build_chrome(cookie_dict: Dict[str, str] | None = None):
    drv = Driver(
        browser="chrome",
        uc=True, 
        headless=True, 
        incognito=True
    )
    
    drv.get("https://google.com/favicon.ico")

    if cookie_dict:
        for name, val in cookie_dict.items():
            try:
                drv.add_cookie({
                    "name": name, "value": val,
                    "domain": ".google.com", "path": "/", "secure": True
                })
            except Exception:
                pass
    return drv

_cached_cookie, _cached_cookie_ts, _cookie_lock = None, 0.0, threading.Lock()
_global_driver, _driver_last_use, _driver_lock  = None, 0.0, threading.Lock()

def _grab_cookies_with_browser() -> Dict[str, Any]:
    drv = _build_chrome()
    try:
        drv.get("https://lens.google.com/")
        jar = {
            c["name"]: c["value"]
            for c in drv.get_cookies()
            if c.get("domain","").endswith(".google.com") or c.get("domain","").endswith("google.com")
        }
        return {"cookies": jar, "_source": "browser"}
    finally:
        try: drv.quit()
        except Exception: pass

async def _cookie_header() -> str:
    global _cached_cookie, _cached_cookie_ts
    now = time.time()

    with _cookie_lock:
        if _cached_cookie:
            ttl = _BROWSER_TTL if _cached_cookie.get("_source") == "browser" else _CACHE_TTL
            if (now - _cached_cookie_ts) < ttl:
                obj = _cached_cookie.get("cookies", _cached_cookie)
                return "; ".join(f"{k}={v}" for k, v in obj.items())

    if COOKIE_JSON_URL:
        try:
            async with httpx.AsyncClient(timeout=4) as cli:
                resp = await cli.get(COOKIE_JSON_URL)
                resp.raise_for_status()
                data = resp.json()
            with _cookie_lock:
                data["_source"] = "remote"
                _cached_cookie, _cached_cookie_ts = data, now
            obj = data.get("cookies", data)
            return "; ".join(f"{k}={v}" for k, v in obj.items())
        except Exception as e:
            LOGGER.warning("fetch COOKIE_JSON_URL failed: %s – fallback to browser", e)

    loop = asyncio.get_running_loop()
    data = await loop.run_in_executor(None, _grab_cookies_with_browser)

    with _cookie_lock:
        _cached_cookie, _cached_cookie_ts = data, now

    obj = data.get("cookies", data)
    return "; ".join(f"{k}={v}" for k, v in obj.items())

def _sap_header(cookie_hdr: str) -> Dict[str, str]:
    origin = "https://lens.google.com"
    sid = None
    for part in cookie_hdr.split("; "):
        if part.startswith("__Secure-3PAPISID=") or part.startswith("SAPISID="):
            sid = part.split("=", 1)[1]
            break
    if not sid:
        return {}
    ts  = int(time.time())
    sig = hashlib.sha1(f"{ts} {sid} {origin}".encode()).hexdigest()
    return {
        "X-Origin": origin,
        "X-Goog-AuthUser": "0",
        "Authorization": f"SAPISIDHASH {ts}_{sig}",
    }

def _is_alive(drv) -> bool:
    try:
        _ = drv.title
        return True
    except Exception:
        return False

def _ensure_driver(cookie_dict: Dict[str, str]):
    global _global_driver, _driver_last_use
    with _driver_lock:
        if _global_driver is None or not _is_alive(_global_driver):
            try:
                if _global_driver:
                    _global_driver.quit()
            except Exception:
                pass
            _global_driver = _build_chrome(cookie_dict)
        _driver_last_use = time.time()
        return _global_driver

def _driver_reaper():
    global _global_driver
    while True:
        time.sleep(1)
        if _global_driver and (time.time() - _driver_last_use) > _IDLE_TIMEOUT:
            LOGGER.info("♻️  quitting idle driver")
            try: _global_driver.quit()
            except Exception: pass
            _global_driver = None

_reaper_started = False
def _ensure_reaper_started():
    global _reaper_started
    if _reaper_started:
        return
    try:
        threading.Thread(target=_driver_reaper, daemon=True).start()
        _reaper_started = True
        LOGGER.debug("text driver reaper started")
    except Exception as e:
        LOGGER.warning("could not start text driver reaper: %s", e)

def _parse_calc_value(calc: str, dim: float) -> float:
    m = re.search(r"calc\(([\d.]+)%\s*([+-])\s*([\d.]+)px\)", calc)
    if not m: return 0.0
    pct, op, off = float(m[1]), m[2], float(m[3])
    base = dim * pct / 100.0
    return base - off if op == "-" else base + off

def _extract_boxes(drv, w: int, h: int) -> List[Dict[str, Any]]:
    drv.wait_for_element_visible("div.lv6PAb", timeout=30)
    nodes = drv.find_elements("xpath", "//div[contains(@class, 'lv6PAb') and @aria-label]")

    out: List[Dict[str,Any]] = []
    for n in nodes:
        if not (n.get_attribute("data-line-index") or "").strip():
            continue

        text  = (n.get_attribute("aria-label") or "").strip()
        style = n.get_attribute("style") or ""
        if not text or "calc(" not in style:
            continue

        kv = {k.strip(): v.strip()
              for k,v in (p.split(":",1) for p in style.split(";") if ":" in p)}

        top, left = _parse_calc_value(kv.get("top",""),   h), _parse_calc_value(kv.get("left",""),  w)
        wid, hei  = _parse_calc_value(kv.get("width",""), w), _parse_calc_value(kv.get("height",""),h)
        rot_m = re.search(r"rotate\(([-\d.]+)deg\)", style); rot = float(rot_m[1]) if rot_m else 0.0

        verts = [
            {"x": int(left),        "y": int(top)},
            {"x": int(left+wid),    "y": int(top)},
            {"x": int(left+wid),    "y": int(top+hei)},
            {"x": int(left),        "y": int(top+hei)},
        ]
        abs_style = f"top: {int(top)}px; left: {int(left)}px; width: {int(wid)}px; height: {int(hei)}px; transform: rotate({rot}deg);"

        out.append({
            "description": text,
            "boundingPoly": {"vertices": verts},
            "rotate": rot,
            "style": abs_style,

            "raw_style": style,
            "top_str": kv.get("top",""),
            "left_str": kv.get("left",""),
            "width_str": kv.get("width",""),
            "height_str": kv.get("height",""),
        })
    return out

def _merge_by_center_line(anns: List[Dict[str,Any]], m_x: int=10, m_y: int=15) -> List[Dict[str,Any]]:
    for a in anns:
        v  = a["boundingPoly"]["vertices"]
        xs, ys = [p["x"] for p in v], [p["y"] for p in v]
        a["_l"], a["_r"], a["_t"], a["_b"] = min(xs), max(xs), min(ys), max(ys)
        a["_cx"], a["_cy"] = (a["_l"]+a["_r"])/2, (a["_t"]+a["_b"])/2

    parent = list(range(len(anns)))
    def find(i):
        while parent[i]!=i:
            parent[i]=parent[parent[i]]
            i=parent[i]
        return i
    def union(i,j):
        ri,rj = find(i),find(j)
        if ri!=rj: parent[rj]=ri

    for i in range(len(anns)):
        for j in range(i+1,len(anns)):
            ai, aj = anns[i], anns[j]
            if (abs(ai["_cx"]-aj["_cx"]) < m_x and
                ai["_t"]-m_y < aj["_b"] and ai["_b"]+m_y > aj["_t"]):
                union(i,j)

    groups: Dict[int,List[Dict[str,Any]]] = {}
    for idx,a in enumerate(anns):
        groups.setdefault(find(idx), []).append(a)

    merged: List[Dict[str,Any]] = []
    for g in groups.values():
        if len(g)==1:
            a = g[0]
            merged.append({
                "description": a["description"],
                "boundingPoly": a["boundingPoly"],
                "rotate": a["rotate"],
                "style": a["style"],
            })
        else:
            txt = "\n".join(aa["description"] for aa in g)
            l,r = min(aa["_l"] for aa in g), max(aa["_r"] for aa in g)
            t,b = min(aa["_t"] for aa in g), max(aa["_b"] for aa in g)
            merged.append({
                "description": txt,
                "boundingPoly": {"vertices":[
                    {"x":l,"y":t}, {"x":r,"y":t}, {"x":r,"y":b}, {"x":l,"y":b}]},
                "rotate": 0.0,
                "style": f"top: {t}px; left: {l}px; width: {r-l}px; height: {b-t}px; transform: rotate(0deg);",
            })
    return merged

async def translate_lens_text(src: Union[str, bytes, BytesIO]) -> Dict[str,Any]:
    _ensure_reaper_started()
    if isinstance(src, (bytes,bytearray)):           img_bytes = bytes(src)
    elif isinstance(src, BytesIO):                   img_bytes = src.getvalue()
    elif isinstance(src, str):
        if src.startswith("data:"):                  img_bytes = base64.b64decode(src.split(",",1)[1])
        else:
            async with httpx.AsyncClient() as cli:
                o = urlparse(src)
                referer = f"{o.scheme}://{o.netloc}/" if o.scheme and o.netloc else None
                hdr_img = {"User-Agent": UA}
                if referer: hdr_img["Referer"] = referer
                try:
                    r = await cli.get(src, headers=hdr_img, timeout=10)
                    r.raise_for_status()
                    img_bytes = r.content
                except httpx.HTTPStatusError as he:
                    code = he.response.status_code if he.response is not None else "NA"
                    raise RuntimeError(f"fetch image HTTP {code}")
                except httpx.TimeoutException:
                    raise RuntimeError("fetch image TIMEOUT")
                except Exception as e:
                    raise RuntimeError(f"fetch image ERROR {type(e).__name__}")
    else: raise TypeError("unsupported src type")

    from io import BytesIO as _B
    with Image.open(_B(img_bytes)) as im:
        w, h = im.size

    ck = await _cookie_header()
    hdr = {"User-Agent": UA, "Cookie": ck, "Referer":"https://lens.google.com/", **_sap_header(ck)}
    async with httpx.AsyncClient(follow_redirects=False) as cli:
        up = await cli.post("https://lens.google.com/v3/upload",
                            files={ "encoded_image": ("file.jpg", img_bytes, "image/jpeg"),
                                    "sbisrc":(None,"browser"), "rt":(None,"j") },
                            headers=hdr, timeout=10)
    if up.status_code not in (302,303):
        raise RuntimeError(f"Lens upload failed: {up.status_code}")
    loc = up.headers.get("location") or ""
    if not loc: raise RuntimeError("no redirect location")

    cookie_dict = {k: v for k, v in (p.split("=", 1) for p in ck.split("; ") if "=" in p)}

    loop = asyncio.get_running_loop()
    drv = await loop.run_in_executor(None, lambda: _ensure_driver(cookie_dict))

    def _blocking() -> List[Dict[str, Any]]:
        nonlocal drv
        with _driver_lock:
            try:
                try:
                    drv.get(loc)
                except:
                    try:
                        drv.quit()
                    except Exception:
                        pass
                    drv = _ensure_driver(cookie_dict)
                    drv.get(loc)
                return _extract_boxes(drv, w, h)
            finally:
                pass

    raw = await loop.run_in_executor(None, _blocking)

    merged  = _merge_by_center_line(raw)
    fulltxt = " ".join(a["description"] for a in raw).strip()

    return {
        "textAnnotations":      merged,
        "rawTextAnnotations":   raw,
        "fullTextAnnotation":   {"text": fulltxt},
        "loc":                  loc,
    }

async def prewarm_driver():
    try:
        cookie_hdr = await _cookie_header()
        cookie_dict = {k: v for k, v in (p.split("=", 1) for p in cookie_hdr.split("; ") if "=" in p)}
        _ = cookie_dict
        LOGGER.info("prewarm_driver: cookies ready")
    except Exception as e:
        LOGGER.warning("prewarm_driver failed: %s", e)