#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
整页滚动截图 (scroll & stitch)  ——  针对 2号人事部 window 整页滚动 + fixed 顶栏

用法:
  python3 capture_full.py <session> <out_path> [max_tiles]

原理:
  1. 隐藏帮助/反馈/AI助手/购买浮层等固定挂件, 避免遮挡
  2. 读取滚动指标, 探测顶部 fixed 顶栏高度
  3. 逐屏滚动截图(视口图), 除首屏外裁掉与固定顶栏重复的部分
  4. 纵向拼接为整页长图
依赖: PIL(Pillow), 通过 daemon(127.0.0.1:10086) 调用 evaluate / screenshot
"""
import sys, os, json, time, base64, io, urllib.request

DAEMON = "http://127.0.0.1:10086/command"
_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}))

def call(action, args, session):
    body = json.dumps({"action": action, "args": args, "session": session}).encode()
    req = urllib.request.Request(DAEMON, data=body, headers={"Content-Type": "application/json"})
    with _OPENER.open(req, timeout=40) as resp:
        return json.loads(resp.read().decode())

def ev(code, session):
    r = call("evaluate", {"code": code}, session)
    return r.get("data", {}).get("value")

def shot(session):
    r = call("screenshot", {"format": "jpeg", "quality": 70}, session)
    b64 = r.get("data", {}).get("data")
    if not b64:
        raise RuntimeError("no screenshot data: " + json.dumps(r)[:200])
    from PIL import Image
    return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")

HIDE_JS = r"""
(()=>{
  const sels = ['[class*="helpBox"]','[class*="aiAssistant"]','[class*="product_"]',
                '[class*="feedback"]','[class*="Feedback"]','[class*="floatBar"]',
                '[class*="sideBar"]','[class*="robot"]','[class*="advert"]'];
  let n=0;
  sels.forEach(s=>document.querySelectorAll(s).forEach(el=>{
     const r=el.getBoundingClientRect();
     if(r.width<600 && getComputedStyle(el).position==='fixed'){el.style.setProperty('display','none','important');n++;}
  }));
  // 反馈竖条
  document.querySelectorAll('div,a,button').forEach(el=>{
     if(el.children.length<=3 && /^反\s*馈$/.test((el.innerText||'').replace(/\s/g,'')) && getComputedStyle(el).position==='fixed'){el.style.setProperty('display','none','important');n++;}
  });
  return n;
})()
"""

METRICS_JS = r"""
(()=>{
  let hdr=0;
  const walk=(el,d)=>{ if(d>5)return; for(const c of el.children){ const s=getComputedStyle(c); const r=c.getBoundingClientRect();
     if((s.position==='fixed'||s.position==='sticky') && r.top<=1 && r.height>10 && r.height<200 && r.width>window.innerWidth*0.6){ if(r.height>hdr)hdr=r.height; }
     walk(c,d+1);} };
  walk(document.body,0);
  return JSON.stringify({winH:window.innerHeight, winW:window.innerWidth,
     docH:Math.max(document.documentElement.scrollHeight,document.body.scrollHeight), hdr:Math.round(hdr)});
})()
"""

def scroll_to(y, session):
    ev(f"(()=>{{window.scrollTo(0,{y});return window.scrollY;}})()", session)
    time.sleep(0.35)
    return ev("(()=>window.scrollY)()", session)

# 探测主内容区的内部滚动容器(表格/列表在固定高度面板内自行滚动的情况)
DETECT_SCROLLER_JS = r"""
(()=>{
  const W=innerWidth,H=innerHeight; let best=null,score=0;
  document.querySelectorAll('div,main,section,ul').forEach(el=>{
    const s=getComputedStyle(el); const oy=s.overflowY;
    if(oy!=='auto'&&oy!=='scroll') return;
    const extra=el.scrollHeight-el.clientHeight; if(extra<40) return;
    const r=el.getBoundingClientRect();
    if(r.width<W*0.4||r.height<H*0.25) return;
    if(r.bottom<H*0.4||r.top>H*0.85) return;
    const sc=r.width*r.height+extra*300;
    if(sc>score){score=sc;best=el;}
  });
  if(!best) return JSON.stringify({found:false});
  document.querySelectorAll('[data-capscroll]').forEach(e=>e.removeAttribute('data-capscroll'));
  best.setAttribute('data-capscroll','1');
  const r=best.getBoundingClientRect();
  let hs=0;
  best.querySelectorAll(':scope>*,:scope>*>*,:scope>*>*>*').forEach(c=>{
    const cs=getComputedStyle(c); const cr=c.getBoundingClientRect();
    if((cs.position==='sticky'||cs.position==='fixed')&&cr.top<=r.top+3&&cr.height>10&&cr.height<160&&cr.width>r.width*0.5){ if(cr.height>hs)hs=cr.height; }
  });
  return JSON.stringify({found:true,left:Math.round(r.left),top:Math.round(r.top),
     width:Math.round(r.width),height:Math.round(r.height),
     scrollHeight:best.scrollHeight,clientHeight:best.clientHeight,hdr:Math.round(hs)});
})()
"""

def scroll_inner(y, session):
    return ev("(()=>{const e=document.querySelector('[data-capscroll]');if(!e)return 0;e.scrollTop=%d;return Math.round(e.scrollTop);})()" % int(y), session)


def _inner_stitch(session, sc, scale, max_tiles):
    """对内部滚动容器逐屏滚动 -> 裁剪容器区域 -> 纵向拼接(含粘性表头去重)。"""
    from PIL import Image
    left, top = sc["left"], sc["top"]
    width, height = sc["width"], sc["height"]
    hs = sc.get("hdr", 0) or 0
    maxScroll = max(0, sc["scrollHeight"] - sc["clientHeight"])
    box = (int(round(left * scale)), int(round(top * scale)),
           int(round((left + width) * scale)), int(round((top + height) * scale)))

    def tile():
        return shot(session).crop(box)

    scroll_inner(0, session); time.sleep(0.35)
    first = tile()
    parts = [first]
    prev_bottom = height
    target = height - hs
    guard = 0
    while guard < max_tiles:
        guard += 1
        actualY = int(scroll_inner(target, session) or 0); time.sleep(0.3)
        img = tile()
        crop_logical = prev_bottom - actualY
        if crop_logical < 0:
            crop_logical = 0
        crop_px = int(round(crop_logical * scale))
        if crop_px >= img.height:
            crop_px = int(round(hs * scale))
        parts.append(img.crop((0, crop_px, img.width, img.height)))
        prev_bottom = actualY + height
        if actualY >= maxScroll - 1:
            break
        target = prev_bottom - hs
    total_h = sum(p.height for p in parts)
    stitched = Image.new("RGB", (first.width, total_h), (255, 255, 255))
    y = 0
    for p in parts:
        stitched.paste(p, (0, y)); y += p.height
    scroll_inner(0, session)
    return stitched, guard + 1

def capture(session, out_path, max_tiles=14):
    from PIL import Image

    ev(HIDE_JS, session)
    m = json.loads(ev(METRICS_JS, session))
    winH, winW, docH = m["winH"], m["winW"], m["docH"]
    hdr = m["hdr"] or 0
    maxY = max(0, docH - winH)

    # 首屏
    scroll_to(0, session)
    first = shot(session)
    scale = first.width / winW  # 物理/逻辑
    hdr_px = int(round(hdr * scale))

    tiles = [(first, 0, first.height)]  # (img, cropTop_px, used content bottom logical)
    prev_bottom = winH  # 已覆盖内容底部(逻辑)
    ntiles = 1
    mode = "single"

    if docH <= winH + 4:
        # window 不滚动: 尝试内部滚动容器(表格/列表页)
        sc = {}
        try:
            sc = json.loads(ev(DETECT_SCROLLER_JS, session) or "{}")
        except Exception:
            sc = {}
        if sc.get("found") and (sc["scrollHeight"] - sc["clientHeight"]) > 40:
            stitched, ntiles = _inner_stitch(session, sc, scale, max_tiles)
            mode = "inner"
        else:
            stitched = first
    else:
        parts = [first]
        target = winH - hdr
        guard = 0
        last_y = 0
        while guard < max_tiles:
            guard += 1
            actualY = scroll_to(target, session)
            # 页面实际没有继续前进(虚拟网格/内滚容器/滚动被重置): 停止, 避免重复拼接同一屏
            if actualY <= last_y + 8:
                break
            img = shot(session)
            crop_logical = prev_bottom - actualY
            if crop_logical < 0:
                crop_logical = 0
            crop_px = int(round(crop_logical * scale))
            if crop_px >= img.height:
                crop_px = hdr_px
            parts.append(img.crop((0, crop_px, img.width, img.height)))
            prev_bottom = actualY + winH
            last_y = actualY
            if actualY >= maxY - 1:
                break
            target = prev_bottom - hdr
        total_h = sum(p.height for p in parts)
        stitched = Image.new("RGB", (first.width, total_h), (255, 255, 255))
        y = 0
        for p in parts:
            stitched.paste(p, (0, y)); y += p.height
        ntiles = len(parts)
        mode = "window" if len(parts) > 1 else "single"

    scroll_to(0, session)
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    stitched.save(out_path, "JPEG", quality=78)
    print(json.dumps({"path": out_path, "size": [stitched.width, stitched.height],
                      "tiles": ntiles, "mode": mode,
                      "docH": docH, "winH": winH, "hdr": hdr}, ensure_ascii=False))

if __name__ == "__main__":
    _session = sys.argv[1]
    _out = sys.argv[2]
    _mt = int(sys.argv[3]) if len(sys.argv) > 3 else 14
    capture(_session, _out, _mt)