import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python"))

import dataclasses
import threading
from typing import Iterator

from gobridge import expose, call_go, run, worker_id, worker_count

# ── worker_id / worker_count ──────────────────────────────────────────────────
# 只有 worker 0 才执行一次性初始化（如监听端口、建立长连接等），
# 其余 worker 跳过，避免端口冲突 / 重复连接。
print(f"[worker {worker_id}/{worker_count}] started", flush=True)
if worker_id == 0:

    def _init_shared_resource():
        # 示例：此处可启动 WebSocket 客户端、监听 TCP 端口等
        print(f"[worker {worker_id}] shared resource initialized", flush=True)

    threading.Thread(target=_init_shared_resource, daemon=True).start()


# ── 基础类型 ─────────────────────────────────────────────────────────────────


@expose
def add(a: int, b: int) -> int:
    return a + b


@expose
def range_gen(start: int, stop: int) -> Iterator[int]:
    """流式输出：对应 Go 侧 Invoke[chan int]"""
    for i in range(start, stop):
        yield i


@expose
def sum_stream(numbers: Iterator[int]) -> int:
    """流式输入：对应 Go 侧传入 chan int 参数"""
    return sum(numbers)


@expose
def double_stream(numbers: Iterator[int]) -> Iterator[int]:
    """双向流：输入每个数，yield 其平方"""
    for n in numbers:
        yield n * n


# ── struct（dataclass / dict）类型 ───────────────────────────────────────────


@dataclasses.dataclass
class User:
    id: int
    name: str
    score: float
    level: str = ""


@expose
def get_user(uid: int) -> dict:
    """普通调用：返回一个 struct（Go 对应 User）"""
    return {"id": uid, "name": f"user_{uid}", "score": uid * 1.5}


@expose
def total_score(users: list) -> float:
    """slice 输入：接收 []User，返回总分"""
    return sum(u["score"] for u in users)


@expose
def enrich_users(users: list) -> list:
    """slice 输入输出：为每个 user 追加 level 字段"""
    result = []
    for u in users:
        u = dict(u)
        u["level"] = "gold" if u["score"] >= 10 else "silver"
        result.append(u)
    return result


@expose
def gen_users(count: int) -> Iterator[dict]:
    """流式输出 struct：yield 多个 User，对应 Go 侧 Invoke[chan User]"""
    for i in range(1, count + 1):
        yield {"id": i, "name": f"user_{i}", "score": float(i * 3)}


@expose
def process_users(users: Iterator[dict]) -> Iterator[dict]:
    """双向流 struct：输入流式 User，yield 处理后的 User"""
    for u in users:
        yield {"id": u["id"], "name": u["name"].upper(), "score": u["score"] * 2}


# ── Server 全双工示例 ────────────────────────────────────────────────────────


@expose
def compute_with_go_mul(a: int, b: int) -> int:
    """示例1：call_go[int] 指定返回类型"""
    return call_go[int]("Multiply", a, b)


@expose
def squared_with_log(n: int) -> Iterator[int]:
    """示例2：流式输出，每次 yield 前 call_go("Log") 回调 Go"""
    for i in range(1, n + 1):
        call_go("Log", f"yielding {i}² = {i * i}")
        yield i * i


@expose
def to_upper(s: str) -> str:
    """辅助方法：被 Go 的 EnrichName handler 内部调用"""
    return s.upper()


@expose
def full_chain(name: str) -> str:
    """示例3：Go→Python→Go→Python 四层链路

    full_chain("world")
      → call_go[str]("EnrichName", "world")             # Python 调 Go
        → Invoke[string](ctx, serv, "to_upper", "world") # Go 再调 Python
        ← "WORLD"
      ← "Hello, WORLD!"
    ← "Hello, WORLD!"
    """
    return call_go[str]("EnrichName", name)


@expose
def get_user_via_go(uid: int) -> dict:
    """示例4：call_go[User] 自动将 Go 返回的 dict 构造为 dataclass 实例"""
    user = call_go[User]("MakeUser", uid)  # Go 返回 {"id":..,"name":..,"score":..}
    user.level = "gold" if user.score >= 10 else "silver"
    return dataclasses.asdict(user)


# ── Session 亲和示例 ──────────────────────────────────────────────────────────
# _sessions 保存每个 session 的状态，key 由调用方提供
_sessions: dict = {}

# _global_counter 是进程级全局变量，同一 worker 的所有 session 共享
_global_counter: int = 0


@expose
def global_increment(delta: int) -> str:
    global _global_counter
    _global_counter += delta
    return f"[worker {worker_id}] counter = {_global_counter}"


@expose
def global_get() -> str:
    return f"[worker {worker_id}] counter = {_global_counter}"


@expose
def session_init(session_id: str, value: int) -> str:
    _sessions[session_id] = {"value": value, "steps": []}
    return f"[worker {worker_id}] session {session_id} init with {value}"


@expose
def session_step(session_id: str, delta: int) -> int:
    s = _sessions[session_id]
    s["value"] += delta
    s["steps"].append(delta)
    return s["value"]


@expose
def session_result(session_id: str) -> dict:
    return _sessions.pop(session_id)


if __name__ == "__main__":
    run()
    print("worker_id", worker_id)
    print("worker_count", worker_count)