package workerclient

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/url"
	"strconv"
	"sync"
	"time"

	"google.golang.org/protobuf/proto"

	altv1 "git.toki-labs.com/toki/alt/packages/contracts/gen/go/alt/v1"
	apiContracts "git.toki-labs.com/toki/alt/services/api/internal/contracts"
	protoSocket "git.toki-labs.com/toki/proto-socket/go"
)

var (
	ErrUnavailable = errors.New("worker is not available")
	ErrTimeout     = errors.New("worker request timeout")
)

type WorkerClient interface {
	Connect(ctx context.Context) error
	Close() error
	IsConnected() bool
	Hello(ctx context.Context, req *altv1.HelloRequest) (*altv1.HelloResponse, error)

	// Backtest command/query surface. The API forwards client requests onto the
	// worker unchanged, so each method mirrors a contract request/response pair.
	StartBacktest(ctx context.Context, req *altv1.StartBacktestRequest) (*altv1.StartBacktestResponse, error)
	GetBacktestRun(ctx context.Context, req *altv1.GetBacktestRunRequest) (*altv1.GetBacktestRunResponse, error)
	ListBacktestRuns(ctx context.Context, req *altv1.ListBacktestRunsRequest) (*altv1.ListBacktestRunsResponse, error)
	GetBacktestRunDetail(ctx context.Context, req *altv1.GetBacktestRunDetailRequest) (*altv1.GetBacktestRunDetailResponse, error)
	GetBacktestResult(ctx context.Context, req *altv1.GetBacktestResultRequest) (*altv1.GetBacktestResultResponse, error)
	CompareBacktestRuns(ctx context.Context, req *altv1.CompareBacktestRunsRequest) (*altv1.CompareBacktestRunsResponse, error)

	// Paper trading command/query surface. The API forwards client requests onto
	// the worker unchanged; paper execution and state stay worker-owned.
	StartPaperTrading(ctx context.Context, req *altv1.StartPaperTradingRequest) (*altv1.StartPaperTradingResponse, error)
	GetPaperTradingState(ctx context.Context, req *altv1.GetPaperTradingStateRequest) (*altv1.GetPaperTradingStateResponse, error)

	// Paper order lifecycle surface. Submit/cancel/fill stay worker-owned; the API
	// only validates request shape and forwards.
	SubmitPaperOrder(ctx context.Context, req *altv1.SubmitPaperOrderRequest) (*altv1.SubmitPaperOrderResponse, error)
	CancelPaperOrder(ctx context.Context, req *altv1.CancelPaperOrderRequest) (*altv1.CancelPaperOrderResponse, error)
	FillPaperOrder(ctx context.Context, req *altv1.FillPaperOrderRequest) (*altv1.FillPaperOrderResponse, error)

	// Market query surface. The API validates request shape and forwards market
	// reads to the worker-owned storage boundary.
	ListInstruments(ctx context.Context, req *altv1.ListInstrumentsRequest) (*altv1.ListInstrumentsResponse, error)
	ListBars(ctx context.Context, req *altv1.ListBarsRequest) (*altv1.ListBarsResponse, error)

	// Market import command surface. Import execution runs inside the worker; the
	// API only validates shape and forwards the command.
	ImportDailyBars(ctx context.Context, req *altv1.ImportDailyBarsRequest) (*altv1.ImportDailyBarsResponse, error)

	// Monthly aggregation command surface. Aggregation execution runs inside the
	// worker; the API only validates shape and forwards the command.
	AggregateMonthlyBars(ctx context.Context, req *altv1.AggregateMonthlyBarsRequest) (*altv1.AggregateMonthlyBarsResponse, error)

	// Scheduler refresh status query surface. The worker holds the runtime
	// refresh status; the API only forwards the query.
	SchedulerRefreshStatus(ctx context.Context, req *altv1.SchedulerRefreshStatusRequest) (*altv1.SchedulerRefreshStatusResponse, error)

	// Live trading surface: broker capability probe. The API forwards client
	// requests onto the worker unchanged; live execution stays worker-owned.
	GetLiveBrokerCapabilities(ctx context.Context, req *altv1.GetLiveBrokerCapabilitiesRequest) (*altv1.GetLiveBrokerCapabilitiesResponse, error)

	// Live order lifecycle surface. Submit/cancel/get stay worker-owned; the API
	// only validates request shape and forwards.
	SubmitLiveOrder(ctx context.Context, req *altv1.SubmitLiveOrderRequest) (*altv1.SubmitLiveOrderResponse, error)
	CancelLiveOrder(ctx context.Context, req *altv1.CancelLiveOrderRequest) (*altv1.CancelLiveOrderResponse, error)
	GetLiveOrder(ctx context.Context, req *altv1.GetLiveOrderRequest) (*altv1.GetLiveOrderResponse, error)

	// Live risk/kill switch surface. Policy query and kill switch get/set stay
	// worker-owned; the API only validates and forwards.
	GetLiveRiskPolicy(ctx context.Context, req *altv1.GetLiveRiskPolicyRequest) (*altv1.GetLiveRiskPolicyResponse, error)
	GetLiveKillSwitch(ctx context.Context, req *altv1.GetLiveKillSwitchRequest) (*altv1.GetLiveKillSwitchResponse, error)
	SetLiveKillSwitch(ctx context.Context, req *altv1.SetLiveKillSwitchRequest) (*altv1.SetLiveKillSwitchResponse, error)

	// Live account sync surface. Fetch/get stay worker-owned; the API only
	// validates request shape and forwards.
	SyncLiveAccount(ctx context.Context, req *altv1.SyncLiveAccountRequest) (*altv1.SyncLiveAccountResponse, error)
	GetLiveAccountSnapshot(ctx context.Context, req *altv1.GetLiveAccountSnapshotRequest) (*altv1.GetLiveAccountSnapshotResponse, error)

	// Live audit trail query surface. The audit store lives in the worker;
	// the API only validates and forwards.
	ListLiveAuditEvents(ctx context.Context, req *altv1.ListLiveAuditEventsRequest) (*altv1.ListLiveAuditEventsResponse, error)
}

