#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
本代码实现了在 Streamlit（默认端口8501）中集成 WebSocket 功能，
利用 Tornado 的 WebSocketHandler 将 WebSocket 服务挂载到 Streamlit 内置服务器上，
从而使前端页面和 WebSocket 服务共享同一端口。

WebSocket 功能：
1. 客户端通过 WebSocket 连接到 /ws 路径，并发送 JSON 格式的连接请求（包含 cmd、address、port）。
2. 服务端根据请求尝试建立 TCP 连接（目标地址和端口），并回复连接状态。
3. 建立成功后实现双向数据转发：客户端发送的数据转发到 TCP 连接，对方数据通过 WebSocket 返回给客户端。

注意：
- 由于 Streamlit 内置 Tornado 服务器接口属于内部接口，可能会随 Streamlit 版本变化，
  此代码使用 Server._instance 获取当前实例（替代已废弃的 Server.get_current()）。
- 部署环境若支持 SSL/TLS（如 HTTPS 环境），请使用 wss:// 开头的地址；否则使用 ws://。
"""

import asyncio
import json
import tornado.websocket
import streamlit as st

# ------------------------------------------------
# Tornado WebSocket 处理器，实现双向数据转发
# ------------------------------------------------
class WSHandler(tornado.websocket.WebSocketHandler):
    """
    WebSocket 处理器：
    1. 第一次收到消息时解析 JSON 数据，要求 cmd 字段为 "connect"，并根据 address 与 port 建立 TCP 连接；
    2. 连接成功后回复客户端，并启动后台任务将 TCP 数据转发至客户端；
    3. 后续收到的数据均转发到目标 TCP 连接。
    """
    def initialize(self):
        self.connected = False  # 标识是否已建立 TCP 连接
        self.tcp_reader = None  # TCP 连接的 reader
        self.tcp_writer = None  # TCP 连接的 writer
        self.tcp_task = None    # 后台转发任务

    def check_origin(self, origin):
        # 允许所有跨域连接
        return True

    async def open(self):
        # WebSocket 连接建立后调用
        print("WebSocket 连接已建立")

    async def on_message(self, message):
        """
        当收到客户端消息时调用：
        - 如果尚未建立 TCP 连接，则解析 JSON，要求 cmd 为 "connect"，获取目标地址和端口后建立连接；
        - 建立连接成功后回复客户端，并启动后台任务将 TCP 数据转发给客户端；
        - 如果已经建立连接，则将收到的数据转发到目标 TCP 连接。
        """
        if not self.connected:
            # 第一次收到消息，认为是建立连接的请求
            try:
                req = json.loads(message)
            except Exception as e:
                await self.write_message(json.dumps({
                    "cmd": "connect",
                    "status": "error",
                    "error": "invalid JSON"
                }))
                self.close()
                return

            if req.get("cmd") != "connect":
                await self.write_message(json.dumps({
                    "cmd": "connect",
                    "status": "error",
                    "error": "invalid command"
                }))
                self.close()
                return

            dest_addr = req.get("address")
            dest_port = req.get("port")
            print(f"收到连接请求：目标 {dest_addr}:{dest_port}")

            # 尝试建立到目标 TCP 的连接
            try:
                self.tcp_reader, self.tcp_writer = await asyncio.open_connection(dest_addr, dest_port)
            except Exception as e:
                error_msg = f"连接目标失败：{e}"
                print(error_msg)
                await self.write_message(json.dumps({
                    "cmd": "connect",
                    "status": "error",
                    "error": str(e)
                }))
                self.close()
                return

            # 连接成功后回复客户端，并启动后台任务转发 TCP 数据到 WebSocket
            self.connected = True
            await self.write_message(json.dumps({
                "cmd": "connect",
                "status": "ok"
            }))
            self.tcp_task = asyncio.create_task(self.tcp_to_ws())
        else:
            # 已建立 TCP 连接，转发客户端发送的数据到目标 TCP 连接
            try:
                if isinstance(message, str):
                    data = message.encode()
                else:
                    data = message
                self.tcp_writer.write(data)
                await self.tcp_writer.drain()
            except Exception as e:
                print("转发数据到 TCP 出现异常：", e)
                self.close()

    async def tcp_to_ws(self):
        """
        后台任务：不断从目标 TCP 连接读取数据，并通过 WebSocket 发送给客户端。
        """
        try:
            while True:
                data = await self.tcp_reader.read(1024)
                if not data:
                    break
                await self.write_message(data, binary=True)
        except Exception as e:
            print("TCP 到 WebSocket 转发出现异常：", e)
        finally:
            self.close()

    def on_close(self):
        # WebSocket 连接关闭时调用，清理相关资源
        print("WebSocket 连接关闭")
        if self.tcp_task:
            self.tcp_task.cancel()
        if self.tcp_writer:
            try:
                self.tcp_writer.close()
            except Exception as e:
                print("关闭 TCP 连接异常：", e)

# ------------------------------------------------
# 将 WebSocket 处理器添加到 Streamlit 内置的 Tornado 服务器中
# ------------------------------------------------
def add_ws_handler():
    """
    获取当前 Streamlit 服务器实例，并向其 Tornado 应用中添加 WebSocket 路由 /ws，
    使 WebSocket 服务与 Streamlit 前端共享同一端口（8501）。
    注意：本方法依赖 Streamlit 内部实现，可能会随版本变化。
    """
    try:
        # 导入 Server 模块，并通过 _instance 获取当前实例
        from streamlit.web.server.server import Server
        server = Server._instance
    except Exception as e:
        print("无法导入 streamlit.web.server.server 模块或获取实例：", e)
        return

    if server is None:
        print("未找到当前的 Streamlit 服务器实例。")
        return

    try:
        tornado_app = server._tornado
        # 添加 WebSocket 路由，路径为 /ws，允许所有域名匹配
        tornado_app.add_handlers(".*$", [(r"/ws", WSHandler)])
        print("已添加 WebSocket 处理器，访问路径为 /ws")
    except Exception as e:
        print("添加 WebSocket 处理器时出现异常：", e)

# ------------------------------------------------
# Streamlit 前端页面
# ------------------------------------------------
def main():
    st.title("Streamlit 集成 WebSocket 服务（同端口8501）")
    st.write("本页面通过 Streamlit 展示信息，同时在内置 Tornado 服务器上挂载了 WebSocket 处理器。")
    st.write("请使用 WebSocket 客户端连接：")
    st.code("wss://<服务器地址>:8501/ws", language="bash")
    st.write("若未启用 SSL/TLS，则使用：")
    st.code("ws://<服务器地址>:8501/ws", language="bash")
    st.write("客户端需发送 JSON 格式的连接请求，例如：")
    st.code(
r'''{
  "cmd": "connect",
  "address": "目标IP或域名",
  "port": 目标端口
}''', language="json"
    )

# 在 Streamlit 应用加载时添加 WebSocket 处理器
add_ws_handler()

# 显示 Streamlit 前端页面
main()