package openai

import (
	"bytes"
	"context"
	"io"
	"net/http"
	"time"

	edgeservice "iop/apps/edge/internal/service"
	iop "iop/proto/gen/iop"
)

// providerFakeRunService extends the normalized fakeRunService with the
// provider tunnel/pool dispatch state used by provider scenario tests.
type providerFakeRunService struct {
	fakeRunService

	// tunnelProviderURL, when set, makes SubmitProviderTunnel relay the tunnel
	// request to this httptest provider fixture and stream the raw provider
	// response back as ordered frames (emulating the Node relay).
	tunnelProviderURL string
	// tunnelServedTarget emulates provider-pool admission: BuildBody is called
	// with this target when non-empty.
	tunnelServedTarget string
	// tunnelFrames, when set, is handed to the tunnel handle as-is instead of
	// contacting tunnelProviderURL.
	tunnelFrames chan *iop.ProviderTunnelFrame
	tunnelErr    error
	tunnelReqs   []edgeservice.SubmitProviderTunnelRequest
	tunnelBodies [][]byte
	// tunnelWaitTimeout overrides the tunnel handle wait timeout for
	// timeout-path fixtures; zero keeps the 5s default.
	tunnelWaitTimeout time.Duration

	// poolDispatchPath records which execution path was returned by
	// SubmitProviderPool: "provider_tunnel" or "normalized".
	poolDispatchPath string
	// poolRunFrames is the frame source used by the normalized path when
	// SubmitProviderPool is in use.
	poolRunFrames chan *iop.RunEvent
	// lastTunnelHandle stores the tunnel handle from the last SubmitProviderPool
	// tunnel path call, so tests can check headers set by the handler.
	lastTunnelHandle *fakeTunnelHandle
	// poolSubmitResults provides custom results for successive SubmitProviderPool
	// calls. When non-empty, results are consumed in order; when empty, the
	// poolDispatchPath / tunnelErr fields are used instead.
	poolSubmitResults []edgeservice.ProviderPoolDispatchResult
	// poolSubmitCount tracks how many times SubmitProviderPool was called
	// (including from the chat handler). Used by tests to assert dispatch
	// surface fidelity.
	poolSubmitCount int
	// poolLastRun records the last SubmitRunRequest passed to SubmitProviderPool
	// (pre-PrepareRun). Tests assert ModelGroupKey is non-empty to catch the
	// "empty Run" regression.
	poolLastRun edgeservice.SubmitRunRequest
	// poolPrepareRunCalled tracks whether the optional PrepareRun hook was
	// invoked by SubmitProviderPool on the normalized path. Only meaningful
	// when poolDispatchPath == "normalized".
	poolPrepareRunCalled bool
}

type fakeTunnelHandle struct {
	dispatch    edgeservice.RunDispatch
	headers     map[string]string
	frames      chan *iop.ProviderTunnelFrame
	waitTimeout time.Duration
}

func (h *fakeTunnelHandle) Headers() map[string]string        { return h.headers }
func (h *fakeTunnelHandle) SetHeaders(hdrs map[string]string) { h.headers = hdrs }
func (h *fakeTunnelHandle) Dispatch() edgeservice.RunDispatch { return h.dispatch }
func (h *fakeTunnelHandle) Stream() edgeservice.ProviderTunnelStream {
	return edgeservice.ProviderTunnelStream{Frames: h.frames}
}
func (h *fakeTunnelHandle) Close() {}
func (h *fakeTunnelHandle) WaitTimeout() time.Duration {
	if h.waitTimeout > 0 {
		return h.waitTimeout
	}
	return 5 * time.Second
}

// fakeRunResult implements edgeservice.RunResult for tests.
type fakeRunResult struct {
	dispatch edgeservice.RunDispatch
	events   chan *iop.RunEvent
}

func (r *fakeRunResult) Dispatch() edgeservice.RunDispatch { return r.dispatch }
func (r *fakeRunResult) Stream() edgeservice.RunStream {
	return edgeservice.RunStream{Events: r.events}
}
func (r *fakeRunResult) Close() {}
func (r *fakeRunResult) WaitTimeout() time.Duration {
	return 5 * time.Second
}

