package main

import (
	"context"
	"fmt"
	"log"
	"path/filepath"
	"runtime"

	"git.fsdpf.net/go/gobridge"
)

type User struct {
	ID    int     `json:"id"`
	Name  string  `json:"name"`
	Score float64 `json:"score"`
	Level string  `json:"level,omitempty"`
}

func main() {
	_, file, _, _ := runtime.Caller(0)
	script := filepath.Join(filepath.Dir(file), "worker.py")

	pool, err := gobridge.NewPool(script,
		gobridge.WithWorkers(2),
		gobridge.WithMaxConns(4),
	)
	if err != nil {
		log.Fatal(err)
	}
	defer pool.Close()

	ctx := context.Background()

	demoPool(ctx, pool)
	demoServer(ctx, script)
	demoSession(ctx, script)
}

func demoPool(ctx context.Context, pool gobridge.Pool) {

	// ── 普通调用 ──────────────────────────────────────────────────────────
	sum, err := gobridge.Invoke[int](ctx, pool, "add", 3, 4)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("add(3, 4) =", sum) // 7

	// ── 流式输出：Python yield → Go channel ──────────────────────────────
	ch, err := gobridge.Invoke[chan int](ctx, pool, "range_gen", 1, 6)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Print("range_gen(1, 6) =")
	for v := range ch {
		fmt.Print(" ", v)
	}
	fmt.Println() // 1 2 3 4 5

	// ── 流式输入：Go channel → Python Iterator ───────────────────────────
	inputCh := make(chan int, 10)
	go func() {
		for i := 1; i <= 5; i++ {
			inputCh <- i
		}
		close(inputCh)
	}()
	total, err := gobridge.Invoke[int](ctx, pool, "sum_stream", inputCh)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("sum_stream(1..5) =", total) // 15

	// ── 双向流：Go channel 输入 + Go channel 输出 ────────────────────────
	inputCh2 := make(chan int, 10)
	go func() {
		for i := 1; i <= 5; i++ {
			inputCh2 <- i
		}
		close(inputCh2)
	}()
	outCh, err := gobridge.Invoke[chan int](ctx, pool, "double_stream", inputCh2)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Print("double_stream(1..5) =")
	for v := range outCh {
		fmt.Print(" ", v)
	}
	fmt.Println() // 1 4 9 16 25

	// ── struct 普通调用 ───────────────────────────────────────────────────
	user, err := gobridge.Invoke[User](ctx, pool, "get_user", 42)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("get_user(42) = %+v\n", user)

	// ── slice 输入，返回标量 ───────────────────────────────────────────────
	users := []User{
		{ID: 1, Name: "alice", Score: 5.0},
		{ID: 2, Name: "bob", Score: 8.0},
		{ID: 3, Name: "carol", Score: 12.0},
	}
	scoreSum, err := gobridge.Invoke[float64](ctx, pool, "total_score", users)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("total_score([alice,bob,carol]) = %.1f\n", scoreSum)

	// ── slice 输入输出 ─────────────────────────────────────────────────────
	enriched, err := gobridge.Invoke[[]User](ctx, pool, "enrich_users", users)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("enrich_users:")
	for _, u := range enriched {
		fmt.Printf("  %+v\n", u)
	}

	// ── 流式输出 struct：Python yield User → Go chan User ─────────────────
	userCh, err := gobridge.Invoke[chan User](ctx, pool, "gen_users", 3)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Print("gen_users(3) =")
	for u := range userCh {
		fmt.Printf(" {%d %s %.0f}", u.ID, u.Name, u.Score)
	}
	fmt.Println()

	// ── 双向流 struct：Go chan User 输入 → Python 处理 → Go chan User 输出 ─
	inCh := make(chan User, 5)
	go func() {
		for _, u := range users {
			inCh <- u
		}
		close(inCh)
	}()
	procCh, err := gobridge.Invoke[chan User](ctx, pool, "process_users", inCh)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("process_users:")
	for u := range procCh {
		fmt.Printf("  %+v\n", u)
	}
}

// goService 实现 Handler 接口，公开方法自动暴露给 Python 通过 call_go() 调用
type goService struct {
	pool gobridge.Pool // 用于 EnrichName 内部再调 Python
}

func (s *goService) Multiply(ctx context.Context, a, b int) (int, error) {
	return a * b, nil
}

func (s *goService) Log(msg string) {
	fmt.Println("[Go Log]", msg)
}

// EnrichName 内部通过 Invoke 调用 Python 的 to_upper，演示 Go→Python→Go→Python 四层链路
func (s *goService) EnrichName(ctx context.Context, name string) (string, error) {
	upper, err := gobridge.Invoke[string](ctx, s.pool, "to_upper", name)
	if err != nil {
		return "", err
	}
	return "Hello, " + upper + "!", nil
}

func (s *goService) MakeUser(ctx context.Context, uid int) (User, error) {
	return User{ID: uid, Name: fmt.Sprintf("user_%d", uid), Score: float64(uid) * 1.5}, nil
}