type socketClient struct {
	socketURL string
	mu        sync.RWMutex
	wsClient  *protoSocket.WsClient
}

func New(socketURL string) WorkerClient {
	return &socketClient{
		socketURL: socketURL,
	}
}

func (c *socketClient) IsConnected() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.wsClient != nil && c.wsClient.IsAlive()
}

func (c *socketClient) Connect(ctx context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.wsClient != nil && c.wsClient.IsAlive() {
		return nil
	}

	u, err := url.Parse(c.socketURL)
	if err != nil {
		return fmt.Errorf("invalid worker socket URL %q: %w", c.socketURL, err)
	}

	host, portStr, err := net.SplitHostPort(u.Host)
	if err != nil {
		host = u.Host
		if u.Scheme == "wss" {
			portStr = "443"
		} else {
			portStr = "80"
		}
	}

	port, err := strconv.Atoi(portStr)
	if err != nil {
		return fmt.Errorf("invalid worker port %q: %w", portStr, err)
	}

	path := u.Path
	if path == "" {
		path = "/"
	}

	var wsClient *protoSocket.WsClient
	if u.Scheme == "wss" {
		wsClient, err = protoSocket.DialWssWithHeartbeat(ctx, host, port, path, nil, 30, 10, apiContracts.ParserMap())
	} else {
		wsClient, err = protoSocket.DialWsWithHeartbeat(ctx, host, port, path, 30, 10, apiContracts.ParserMap())
	}

	if err != nil {
		return fmt.Errorf("%w: %v", ErrUnavailable, err)
	}

	c.wsClient = wsClient
	return nil
}

func (c *socketClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.wsClient != nil {
		err := c.wsClient.Close()
		c.wsClient = nil
		return err
	}
	return nil
}

func (c *socketClient) Hello(ctx context.Context, req *altv1.HelloRequest) (*altv1.HelloResponse, error) {
	return sendTyped[*altv1.HelloRequest, *altv1.HelloResponse](c, ctx, req)
}

func (c *socketClient) StartBacktest(ctx context.Context, req *altv1.StartBacktestRequest) (*altv1.StartBacktestResponse, error) {
	return sendTyped[*altv1.StartBacktestRequest, *altv1.StartBacktestResponse](c, ctx, req)
}

func (c *socketClient) GetBacktestRun(ctx context.Context, req *altv1.GetBacktestRunRequest) (*altv1.GetBacktestRunResponse, error) {
	return sendTyped[*altv1.GetBacktestRunRequest, *altv1.GetBacktestRunResponse](c, ctx, req)
}

func (c *socketClient) ListBacktestRuns(ctx context.Context, req *altv1.ListBacktestRunsRequest) (*altv1.ListBacktestRunsResponse, error) {
	return sendTyped[*altv1.ListBacktestRunsRequest, *altv1.ListBacktestRunsResponse](c, ctx, req)
}

