#!/usr/bin/env python3
"""
NoScrible v3.0-math2latex PC Companion Server
==============================================

一体化 PC 端服务器，提供：
1. HTTP 服务 (端口 48158)
   - POST /math2latex — 本地 pix2tex LaTeX-OCR 公式识别
   - POST /v1/chat/completions — AI 代理（转发到 OpenAI 兼容 API）
   - GET  / — 状态页面 / 静态文件
2. TCP 服务 (端口 48159)
   - 接收平板发送的 LaTeX 代码，写入 PC 剪贴板

用法:
    python clipboard_server.py [--http-port 48158] [--tcp-port 48159]

依赖:
    pip install pix2tex pillow pyperclip
"""

import argparse
import base64
import io
import json
import os
import signal
import socket
import struct
import sys
import threading
import traceback
from functools import partial
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from io import BytesIO

# ──────────────────────────────────────────────
# 可选依赖导入（pix2tex + PIL + pyperclip）
# ──────────────────────────────────────────────

try:
    from PIL import Image
    HAS_PIL = True
except ImportError:
    HAS_PIL = False

try:
    import pyperclip
    HAS_CLIPBOARD = True
except ImportError:
    HAS_CLIPBOARD = False

# ─── pix2tex（LaTeX-OCR）───

MODEL = None  # 全局单例

def get_pix2tex_model():
    """懒加载初始化 pix2tex LaTeXOCR 模型（全局单例）"""
    global MODEL
    if MODEL is not None:
        return MODEL
    try:
        # noinspection PyUnresolvedReferences
        from pix2tex.cli import LaTeXOCR
        print("[pix2tex] 正在加载 LaTeX-OCR 模型（首次加载较慢）...")
        MODEL = LaTeXOCR()
        print("[pix2tex] ✓ 模型加载完毕")
        return MODEL
    except ImportError as e:
        print(f"[pix2tex] ✗ 导入失败: {e}")
        print("[pix2tex]   请执行: pip install pix2tex")
        return None
    except Exception as e:
        print(f"[pix2tex] ✗ 模型初始化异常: {e}")
        return None


# ──────────────────────────────────────────────
# HTTP 请求处理器
# ──────────────────────────────────────────────

