"""
语料采集 · Mac客户端
====================
在Mac上运行，自动截图OCR+滚屏采集，发送到服务端处理

用法:
  # 实时采集模式（后台监控屏幕变化）
  python3 mac-corpus-agent.py --token YOUR_TOKEN --mode auto
  
  # 滚屏采集模式（采集历史对话，先手动滚到顶部）
  python3 mac-corpus-agent.py --token YOUR_TOKEN --mode scroll
  
  # 剪贴板采集模式（复制即采集）
  python3 mac-corpus-agent.py --token YOUR_TOKEN --mode clipboard

选项:
  --server    服务端地址 (默认 http://localhost:8084)
  --interval  采集间隔秒数 (默认 auto=5s, scroll=3s)
  --region    屏幕采集区域 (默认 全屏)
  --help      显示帮助
"""
import os
import sys
import json
import time
import argparse
import subprocess
import tempfile
import threading
from pathlib import Path
from typing import Optional

try:
    import websocket
except ImportError:
    print("❌ 缺少依赖: pip3 install websocket-client")
    sys.exit(1)

try:
    import pyautogui
except ImportError:
    print("❌ 缺少依赖: pip3 install pyautogui pillow")
    sys.exit(1)

# ============================================================
# 配置
# ============================================================

VERSION = "1.0.0"
SERVER_URL = "http://localhost:8084"
RECONNECT_DELAY = 5  # 断线重连等待秒数

# 键盘快捷键（全局热键用，需管理员权限）
HOTKEY_STOP = "esc"       # 停止采集
HOTKEY_PAUSE = "f6"       # 暂停/继续

# ============================================================
# macOS工具函数
# ============================================================

def screenshot(region: Optional[tuple] = None) -> bytes:
    """截取屏幕指定区域，返回PNG字节"""
    if region:
        img = pyautogui.screenshot(region=region)
    else:
        img = pyautogui.screenshot()
    
    buf = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    img.save(buf, format="PNG")
    buf.close()
    
    with open(buf.name, "rb") as f:
        data = f.read()
    
    os.unlink(buf.name)
    return data


def ocr_text(image_path: str) -> str:
    """使用macOS原生Vision框架做OCR（本地，不上传）"""
    script = f'''
    use framework "Vision"
    use scripting additions
    
    set theImage to (current application's NSImage's alloc()'s initWithContentsOfFile:"{image_path}")
    set requestHandler to (current application's VNImageRequestHandler's alloc()'s initWithData:(theImage's TIFFRepresentation()) options:(missing value))
    
    set textResult to ""
    set theRequest to (current application's VNRecognizeTextRequest's alloc()'s initWithCompletionHandler:(lambda request, error
        if error ≠ missing value then return
        set observations to request's results()
        repeat with obs in observations
            set topCandidate to (obs's topCandidates:(1))
            if topCandidate's count() > 0 then
                set candidate to (topCandidate's objectAtIndex:(0))
                set recognizedText to candidate's string() as text
                set textResult to textResult & recognizedText & linefeed
            end if
        end repeat
    end))
    
    theRequest's setRecognitionLevel:(VNRequestTextRecognitionLevel1)  -- Accurate
    requestHandler's performRequests:({{theRequest}}) |error|:(missing value)
    return textResult
    '''
    
    result = subprocess.run(
        ["osascript", "-e", script],
        capture_output=True, text=True, timeout=30
    )
    return result.stdout.strip()


def scroll_page(direction: str = "down", amount: int = 3):
    """模拟鼠标滚轮"""
    pyautogui.scroll(-amount if direction == "down" else amount)


def get_active_window_title() -> str:
    """获取当前活跃窗口标题"""
    script = 'tell application "System Events" to get name of first application process whose frontmost is true'
    result = subprocess.run(["osascript", "-e", script], capture_output=True, text=True)
    return result.stdout.strip()


def notify(title: str, message: str):
    """MacOS系统通知"""
    script = f'display notification "{message}" with title "{title}"'
    subprocess.run(["osascript", "-e", script], capture_output=True)


# ============================================================
# 采集器核心
# ============================================================