func (s *providerFakeRunService) SubmitProviderTunnel(_ context.Context, req edgeservice.SubmitProviderTunnelRequest) (edgeservice.ProviderTunnelResult, error) {
	s.submitMu.Lock()
	s.tunnelReqs = append(s.tunnelReqs, req)
	s.submitMu.Unlock()
	if s.tunnelErr != nil {
		return nil, s.tunnelErr
	}
	target := req.Target
	if s.tunnelServedTarget != "" {
		target = s.tunnelServedTarget
	}
	body := req.Body
	if req.BuildBody != nil {
		built, err := req.BuildBody(target)
		if err != nil {
			return nil, err
		}
		body = built
	}
	s.submitMu.Lock()
	s.tunnelBodies = append(s.tunnelBodies, body)
	s.submitMu.Unlock()

	frames := s.tunnelFrames
	if frames == nil && s.tunnelProviderURL != "" {
		relayed, err := relayProviderFixtureFrames(s.tunnelProviderURL, req.Method, req.Path, body)
		if err != nil {
			return nil, err
		}
		frames = relayed
	}
	return &fakeTunnelHandle{
		dispatch: edgeservice.RunDispatch{
			RunID:         "run-tunnel",
			NodeID:        "node-1",
			ModelGroupKey: req.ModelGroupKey,
			Adapter:       req.Adapter,
			Target:        target,
			SessionID:     req.SessionID,
			TimeoutSec:    5,
		},
		frames:      frames,
		waitTimeout: s.tunnelWaitTimeout,
	}, nil
}

