# -*- coding: utf-8 -*-
"""Google Lens 搜图 HTTP 服务

启动:
    python lens_server.py                    # 默认 0.0.0.0:8900
    python lens_server.py --port 9000        # 自定义端口
    python lens_server.py --warmup           # 启动时预热 cookie

API:
    POST /search  上传图片文件搜图
        curl -F "image=@product.png" http://localhost:8900/search

    POST /search  JSON 传图片 URL 搜图
        curl -X POST http://localhost:8900/search -H "Content-Type: application/json" -d '{"image_url": "https://example.com/img.jpg"}'

    GET  /search?image_url=https://...  URL 参数搜图
        curl "http://localhost:8900/search?image_url=https://example.com/img.jpg"

    GET  /health  健康检查
    POST /warmup  手动预热 cookie
"""
import os, sys, json, time, tempfile, threading

sys.path.insert(0, os.path.dirname(__file__))
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs

os.environ.setdefault("LLM_GPT4OMINI_API_KEY", "sk-sDU8gSwyoqUsH9OVVTaNObQ176PZ134I34SCpJNmnBTsOcBx")
os.environ.setdefault("LLM_GPT4OMINI_BASE_URL", "https://api.tu-zi.com/v1")
os.environ.setdefault("LLM_GPT4OMINI_MODEL", "gpt-4o-mini")

from lens_pure_python import google_lens_search, warmup_cookie_pool

_lock = threading.Lock()


class LensHandler(BaseHTTPRequestHandler):
    def _send_json(self, code, data):
        body = json.dumps(data, ensure_ascii=False, indent=2).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        parsed = urlparse(self.path)
        path = parsed.path.rstrip("/")

        if path == "/health":
            self._send_json(200, {"status": "ok", "time": time.strftime("%Y-%m-%d %H:%M:%S")})
            return

        if path == "/search":
            qs = parse_qs(parsed.query)
            image_url = qs.get("image_url", [None])[0]
            if not image_url:
                self._send_json(400, {"error": "需要 image_url 参数"})
                return
            hl = qs.get("hl", ["zh-CN"])[0]
            self._do_search(image_url=image_url, hl=hl)
            return

        self._send_json(404, {"error": "未知路径", "endpoints": ["/search", "/health", "/warmup"]})

    def do_POST(self):
        parsed = urlparse(self.path)
        path = parsed.path.rstrip("/")

        if path == "/warmup":
            ok = warmup_cookie_pool(force=True)
            self._send_json(200, {"status": "ok" if ok else "failed"})
            return

        if path != "/search":
            self._send_json(404, {"error": "未知路径"})
            return

        content_type = self.headers.get("Content-Type", "")
        content_length = int(self.headers.get("Content-Length", 0))

        # JSON body: {"image_url": "https://..."}
        if "application/json" in content_type:
            try:
                body = json.loads(self.rfile.read(content_length))
            except Exception:
                self._send_json(400, {"error": "无效 JSON"})
                return
            image_url = body.get("image_url", "")
            hl = body.get("hl", "zh-CN")
            if not image_url:
                self._send_json(400, {"error": "需要 image_url 字段"})
                return
            self._do_search(image_url=image_url, hl=hl)
            return

        # multipart/form-data: 文件上传
        if "multipart/form-data" in content_type:
            tmp_path = self._parse_multipart(content_type, content_length)
            if not tmp_path:
                self._send_json(400, {"error": "未找到 image 文件字段"})
                return
            try:
                self._do_search(image_path=tmp_path)
            finally:
                try:
                    os.remove(tmp_path)
                except Exception:
                    pass
            return

        self._send_json(400, {"error": "不支持的 Content-Type，用 multipart/form-data 或 application/json"})

    def _parse_multipart(self, content_type, content_length):
        boundary = ""
        for part in content_type.split(";"):
            part = part.strip()
            if part.startswith("boundary="):
                boundary = part[9:].strip('"')
                break
        if not boundary:
            return None

        raw = self.rfile.read(content_length)
        sep = f"--{boundary}".encode()
        parts = raw.split(sep)

        for part in parts:
            if b"Content-Disposition" not in part:
                continue
            header_end = part.find(b"\r\n\r\n")
            if header_end < 0:
                continue
            header = part[:header_end].decode("utf-8", errors="replace")
            if 'name="image"' not in header and 'name="file"' not in header:
                continue
            file_data = part[header_end + 4:]
            if file_data.endswith(b"\r\n"):
                file_data = file_data[:-2]

            ext = ".png"
            if b"\xff\xd8" in file_data[:4]:
                ext = ".jpg"
            elif b"WEBP" in file_data[:16]:
                ext = ".webp"

            tmp = os.path.join(tempfile.gettempdir(), f"lens_upload_{int(time.time()*1000)}{ext}")
            with open(tmp, "wb") as f:
                f.write(file_data)
            return tmp
        return None

    def _do_search(self, image_path=None, image_url=None, hl="zh-CN"):
        t0 = time.time()
        try:
            with _lock:
                result = google_lens_search(image_path=image_path, image_url=image_url, hl=hl)
            elapsed = time.time() - t0
            result["elapsed"] = f"{elapsed:.1f}s"
            self._send_json(200, result)
        except Exception as e:
            elapsed = time.time() - t0
            self._send_json(500, {"error": str(e), "elapsed": f"{elapsed:.1f}s"})

    def log_message(self, fmt, *args):
        ts = time.strftime("%H:%M:%S")
        print(f"[{ts}] {args[0]}" if args else "")


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Google Lens 搜图 HTTP 服务")
    parser.add_argument("--port", type=int, default=8900)
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--warmup", action="store_true", help="启动时预热 cookie")
    args = parser.parse_args()

    if args.warmup:
        print("[server] 预热 cookie ...")
        warmup_cookie_pool(force=True)

    server = HTTPServer((args.host, args.port), LensHandler)
    print(f"[server] Google Lens 搜图服务已启动: http://{args.host}:{args.port}")
    print(f"[server] POST /search  上传图片: curl -F 'image=@img.png' http://localhost:{args.port}/search")
    print(f"[server] POST /search  图片URL:  curl -X POST -H 'Content-Type: application/json' -d '{{\"image_url\":\"...\"}}'  http://localhost:{args.port}/search")
    print(f"[server] GET  /search  URL参数:  curl 'http://localhost:{args.port}/search?image_url=...'")
    print(f"[server] GET  /health  健康检查")
    print(f"[server] POST /warmup  刷新cookie")

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n[server] 已停止")
        server.server_close()


if __name__ == "__main__":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass
    main()