class HTTPHandler(BaseHTTPRequestHandler):
    """处理所有 HTTP 请求"""

    # ── 基础配置 ──
    server_version = "NoScrible/3.0"
    sys_version = ""
    # 静态文件根目录（与脚本同目录下的 www / static）
    STATIC_DIRS = []

    @classmethod
    def configure(cls, tcp_port, static_dirs=None):
        cls.TCP_PORT = tcp_port
        if static_dirs:
            cls.STATIC_DIRS = static_dirs

    # ── CORS ──

    def _set_cors(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization")

    def _send_json(self, status_code, data):
        """发送 JSON 响应"""
        body = json.dumps(data, ensure_ascii=False).encode("utf-8")
        self.send_response(status_code)
        self._set_cors()
        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 _send_text(self, status_code, text, content_type="text/plain; charset=utf-8"):
        """发送纯文本响应"""
        body = text.encode("utf-8")
        self.send_response(status_code)
        self._set_cors()
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_error_json(self, status_code, message):
        self._send_json(status_code, {"error": True, "message": message})

    # ── 路由分发 ──

    def do_OPTIONS(self):
        """预检请求"""
        self.send_response(204)
        self._set_cors()
        self.end_headers()

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

        if path == "/":
            self._handle_status()
        elif path.startswith("/static/"):
            self._serve_static(path[8:])  # 去掉 "/static/"
        else:
            self._serve_static(path.lstrip("/"))

    def do_POST(self):
        parsed = urlparse(self.path)
        path = parsed.path

        if path == "/math2latex":
            self._handle_math2latex()
        elif path in ("/v1/chat/completions", "/chat/completions"):
            self._handle_ai_proxy()
        else:
            self._send_error_json(404, f"未知端点: {path}")

    # ── 端点：状态 ──

    def _handle_status(self):
        model_status = "loaded" if MODEL is not None else "not_loaded"
        if MODEL is None:
            # 尝试加载一次
            m = get_pix2tex_model()
            model_status = "loaded" if m is not None else "unavailable"

        status = {
            "service": "NoScrible v3.0-math2latex",
            "status": "running",
            "endpoints": {
                "GET  /": "服务状态",
                "GET  /static/...": "静态文件",
                "POST /math2latex": "pix2tex LaTeX-OCR 识别",
                "POST /v1/chat/completions": "AI 代理（OpenAI 兼容）",
            },
            "tcp_server": {"port": self.TCP_PORT, "description": "剪贴板接收"},
            "pix2tex_model": model_status,
            "clipboard": "available" if HAS_CLIPBOARD else "unavailable",
        }
        self._send_json(200, status)

    # ── 端点：/math2latex ──

    def _handle_math2latex(self):
        """接收 base64 二值图像，调用 pix2tex 识别，返回 LaTeX"""
        if not HAS_PIL:
            self._send_error_json(500, "PIL (Pillow) 未安装，无法处理图像")
            return

        # 1. 解析请求体
        try:
            content_length = int(self.headers.get("Content-Length", 0))
            if content_length == 0:
                self._send_error_json(400, "请求体为空")
                return
            raw_body = self.rfile.read(content_length)
            body = json.loads(raw_body)
        except json.JSONDecodeError:
            self._send_error_json(400, "请求体不是有效的 JSON")
            return
        except Exception as e:
            self._send_error_json(400, f"读取请求体失败: {e}")
            return

        # 2. 获取 base64 图像数据
        b64_data = body.get("image") or body.get("base64") or body.get("data")
        if not b64_data:
            self._send_error_json(400, "缺少 image/base64/data 字段")
            return

        # 去掉可能的 data: URL 前缀
        if b64_data.startswith("data:"):
            _, b64_data = b64_data.split(",", 1)

        # 3. 解码图像
        try:
            raw_bytes = base64.b64decode(b64_data)
            img = Image.open(BytesIO(raw_bytes))
            # 转为灰度图（pix2tex 需要单通道输入）
            if img.mode != "L":
                img = img.convert("L")
        except Exception as e:
            self._send_error_json(400, f"图像解码失败: {e}")
            return

        # 4. 调用 pix2tex
        model = get_pix2tex_model()
        if model is None:
            self._send_error_json(503, "pix2tex 模型未加载，请先安装: pip install pix2tex")
            return

        try:
            latex = model(img)
            self._send_json(200, {
                "success": True,
                "latex": latex.strip(),
                "engine": "pix2tex",
            })
        except Exception as e:
            tb = traceback.format_exc()
            print(f"[/math2latex] ✗ pix2tex 识别失败: {e}")
            print(tb)
            self._send_error_json(500, f"pix2tex 识别失败: {e}")

    # ── 端点：AI 代理 ──

    def _handle_ai_proxy(self):
        """AI 代理 — 将请求转发到 OpenAI 兼容的 API"""
        # 1. 解析请求体
        try:
            content_length = int(self.headers.get("Content-Length", 0))
            if content_length == 0:
                self._send_error_json(400, "请求体为空")
                return
            raw_body = self.rfile.read(content_length)
            body = json.loads(raw_body)
        except json.JSONDecodeError:
            self._send_error_json(400, "请求体不是有效的 JSON")
            return
        except Exception as e:
            self._send_error_json(400, f"读取请求体失败: {e}")
            return

        # 2. 从请求头或 body 配置字段获取目标 API 地址和 Key
        upstream_url = self.headers.get("X-Upstream-URL") or body.pop("upstream_url", None)
        upstream_key = self.headers.get("X-Upstream-Key") or body.pop("upstream_key", None)

        if not upstream_url:
            self._send_error_json(400, "缺少目标 API URL（请设置 X-Upstream-URL 请求头或 upstream_url 字段）")
            return

        # 去除末尾斜杠
        upstream_url = upstream_url.rstrip("/")
        # 补齐路径
        if not upstream_url.endswith("/chat/completions") and not upstream_url.endswith("/v1/chat/completions"):
            if upstream_url.endswith("/v1"):
                upstream_url += "/chat/completions"
            else:
                upstream_url += "/v1/chat/completions"

        # 3. 转发请求
        import urllib.request

        headers = {
            "Content-Type": "application/json",
        }
        if upstream_key:
            headers["Authorization"] = f"Bearer {upstream_key}"

        req_body = json.dumps(body, ensure_ascii=False).encode("utf-8")
        req = urllib.request.Request(
            upstream_url,
            data=req_body,
            headers=headers,
            method="POST",
        )

        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                upstream_body = resp.read()
                # 透传上游响应
                upstream_data = json.loads(upstream_body)
                self._send_json(resp.status, upstream_data)
        except urllib.error.HTTPError as e:
            error_body = e.read().decode("utf-8", errors="replace")
            try:
                error_data = json.loads(error_body)
            except json.JSONDecodeError:
                error_data = {"error": {"message": error_body[:500]}}
            self._send_json(e.code, error_data)
        except urllib.error.URLError as e:
            self._send_error_json(502, f"代理请求失败: {e.reason}")
        except Exception as e:
            self._send_error_json(500, f"代理异常: {e}")

    # ── 静态文件服务 ──

    def _serve_static(self, relative_path):
        """提供静态文件"""
        if not relative_path or relative_path.endswith("/"):
            relative_path = "index.html"

        # 防止路径穿越
        relative_path = os.path.normpath(relative_path.lstrip("/"))
        if relative_path.startswith("..") or relative_path.startswith("/"):
            self._send_error_json(403, "禁止访问")
            return

        for static_dir in self.STATIC_DIRS:
            filepath = os.path.join(static_dir, relative_path)
            if os.path.isfile(filepath):
                try:
                    with open(filepath, "rb") as f:
                        content = f.read()
                    content_type, _ = mimetypes.guess_type(filepath, strict=False)
                    if content_type is None:
                        content_type = "application/octet-stream"
                    self.send_response(200)
                    self._set_cors()
                    self.send_header("Content-Type", content_type)
                    self.send_header("Content-Length", str(len(content)))
                    self.end_headers()
                    self.wfile.write(content)
                    return
                except Exception as e:
                    self._send_error_json(500, f"读取文件失败: {e}")
                    return

        self._send_error_json(404, f"文件未找到: {relative_path}")

    # ── 日志抑制 ──

    def log_message(self, fmt, *args):
        """覆写以使用自定义格式"""
        msg = fmt % args
        print(f"[HTTP] {self.command} {self.path} → {msg}")


# 尝试导入 mimetypes (标准库)
try:
    import mimetypes
except ImportError:
    mimetypes = None


# ──────────────────────────────────────────────
# TCP 剪贴板服务器
# ──────────────────────────────────────────────

class ClipboardTCPServer:
    """TCP 剪贴板接收服务器（端口 48159）"""

    def __init__(self, host="0.0.0.0", port=48159):
        self.host = host
        self.port = port
        self.server = None
        self.running = False

    def start(self):
        self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server.bind((self.host, self.port))
        self.server.listen(5)
        self.running = True

        print(f"[TCP 剪贴板服务器] ✅ 已启动 (端口 {self.port})")

        while self.running:
            try:
                conn, addr = self.server.accept()
                thread = threading.Thread(
                    target=self._handle_client,
                    args=(conn, addr),
                    daemon=True,
                )
                thread.start()
            except OSError:
                break  # 服务器关闭

        print("[TCP 剪贴板服务器] 已停止")

    def stop(self):
        self.running = False
        if self.server:
            try:
                self.server.close()
            except OSError:
                pass

    @staticmethod
    def _handle_client(conn, addr):
        """处理一个 TCP 客户端连接"""
        print(f"[TCP] [+] 平板已连接: {addr}")
        try:
            # 读取 4 字节长度前缀 (大端序)
            raw_len = conn.recv(4)
            if not raw_len:
                return
            msg_len = struct.unpack("!I", raw_len)[0]

            # 读取报文内容
            data = b""
            remaining = msg_len
            while remaining > 0:
                chunk = conn.recv(min(remaining, 4096))
                if not chunk:
                    break
                data += chunk
                remaining -= len(chunk)

            if data:
                latex = data.decode("utf-8")
                if HAS_CLIPBOARD:
                    pyperclip.copy(latex)
                    print(f"[TCP] [✓] 已写入剪贴板: {latex[:80]}...")
                else:
                    print(f"[TCP] [✓] 收到 LaTeX（pyperclip 未安装，未写入剪贴板）: {latex[:80]}...")
                    print(f"[TCP]    请手工复制: {latex}")
        except Exception as e:
            print(f"[TCP] [!] 处理客户端 {addr} 出错: {e}")
        finally:
            conn.close()


# ──────────────────────────────────────────────
# 主控逻辑
# ──────────────────────────────────────────────

def get_local_ip():
    """获取本机局域网 IP"""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.settimeout(1)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return "127.0.0.1"


def run_http_server(http_port, tcp_port, static_dirs):
    """启动 HTTP 服务器"""
    HTTPHandler.configure(tcp_port=tcp_port, static_dirs=static_dirs)

    server = HTTPServer(("0.0.0.0", http_port), HTTPHandler)
    print(f"[HTTP 服务器] ✅ 已启动 (端口 {http_port})")

    # 优雅关闭
    def shutdown(sig, frame):
        print("\n[HTTP 服务器] 正在关闭...")
        server.shutdown()
        print("[HTTP 服务器] 已停止")

    signal.signal(signal.SIGINT, shutdown)
    signal.signal(signal.SIGTERM, shutdown)

    server.serve_forever()


def main():
    parser = argparse.ArgumentParser(
        description="NoScrible v3.0-math2latex PC Companion Server",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
    python clipboard_server.py
    python clipboard_server.py --http-port 48158 --tcp-port 48159
        """,
    )
    parser.add_argument(
        "--http-port", type=int, default=48158,
        help="HTTP 服务端口（默认 48158）",
    )
    parser.add_argument(
        "--tcp-port", type=int, default=48159,
        help="TCP 剪贴板服务端口（默认 48159）",
    )
    parser.add_argument(
        "--static-dir", type=str, action="append", default=[],
        help="静态文件目录（可多次指定）",
    )
    args = parser.parse_args()

    local_ip = get_local_ip()
    static_dirs = args.static_dir or [os.path.join(os.path.dirname(__file__), "www")]

    # ── 打印启动信息 ──
    print("=" * 54)
    print("  NoScrible v3.0-math2latex — PC Companion Server")
    print("=" * 54)
    print()
    print(f"  🌐 HTTP 服务:      http://0.0.0.0:{args.http_port}")
    print(f"    局域网访问:      http://{local_ip}:{args.http_port}")
    print(f"    端點列表:")
    print(f"      GET  /                 服务状态")
    print(f"      POST /math2latex       pix2tex LaTeX 公式识别")
    print(f"      POST /v1/chat/completions  AI 代理")
    print()
    print(f"  📋 TCP 剪贴板:    0.0.0.0:{args.tcp_port}")
    print(f"    在平板 APP 设置中输入:    {local_ip}")
    print()
    print(f"  📦 组件状态:")
    print(f"      Pillow:      {'✅' if HAS_PIL else '❌ 未安装'}")
    print(f"      pyperclip:   {'✅' if HAS_CLIPBOARD else '❌ 未安装（剪贴板功能受限）'}")
    print(f"      pix2tex:     {'🔄 懒加载（首次调用时初始化）' if get_pix2tex_model() is None else '✅ 已加载'}")
    print()

    # ── 预加载 pix2tex（后台线程） ──
    def preload_model():
        m = get_pix2tex_model()
        if m:
            print("[pix2tex] ✅ 后台预加载完成")
        else:
            print("[pix2tex] ⚠️  后台预加载失败，/math2latex 端点将不可用")

    preload_thread = threading.Thread(target=preload_model, daemon=True)
    preload_thread.start()

    # ── 启动 TCP 剪贴板服务器（后台线程） ──
    tcp_server = ClipboardTCPServer(port=args.tcp_port)
    tcp_thread = threading.Thread(target=tcp_server.start, daemon=True)
    tcp_thread.start()

    # ── 启动 HTTP 服务器（主线程） ──
    try:
        run_http_server(args.http_port, args.tcp_port, static_dirs)
    except KeyboardInterrupt:
        print("\n[主程序] 正在关闭所有服务...")
    finally:
        tcp_server.stop()
        print("[主程序] ✅ 所有服务已停止")


if __name__ == "__main__":
    main()