func demoSession(ctx context.Context, script string) {
	fmt.Println("\n── Session 亲和示例（workers=2）─────────────────────────────────")

	pool, err := gobridge.NewPool(script, gobridge.WithWorkers(2))
	if err != nil {
		log.Fatal(err)
	}
	defer pool.Close()

	// ── NewSession：三个 session，C 与 A 落在同一 worker，验证互不干扰 ───────
	sessA := gobridge.NewSession(pool) // worker 1
	sessB := gobridge.NewSession(pool) // worker 0
	sessC := gobridge.NewSession(pool) // worker 1（与 A 同进程，不同 session_id）

	msgA, err := gobridge.Invoke[string](ctx, sessA, "session_init", "A", 100)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("A init:", msgA)

	msgB, err := gobridge.Invoke[string](ctx, sessB, "session_init", "B", 200)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("B init:", msgB)

	msgC, err := gobridge.Invoke[string](ctx, sessC, "session_init", "C", 300)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("C init:", msgC) // 应与 A 在同一 worker

	// 对 A/B/C 各自做 step，验证三者状态完全独立
	v, _ := gobridge.Invoke[int](ctx, sessA, "session_step", "A", 10)
	fmt.Println("A step(+10) =", v) // 110
	v, _ = gobridge.Invoke[int](ctx, sessA, "session_step", "A", 5)
	fmt.Println("A step(+5)  =", v) // 115

	v, _ = gobridge.Invoke[int](ctx, sessB, "session_step", "B", 50)
	fmt.Println("B step(+50) =", v) // 250

	v, _ = gobridge.Invoke[int](ctx, sessC, "session_step", "C", 99)
	fmt.Println("C step(+99) =", v) // 399（与 A 同 worker 但不受影响）

	rA, _ := gobridge.Invoke[map[string]any](ctx, sessA, "session_result", "A")
	rB, _ := gobridge.Invoke[map[string]any](ctx, sessB, "session_result", "B")
	rC, _ := gobridge.Invoke[map[string]any](ctx, sessC, "session_result", "C")
	fmt.Printf("A result = %v\n", rA) // map[steps:[10 5] value:115]
	fmt.Printf("B result = %v\n", rB) // map[steps:[50]   value:250]
	fmt.Printf("C result = %v\n", rC) // map[steps:[99]   value:399]

	// ── 全局变量测试：sessA 和 sessC 同 worker，共享进程级 counter ──────────
	fmt.Println()
	r, _ := gobridge.Invoke[string](ctx, sessA, "global_increment", 10)
	fmt.Println("sessA +10:", r) // worker 1 counter = 10
	r, _ = gobridge.Invoke[string](ctx, sessC, "global_increment", 5)
	fmt.Println("sessC +5 :", r) // worker 1 counter = 15（与 sessA 共享）
	r, _ = gobridge.Invoke[string](ctx, sessB, "global_increment", 99)
	fmt.Println("sessB +99:", r) // worker 0 counter = 99（独立进程，从 0 开始）
	r, _ = gobridge.Invoke[string](ctx, sessA, "global_get")
	fmt.Println("sessA get:", r) // worker 1 counter = 15（不受 sessB 影响）

	// ── WithAffinity：相同 key 跨调用始终路由同一 worker ────────────────────
	fmt.Println()
	for i := range 4 {
		affinityCtx := gobridge.WithAffinity(ctx, "sticky-key")
		msg, _ := gobridge.Invoke[string](affinityCtx, pool, "session_init",
			fmt.Sprintf("aff-%d", i), i)
		fmt.Printf("WithAffinity(sticky-key) #%d → %s\n", i, msg)
	}

	// ── 对照组：不带亲和，轮询分配给两个 worker ─────────────────────────────
	fmt.Println()
	for i := range 4 {
		msg, _ := gobridge.Invoke[string](ctx, pool, "session_init",
			fmt.Sprintf("rr-%d", i), i)
		fmt.Printf("round-robin          #%d → %s\n", i, msg)
	}
}

func demoServer(ctx context.Context, script string) {
	fmt.Println("\n── Server 全双工示例 ─────────────────────────────────────────────")

	svc := &goService{}
	serv, err := gobridge.NewPool(script,
		gobridge.WithWorkers(1),
		gobridge.WithHandlers(svc),
	)
	if err != nil {
		log.Fatal(err)
	}
	defer serv.Close()
	svc.pool = serv // 注入 pool 供 EnrichName 内部调用

	// ── 示例1：Python 调用 Go Multiply ───────────────────────────────────────
	result, err := gobridge.Invoke[int](ctx, serv, "compute_with_go_mul", 6, 7)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("compute_with_go_mul(6, 7) =", result) // 42

	// ── 示例2：流式输出 + Go Log 回调 ────────────────────────────────────────
	ch, err := gobridge.Invoke[chan int](ctx, serv, "squared_with_log", 4)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Print("squared_with_log(4) =")
	for v := range ch {
		fmt.Print(" ", v)
	}
	fmt.Println() // 1 4 9 16

	// ── 示例3：Go→Python→Go→Python 四层全双工链路 ───────────────────────────
	// full_chain("world") → call_go[str]("EnrichName","world") → Invoke to_upper("world") → "WORLD"
	//                     ← "Hello, WORLD!"                    ← "Hello, WORLD!"
	greeting, err := gobridge.Invoke[string](ctx, serv, "full_chain", "world")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("full_chain(world) =", greeting) // Hello, WORLD!

	// ── 示例4：call_go[User] 将 Go 返回的 dict 自动构造为 dataclass ──────────
	enriched, err := gobridge.Invoke[User](ctx, serv, "get_user_via_go", 12)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("get_user_via_go(12) = %+v\n", enriched) // {ID:12 Name:user_12 Score:18 Level:gold}
}