// SubmitProviderPool is a minimal test double for the provider-pool one-shot
// dispatch. It records the selected path (tunnel vs normalized) and returns a
// stub result so the OpenAI handler can exercise the selection-first flow.
// It also tracks whether the optional PrepareRun hook was invoked, so tests
// can assert the handler only calls PrepareRun on the normalized path.
func (s *providerFakeRunService) SubmitProviderPool(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
	s.submitMu.Lock()
	s.poolSubmitCount++
	s.poolLastRun = req.Run
	s.poolPrepareRunCalled = false
	s.submitMu.Unlock()

	// Custom results for successive calls (e.g. tool-validation retry tests).
	s.submitMu.Lock()
	if len(s.poolSubmitResults) > 0 {
		result := s.poolSubmitResults[0]
		s.poolSubmitResults = s.poolSubmitResults[1:]
		s.submitMu.Unlock()
		// Apply pre-dispatch tunnel preparation for tunnel path when custom
		// results are used.
		if result.Path == edgeservice.ProviderPoolPathTunnel && req.PrepareTunnel != nil {
			tunnelReqPrepared, prepErr := req.PrepareTunnel(req.Tunnel)
			if prepErr != nil {
				return nil, prepErr
			}
			req.Tunnel = tunnelReqPrepared
		}
		return s.buildPoolResult(req, result)
	}
	err := s.tunnelErr
	path := s.poolDispatchPath
	s.submitMu.Unlock()

	// Apply pre-dispatch tunnel preparation (mirrors real SubmitProviderPool
	// behavior: PrepareTunnel runs BEFORE buildProviderTunnelRequest and the
	// Node Send step). On failure, no tunnel request is recorded or sent.
	tunnelReq := req.Tunnel
	if path == string(edgeservice.ProviderPoolPathTunnel) || path == "" {
		if req.PrepareTunnel != nil {
			tunnelReqPrepared, prepErr := req.PrepareTunnel(req.Tunnel)
			if prepErr != nil {
				return nil, prepErr
			}
			tunnelReq = tunnelReqPrepared
		}
	}

	// Record tunnel request for test assertions.
	if path == string(edgeservice.ProviderPoolPathTunnel) || path == "" {
		s.submitMu.Lock()
		s.tunnelReqs = append(s.tunnelReqs, tunnelReq)
		s.submitMu.Unlock()
	}

	if err != nil {
		return nil, err
	}

	disp := edgeservice.RunDispatch{
		RunID:         "run-pool-tunnel",
		NodeID:        "node-pool",
		ModelGroupKey: req.Run.ModelGroupKey,
		Adapter:       "selected",
		Target:        "served",
		ProviderID:    "prov-fake",
		ProviderType:  "vllm",
		ExecutionPath: string(edgeservice.ProviderPoolPathTunnel),
		QueueReason:   "dispatched",
	}
	if path == string(edgeservice.ProviderPoolPathNormalized) {
		disp.RunID = "run-pool-normalized"
	}

	if path == string(edgeservice.ProviderPoolPathTunnel) || path == "" {
		frames := s.tunnelFrames
		if frames == nil && s.tunnelProviderURL != "" {
			body := req.Tunnel.Body
			if req.Tunnel.BuildBody != nil {
				target := req.Tunnel.Target
				if s.tunnelServedTarget != "" {
					target = s.tunnelServedTarget
				}
				built, err := req.Tunnel.BuildBody(target)
				if err != nil {
					return nil, err
				}
				body = built
			}
			s.submitMu.Lock()
			s.tunnelBodies = append(s.tunnelBodies, body)
			s.submitMu.Unlock()
			relayed, err := relayProviderFixtureFrames(s.tunnelProviderURL, req.Tunnel.Method, req.Tunnel.Path, body)
			if err != nil {
				return nil, err
			}
			frames = relayed
		} else if frames == nil {
			frames = staticProviderTunnelFrames(`{"ok":true}`)
			// Record the built tunnel body even when using a static frame source.
			body := req.Tunnel.Body
			if req.Tunnel.BuildBody != nil {
				target := req.Tunnel.Target
				if s.tunnelServedTarget != "" {
					target = s.tunnelServedTarget
				}
				if built, buildErr := req.Tunnel.BuildBody(target); buildErr == nil {
					body = built
				}
			}
			s.submitMu.Lock()
			s.tunnelBodies = append(s.tunnelBodies, body)
			s.submitMu.Unlock()
		} else {
			body := req.Tunnel.Body
			if req.Tunnel.BuildBody != nil {
				target := req.Tunnel.Target
				if s.tunnelServedTarget != "" {
					target = s.tunnelServedTarget
				}
				if built, err := req.Tunnel.BuildBody(target); err == nil {
					body = built
				}
			}
			s.submitMu.Lock()
			s.tunnelBodies = append(s.tunnelBodies, body)
			s.submitMu.Unlock()
		}

		handle := &fakeTunnelHandle{
			dispatch: edgeservice.RunDispatch{
				RunID:         "run-tunnel",
				NodeID:        "node-1",
				ModelGroupKey: req.Tunnel.ModelGroupKey,
				Adapter:       req.Tunnel.Adapter,
				Target:        req.Tunnel.Target,
				SessionID:     req.Tunnel.SessionID,
				TimeoutSec:    5,
			},
			headers: req.Tunnel.Headers,
			frames:  frames,
		}
		// Store for later inspection by tests.
		s.submitMu.Lock()
		s.lastTunnelHandle = handle
		s.submitMu.Unlock()
		return &edgeservice.ProviderPoolDispatchResult{
			Path:         edgeservice.ProviderPoolPathTunnel,
			Tunnel:       handle,
			DispatchInfo: disp,
		}, nil
	}

	// Normalized path (mimics Ollama/CLI behavior).
	// Invoke the optional PrepareRun hook so tests can assert it only runs
	// on the normalized path (mirrors real SubmitProviderPool behavior).
	if req.PrepareRun != nil {
		s.submitMu.Lock()
		s.poolPrepareRunCalled = true
		s.submitMu.Unlock()
		prepared, prepErr := req.PrepareRun(req.Run)
		if prepErr != nil {
			return nil, prepErr
		}
		req.Run = prepared
	}
	s.submitMu.Lock()
	s.reqs = append(s.reqs, req.Run)
	s.submitMu.Unlock()

	runChan := s.poolRunFrames
	if runChan == nil {
		runChan = make(chan *iop.RunEvent, 1)
		runChan <- &iop.RunEvent{RunId: disp.RunID, Type: "complete"}
		close(runChan)
	}
	return &edgeservice.ProviderPoolDispatchResult{
		Path: edgeservice.ProviderPoolPathNormalized,
		Run: &fakeRunResult{
			dispatch: disp,
			events:   runChan,
		},
		DispatchInfo: disp,
	}, nil
}