class CorpusCollector:
    """语料采集器"""
    
    def __init__(self, token: str, server: str, interval: float):
        self.token = token
        self.server = server
        self.interval = interval
        self.ws = None
        self.running = False
        self.paused = False
        self.last_text = ""  # 去重用
        self.collected_count = 0
        self.filtered_count = 0
    
    def connect_ws(self) -> bool:
        """连接WebSocket"""
        try:
            ws_url = self.server.replace("http://", "ws://").replace("https://", "wss://")
            ws_url = f"{ws_url}/ws/collect"
            
            self.ws = websocket.create_connection(
                f"{ws_url}?token={self.token}",
                timeout=10
            )
            
            # 验证连接
            self.ws.send(json.dumps({"type": "ping"}))
            resp = json.loads(self.ws.recv())
            
            if resp.get("type") == "pong":
                print(f"  ✅ WebSocket已连接")
                return True
        except Exception as e:
            print(f"  ❌ WebSocket连接失败: {e}")
        return False
    
    def send_text(self, text: str, source: str = "screen_capture"):
        """发送文本到服务器"""
        if not text or len(text) < 10:
            return
        
        # 简单去重（连续相同内容跳过）
        if text == self.last_text:
            return
        self.last_text = text
        
        try:
            if self.ws:
                self.ws.send(json.dumps({
                    "type": "text",
                    "text": text,
                    "source": source
                }))
                resp = json.loads(self.ws.recv())
                
                if resp.get("collected", 0) > 0:
                    self.collected_count += resp["collected"]
                    print(f"    ✅ 采集 {resp['collected']} 条 | 总计: {self.collected_count}")
                    notify("语料采集", f"已采集 {resp['collected']} 条对话")
                elif resp.get("valuable"):
                    pass  # 有价值但未成对
                else:
                    self.filtered_count += 1
        except Exception as e:
            print(f"    ⚠️ 发送失败: {e}")
    
    def process_screenshot(self):
        """截屏→OCR→发送"""
        try:
            # 截屏
            img_data = screenshot()
            
            # 存临时文件
            tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
            tmp.write(img_data)
            tmp.close()
            
            # OCR
            text = ocr_text(tmp.name)
            os.unlink(tmp.name)
            
            if text.strip():
                self.send_text(text.strip())
        except Exception as e:
            print(f"    ⚠️ 截图处理失败: {e}")
    
    # === 模式1: 实时采集 ===
    def run_auto_mode(self):
        """实时模式：后台监控屏幕变化"""
        print("\n🔴 实时采集模式启动中...")
        print(f"   采集间隔: {self.interval}秒")
        print(f"   按 ESC 停止, F6 暂停/继续")
        print(f"   活跃窗口: {get_active_window_title()}")
        
        self.running = True
        
        while self.running:
            if self.paused:
                time.sleep(1)
                continue
            
            self.process_screenshot()
            
            # 检查快捷键（每轮检查一次）
            # Mac上需要更好的热键方案，这里简化处理
            for _ in range(int(self.interval * 2)):
                if not self.running:
                    break
                time.sleep(0.5)
        
        self.cleanup()
    
    # === 模式2: 滚屏采集 ===
    def run_scroll_mode(self):
        """滚屏模式：自动向下滚动采集历史对话"""
        print("\n📜 滚屏采集模式启动")
        print(f"   采集间隔: {self.interval}秒")
        print(f"   请确保已手动滚动到页面顶部!")
        print(f"   5秒后开始...")
        
        time.sleep(5)
        
        self.running = True
        scroll_count = 0
        empty_rounds = 0
        
        while self.running and empty_rounds < 10:
            # 截图+OCR
            self.process_screenshot()
            
            # 滚动
            scroll_page("down", 5)
            scroll_count += 1
            
            # 检测是否到底（连续N次没有新内容）
            time.sleep(self.interval)
            
            if scroll_count % 10 == 0:
                print(f"   已滚动 {scroll_count} 次 | 采集: {self.collected_count} | 过滤: {self.filtered_count}")
        
        print(f"\n✅ 滚屏采集完成")
        print(f"   共滚动 {scroll_count} 次")
        print(f"   采集: {self.collected_count} 条")
        print(f"   过滤: {self.filtered_count} 条")
        
        notify("语料采集完成", f"共采集 {self.collected_count} 条对话")
        self.cleanup()
    
    # === 模式3: 剪贴板采集 ===
    def run_clipboard_mode(self):
        """剪贴板模式：监控剪贴板变化，自动采集"""
        print("\n📋 剪贴板采集模式启动")
        print(f"   监控间隔: {self.interval}秒")
        print(f"   复制内容后自动采集...")
        
        import subprocess
        
        self.running = True
        last_clip = ""
        
        while self.running:
            # 获取剪贴板内容
            result = subprocess.run(
                ["pbpaste"],
                capture_output=True, text=True
            )
            current = result.stdout.strip()
            
            if current and current != last_clip:
                print(f"\n  📋 检测到新内容 ({len(current)}字)")
                self.send_text(current, "clipboard")
                last_clip = current
            
            time.sleep(self.interval)
        
        self.cleanup()
    
    def cleanup(self):
        """清理"""
        if self.ws:
            try:
                self.ws.close()
            except:
                pass


# ============================================================
# 主入口
# ============================================================

def main():
    parser = argparse.ArgumentParser(description="语料采集 Mac 客户端")
    parser.add_argument("--token", required=True, help="Gitea Access Token")
    parser.add_argument("--server", default=SERVER_URL, help=f"服务端地址 (默认 {SERVER_URL})")
    parser.add_argument("--mode", choices=["auto", "scroll", "clipboard"], default="auto",
                        help="采集模式: auto=实时, scroll=滚屏, clipboard=剪贴板")
    parser.add_argument("--interval", type=float, default=0,
                        help="采集间隔秒数 (默认 auto=5, scroll=3, clipboard=2)")
    
    args = parser.parse_args()
    
    # 默认间隔
    if args.interval <= 0:
        intervals = {"auto": 5.0, "scroll": 3.0, "clipboard": 2.0}
        args.interval = intervals[args.mode]
    
    print(f"\n{'='*50}")
    print(f"🧠 语料采集 Mac 客户端 v{VERSION}")
    print(f"{'='*50}")
    print(f"  模式: {args.mode}")
    print(f"  服务端: {args.server}")
    print(f"  间隔: {args.interval}s")
    print(f"{'='*50}\n")
    
    collector = CorpusCollector(args.token, args.server, args.interval)
    
    # 连接服务端
    print("🔄 连接服务端...")
    if not collector.connect_ws():
        print("  ⚠️ WebSocket连接失败，将使用HTTP回退")
        print("  请确保服务端已启动: python3 server.py")
        print(f"  服务端地址: {args.server}")
    
    try:
        if args.mode == "auto":
            collector.run_auto_mode()
        elif args.mode == "scroll":
            collector.run_scroll_mode()
        elif args.mode == "clipboard":
            collector.run_clipboard_mode()
    except KeyboardInterrupt:
        print("\n\n⏹️ 用户停止")
        collector.cleanup()
    
    print("\n👋 采集结束")


if __name__ == "__main__":
    main()