func (c *socketClient) GetBacktestRunDetail(ctx context.Context, req *altv1.GetBacktestRunDetailRequest) (*altv1.GetBacktestRunDetailResponse, error) {
	return sendTyped[*altv1.GetBacktestRunDetailRequest, *altv1.GetBacktestRunDetailResponse](c, ctx, req)
}

func (c *socketClient) GetBacktestResult(ctx context.Context, req *altv1.GetBacktestResultRequest) (*altv1.GetBacktestResultResponse, error) {
	return sendTyped[*altv1.GetBacktestResultRequest, *altv1.GetBacktestResultResponse](c, ctx, req)
}

func (c *socketClient) CompareBacktestRuns(ctx context.Context, req *altv1.CompareBacktestRunsRequest) (*altv1.CompareBacktestRunsResponse, error) {
	return sendTyped[*altv1.CompareBacktestRunsRequest, *altv1.CompareBacktestRunsResponse](c, ctx, req)
}

func (c *socketClient) StartPaperTrading(ctx context.Context, req *altv1.StartPaperTradingRequest) (*altv1.StartPaperTradingResponse, error) {
	return sendTyped[*altv1.StartPaperTradingRequest, *altv1.StartPaperTradingResponse](c, ctx, req)
}

func (c *socketClient) GetPaperTradingState(ctx context.Context, req *altv1.GetPaperTradingStateRequest) (*altv1.GetPaperTradingStateResponse, error) {
	return sendTyped[*altv1.GetPaperTradingStateRequest, *altv1.GetPaperTradingStateResponse](c, ctx, req)
}

func (c *socketClient) SubmitPaperOrder(ctx context.Context, req *altv1.SubmitPaperOrderRequest) (*altv1.SubmitPaperOrderResponse, error) {
	return sendTyped[*altv1.SubmitPaperOrderRequest, *altv1.SubmitPaperOrderResponse](c, ctx, req)
}

func (c *socketClient) CancelPaperOrder(ctx context.Context, req *altv1.CancelPaperOrderRequest) (*altv1.CancelPaperOrderResponse, error) {
	return sendTyped[*altv1.CancelPaperOrderRequest, *altv1.CancelPaperOrderResponse](c, ctx, req)
}

func (c *socketClient) FillPaperOrder(ctx context.Context, req *altv1.FillPaperOrderRequest) (*altv1.FillPaperOrderResponse, error) {
	return sendTyped[*altv1.FillPaperOrderRequest, *altv1.FillPaperOrderResponse](c, ctx, req)
}

func (c *socketClient) ListInstruments(ctx context.Context, req *altv1.ListInstrumentsRequest) (*altv1.ListInstrumentsResponse, error) {
	return sendTyped[*altv1.ListInstrumentsRequest, *altv1.ListInstrumentsResponse](c, ctx, req)
}

func (c *socketClient) ListBars(ctx context.Context, req *altv1.ListBarsRequest) (*altv1.ListBarsResponse, error) {
	return sendTyped[*altv1.ListBarsRequest, *altv1.ListBarsResponse](c, ctx, req)
}

func (c *socketClient) ImportDailyBars(ctx context.Context, req *altv1.ImportDailyBarsRequest) (*altv1.ImportDailyBarsResponse, error) {
	return sendTyped[*altv1.ImportDailyBarsRequest, *altv1.ImportDailyBarsResponse](c, ctx, req)
}

func (c *socketClient) AggregateMonthlyBars(ctx context.Context, req *altv1.AggregateMonthlyBarsRequest) (*altv1.AggregateMonthlyBarsResponse, error) {
	return sendTyped[*altv1.AggregateMonthlyBarsRequest, *altv1.AggregateMonthlyBarsResponse](c, ctx, req)
}

func (c *socketClient) SchedulerRefreshStatus(ctx context.Context, req *altv1.SchedulerRefreshStatusRequest) (*altv1.SchedulerRefreshStatusResponse, error) {
	return sendTyped[*altv1.SchedulerRefreshStatusRequest, *altv1.SchedulerRefreshStatusResponse](c, ctx, req)
}

func (c *socketClient) GetLiveBrokerCapabilities(ctx context.Context, req *altv1.GetLiveBrokerCapabilitiesRequest) (*altv1.GetLiveBrokerCapabilitiesResponse, error) {
	return sendTyped[*altv1.GetLiveBrokerCapabilitiesRequest, *altv1.GetLiveBrokerCapabilitiesResponse](c, ctx, req)
}

func (c *socketClient) SubmitLiveOrder(ctx context.Context, req *altv1.SubmitLiveOrderRequest) (*altv1.SubmitLiveOrderResponse, error) {
	return sendTyped[*altv1.SubmitLiveOrderRequest, *altv1.SubmitLiveOrderResponse](c, ctx, req)
}