// buildPoolResult assembles a ProviderPoolDispatchResult from a pre-configured
// custom result. It handles tunnel request recording, frame relay, and
// normalized path run construction.
func (s *providerFakeRunService) buildPoolResult(req edgeservice.ProviderPoolDispatchRequest, result edgeservice.ProviderPoolDispatchResult) (*edgeservice.ProviderPoolDispatchResult, error) {
	if result.Path == edgeservice.ProviderPoolPathTunnel || result.Path == "" {
		s.submitMu.Lock()
		s.tunnelReqs = append(s.tunnelReqs, req.Tunnel)
		s.submitMu.Unlock()
		var frames chan *iop.ProviderTunnelFrame
		if frames == nil && s.tunnelProviderURL != "" {
			body := req.Tunnel.Body
			if req.Tunnel.BuildBody != nil {
				target := req.Tunnel.Target
				if s.tunnelServedTarget != "" {
					target = s.tunnelServedTarget
				}
				if built, err := req.Tunnel.BuildBody(target); err == nil {
					body = built
				}
			}
			s.submitMu.Lock()
			s.tunnelBodies = append(s.tunnelBodies, body)
			s.submitMu.Unlock()
			relayed, err := relayProviderFixtureFrames(s.tunnelProviderURL, req.Tunnel.Method, req.Tunnel.Path, body)
			if err != nil {
				return nil, err
			}
			frames = relayed
		} else if frames == nil {
			frames = staticProviderTunnelFrames(`{"ok":true}`)
		}
		disp := result.DispatchInfo
		if disp.RunID == "" {
			disp.RunID = "run-tunnel"
		}
		if disp.NodeID == "" {
			disp.NodeID = "node-1"
		}
		handle := &fakeTunnelHandle{
			dispatch:    disp,
			headers:     req.Tunnel.Headers,
			frames:      frames,
			waitTimeout: s.tunnelWaitTimeout,
		}
		s.submitMu.Lock()
		s.lastTunnelHandle = handle
		s.submitMu.Unlock()
		return &edgeservice.ProviderPoolDispatchResult{
			Path:         edgeservice.ProviderPoolPathTunnel,
			Tunnel:       handle,
			DispatchInfo: disp,
		}, nil
	}

	// Normalized path.
	disp := result.DispatchInfo
	if disp.RunID == "" {
		disp.RunID = "run-pool-normalized"
	}
	if disp.NodeID == "" {
		disp.NodeID = "node-pool"
	}
	// Invoke the optional PrepareRun hook on the normalized path (mirrors
	// real SubmitProviderPool behavior).
	if req.PrepareRun != nil {
		prepared, prepErr := req.PrepareRun(req.Run)
		if prepErr != nil {
			return nil, prepErr
		}
		req.Run = prepared
	}
	s.submitMu.Lock()
	s.poolPrepareRunCalled = true
	s.submitMu.Unlock()
	s.submitMu.Lock()
	s.reqs = append(s.reqs, req.Run)
	s.submitMu.Unlock()
	// Use the Run from the custom result if provided, otherwise create a
	// default completing run.
	if result.Run != nil {
		return &edgeservice.ProviderPoolDispatchResult{
			Path:         edgeservice.ProviderPoolPathNormalized,
			Run:          result.Run,
			DispatchInfo: disp,
		}, nil
	}
	runChan := make(chan *iop.RunEvent, 1)
	runChan <- &iop.RunEvent{RunId: disp.RunID, Type: "complete"}
	close(runChan)
	return &edgeservice.ProviderPoolDispatchResult{
		Path: edgeservice.ProviderPoolPathNormalized,
		Run: &fakeRunResult{
			dispatch: disp,
			events:   runChan,
		},
		DispatchInfo: disp,
	}, nil
}

// tunnelHandleSnapshot returns the tunnel handle from the last SubmitProviderPool
// call, allowing tests to check headers set by the handler.
type tunnelHandleSnapshot struct {
	headers map[string]string
	runID   string
	nodeID  string
}

func (s *providerFakeRunService) tunnelHandleSnapshot() tunnelHandleSnapshot {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	h := s.lastTunnelHandle
	if h == nil {
		return tunnelHandleSnapshot{}
	}
	return tunnelHandleSnapshot{
		headers: h.Headers(),
		runID:   h.Dispatch().RunID,
		nodeID:  h.Dispatch().NodeID,
	}
}

// poolSubmitCountSnapshot returns the number of times SubmitProviderPool was
// called. Tests assert exactly-once dispatch from the handler.
func (s *providerFakeRunService) poolSubmitCountSnapshot() int {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	return s.poolSubmitCount
}

// poolLastRunSnapshot returns a copy of the last SubmitRunRequest passed to
// SubmitProviderPool. Tests assert ModelGroupKey is non-empty to catch the
// "empty Run" regression.
func (s *providerFakeRunService) poolLastRunSnapshot() edgeservice.SubmitRunRequest {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	return s.poolLastRun
}

// poolPrepareRunCalledSnapshot returns whether the PrepareRun hook was
// invoked by the last SubmitProviderPool call. Only meaningful when the last
// call was on the normalized path.
func (s *providerFakeRunService) poolPrepareRunCalledSnapshot() bool {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	return s.poolPrepareRunCalled
}

func (s *providerFakeRunService) tunnelReqsSnapshot() []edgeservice.SubmitProviderTunnelRequest {
	s.submitMu.Lock()
	reqs := make([]edgeservice.SubmitProviderTunnelRequest, len(s.tunnelReqs))
	copy(reqs, s.tunnelReqs)
	s.submitMu.Unlock()
	return reqs
}

func (s *providerFakeRunService) tunnelBodiesSnapshot() [][]byte {
	s.submitMu.Lock()
	defer s.submitMu.Unlock()
	return append([][]byte(nil), s.tunnelBodies...)
}

// relayProviderFixtureFrames performs the provider HTTP request the way the
// Node relay does and converts status/header/body bytes into ordered tunnel
// frames, chunking the body at odd boundaries to exercise ordered writes.
func relayProviderFixtureFrames(providerURL, method, path string, body []byte) (chan *iop.ProviderTunnelFrame, error) {
	httpReq, err := http.NewRequest(method, providerURL+path, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	headers := make(map[string]string)
	for k, values := range resp.Header {
		if len(values) > 0 {
			headers[k] = values[0]
		}
	}

	const chunkSize = 7
	frames := make(chan *iop.ProviderTunnelFrame, len(respBody)/chunkSize+3)
	seq := int64(0)
	frames <- &iop.ProviderTunnelFrame{
		RunId:      "run-tunnel",
		TunnelId:   "tunnel-1",
		Sequence:   seq,
		Kind:       iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
		StatusCode: int32(resp.StatusCode),
		Headers:    headers,
	}
	seq++
	for off := 0; off < len(respBody); off += chunkSize {
		end := off + chunkSize
		if end > len(respBody) {
			end = len(respBody)
		}
		frames <- &iop.ProviderTunnelFrame{
			RunId:    "run-tunnel",
			TunnelId: "tunnel-1",
			Sequence: seq,
			Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
			Body:     append([]byte(nil), respBody[off:end]...),
		}
		seq++
	}
	frames <- &iop.ProviderTunnelFrame{
		RunId:    "run-tunnel",
		TunnelId: "tunnel-1",
		Sequence: seq,
		Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
		End:      true,
	}
	close(frames)
	return frames, nil
}

func staticProviderTunnelFrames(body string) chan *iop.ProviderTunnelFrame {
	frames := make(chan *iop.ProviderTunnelFrame, 3)
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(body)}
	frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
	close(frames)
	return frames
}