func (c *socketClient) CancelLiveOrder(ctx context.Context, req *altv1.CancelLiveOrderRequest) (*altv1.CancelLiveOrderResponse, error) {
	return sendTyped[*altv1.CancelLiveOrderRequest, *altv1.CancelLiveOrderResponse](c, ctx, req)
}

func (c *socketClient) GetLiveOrder(ctx context.Context, req *altv1.GetLiveOrderRequest) (*altv1.GetLiveOrderResponse, error) {
	return sendTyped[*altv1.GetLiveOrderRequest, *altv1.GetLiveOrderResponse](c, ctx, req)
}

func (c *socketClient) GetLiveRiskPolicy(ctx context.Context, req *altv1.GetLiveRiskPolicyRequest) (*altv1.GetLiveRiskPolicyResponse, error) {
	return sendTyped[*altv1.GetLiveRiskPolicyRequest, *altv1.GetLiveRiskPolicyResponse](c, ctx, req)
}

func (c *socketClient) GetLiveKillSwitch(ctx context.Context, req *altv1.GetLiveKillSwitchRequest) (*altv1.GetLiveKillSwitchResponse, error) {
	return sendTyped[*altv1.GetLiveKillSwitchRequest, *altv1.GetLiveKillSwitchResponse](c, ctx, req)
}

func (c *socketClient) SetLiveKillSwitch(ctx context.Context, req *altv1.SetLiveKillSwitchRequest) (*altv1.SetLiveKillSwitchResponse, error) {
	return sendTyped[*altv1.SetLiveKillSwitchRequest, *altv1.SetLiveKillSwitchResponse](c, ctx, req)
}

func (c *socketClient) SyncLiveAccount(ctx context.Context, req *altv1.SyncLiveAccountRequest) (*altv1.SyncLiveAccountResponse, error) {
	return sendTyped[*altv1.SyncLiveAccountRequest, *altv1.SyncLiveAccountResponse](c, ctx, req)
}

func (c *socketClient) GetLiveAccountSnapshot(ctx context.Context, req *altv1.GetLiveAccountSnapshotRequest) (*altv1.GetLiveAccountSnapshotResponse, error) {
	return sendTyped[*altv1.GetLiveAccountSnapshotRequest, *altv1.GetLiveAccountSnapshotResponse](c, ctx, req)
}

func (c *socketClient) ListLiveAuditEvents(ctx context.Context, req *altv1.ListLiveAuditEventsRequest) (*altv1.ListLiveAuditEventsResponse, error) {
	return sendTyped[*altv1.ListLiveAuditEventsRequest, *altv1.ListLiveAuditEventsResponse](c, ctx, req)
}

// sendTyped is the shared request path for every worker call. It centralises
// context-cancellation, deadline-derived timeouts, and the unavailable/timeout
// error mapping so each new request method stays a one-line forwarder and the
// behaviour cannot drift between them.
func sendTyped[Req proto.Message, Res proto.Message](c *socketClient, ctx context.Context, req Req) (Res, error) {
	var zero Res

	if err := ctx.Err(); err != nil {
		return zero, mapContextError(err)
	}

	c.mu.RLock()
	client := c.wsClient
	c.mu.RUnlock()

	if client == nil || !client.IsAlive() {
		return zero, ErrUnavailable
	}

	timeout := 5 * time.Second
	if dl, ok := ctx.Deadline(); ok {
		timeout = time.Until(dl)
		if timeout <= 0 {
			return zero, fmt.Errorf("%w: deadline already exceeded", ErrTimeout)
		}
	}

	type result struct {
		res Res
		err error
	}
	ch := make(chan result, 1)

	go func() {
		res, err := protoSocket.SendRequestTyped[Req, Res](&client.Communicator, req, timeout)
		ch <- result{res: res, err: err}
	}()

	select {
	case <-ctx.Done():
		return zero, mapContextError(ctx.Err())
	case r := <-ch:
		if r.err != nil {
			if errors.Is(r.err, protoSocket.ErrNotConnected) {
				return zero, ErrUnavailable
			}
			return zero, fmt.Errorf("%w: %v", ErrTimeout, r.err)
		}
		return r.res, nil
	}
}

// mapContextError normalises context errors into the worker client's error
// vocabulary: cancellation propagates as-is, deadlines surface as ErrTimeout.
func mapContextError(err error) error {
	if errors.Is(err, context.Canceled) {
		return err
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return fmt.Errorf("%w: %v", ErrTimeout, err)
	}
	return err
}