package runner

import (
	"bufio"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
	"github.com/google/uuid"
)

const (
	workerExecutionAudience = "easyai-worker-execution"
	workerExecutionPath     = "/internal/v1/executions"
)

type InternalExecutionRequest struct {
	TaskID  string `json:"task_id"`
	LeaseID string `json:"lease_id"`
	Stream  bool   `json:"stream"`
}

type InternalExecutionFrame struct {
	Type      string                    `json:"type"`
	Delta     *clients.StreamDeltaEvent `json:"delta,omitempty"`
	Output    map[string]any            `json:"output,omitempty"`
	Wire      *clients.WireResponse     `json:"wire,omitempty"`
	Code      string                    `json:"code,omitempty"`
	Message   string                    `json:"message,omitempty"`
	Status    int                       `json:"status,omitempty"`
	Retryable bool                      `json:"retryable,omitempty"`
	Details   map[string]any            `json:"details,omitempty"`
}

func (s *Service) executeRouted(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) {
	if task.AsyncMode || !s.cfg.RunsPublicHTTP() || !s.routingEnabled() || task.RunMode == "simulation" {
		return s.executeLocal(ctx, task, user, onDelta)
	}
	candidate, err := s.routingCandidateForTask(ctx, task, user)
	if err != nil {
		if s.routingEnforced() {
			s.observeExecutionPoolRouting("unavailable")
			return Result{}, upstreamRouteUnavailable(err)
		}
		return s.executeLocal(ctx, task, user, onDelta)
	}
	decision, profile, err := s.selectExecutionPool(ctx, task, candidate)
	if err != nil {
		_ = s.coordinationStore.RequestRouteProbe(context.WithoutCancel(ctx), profile.Key)
		if s.routingEnforced() {
			s.observeExecutionPoolRouting("unavailable")
			return Result{}, upstreamRouteUnavailable(err)
		}
		_ = s.store.AssignTaskRouting(ctx, store.TaskRoutingDecision{
			TaskID: task.ID, RouteProfileKey: profile.Key, RoutingVersion: routingVersion,
			PlatformID: candidate.PlatformID, PlatformModelID: candidate.PlatformModelID,
			Reason: "shadow_no_eligible_pool", Snapshot: map[string]any{"mode": "shadow", "error": safeRoutingError(err)},
		})
		s.observeExecutionPoolRouting("shadow")
		return s.executeLocal(ctx, task, user, onDelta)
	}
	snapshot := routingDecisionSnapshot(decision)
	if !s.routingEnforced() {
		snapshot["suggestedPoolId"] = decision.PoolID
		_ = s.store.AssignTaskRouting(ctx, store.TaskRoutingDecision{
			TaskID: task.ID, RouteProfileKey: profile.Key, RoutingVersion: routingVersion,
			PlatformID: candidate.PlatformID, PlatformModelID: candidate.PlatformModelID,
			Reason: decision.Reason, Snapshot: snapshot,
		})
		s.observeExecutionPoolRouting("shadow")
		return s.executeLocal(ctx, task, user, onDelta)
	}
	if err := s.store.AssignTaskRouting(ctx, store.TaskRoutingDecision{
		TaskID: task.ID, PoolID: decision.PoolID, RouteProfileKey: profile.Key,
		PlatformID: candidate.PlatformID, PlatformModelID: candidate.PlatformModelID,
		RoutingVersion: routingVersion, Reason: decision.Reason, Snapshot: snapshot,
	}); err != nil {
		return Result{}, err
	}
	s.observeExecutionPoolRouting("selected")
	return s.executeThroughWorker(ctx, task, decision.PoolID, onDelta)
}

func (s *Service) executeLocal(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) {
	return s.executeWithToken(ctx, task, user, onDelta, uuid.NewString())
}

func (s *Service) executeThroughWorker(ctx context.Context, task store.GatewayTask, poolID string, onDelta clients.StreamDelta) (Result, error) {
	nonce := uuid.NewString()
	lease, err := s.coordinationStore.ReserveWorkerExecution(ctx, task.ID, poolID, nonce, 30*time.Second)
	if err != nil {
		s.observeExecutionPoolRouting("capacity_rejected")
		return Result{}, upstreamRouteUnavailable(err)
	}
	if err := executionpool.ValidateAdvertisedEndpoint(
		lease.Endpoint, splitTrimmed(s.cfg.WorkerEndpointAllowedSuffixes), s.cfg.WorkerEndpointAllowPrivate,
	); err != nil {
		_ = s.coordinationStore.ReleaseWorkerExecutionLease(context.WithoutCancel(ctx), lease.LeaseID)
		return Result{}, &clients.ClientError{Code: "worker_endpoint_untrusted", Message: err.Error(), StatusCode: http.StatusServiceUnavailable, Retryable: false}
	}
	signer := executionpool.TokenSigner{Secret: []byte(s.cfg.WorkerExecutionSecret)}
	token, err := signer.Sign(executionpool.ExecutionClaims{
		Audience: workerExecutionAudience, TaskID: task.ID, PoolID: poolID,
		WorkerID: lease.WorkerID, Nonce: nonce, ExpiresAt: time.Now().Add(30 * time.Second).Unix(),
	})
	if err != nil {
		_ = s.coordinationStore.ReleaseWorkerExecutionLease(context.WithoutCancel(ctx), lease.LeaseID)
		return Result{}, err
	}
	client, err := s.workerExecutionHTTPClient()
	if err != nil {
		return Result{}, err
	}
	transport := httpExecutionTransport{client: client}
	response, err := transport.Execute(ctx, executionpool.WorkerDescriptor{
		WorkerID: lease.WorkerID, InstanceID: lease.InstanceID, PoolID: poolID, Endpoint: lease.Endpoint,
	}, executionpool.ExecutionRequest{
		TaskID: task.ID, PoolID: poolID, WorkerID: lease.WorkerID, LeaseID: lease.LeaseID,
		AuthorizationToken: token, Stream: onDelta != nil,
	})
	if err != nil {
		latest, readErr := s.store.GetTask(context.WithoutCancel(ctx), task.ID)
		if readErr == nil && latest.SubmissionState != "not_started" {
			return Result{Task: latest}, &clients.ClientError{
				Code: "upstream_timeout", Message: "worker transport ended after upstream submission began",
				StatusCode: http.StatusGatewayTimeout, Retryable: true,
			}
		}
		return Result{}, &clients.ClientError{Code: "worker_transport_error", Message: err.Error(), StatusCode: http.StatusServiceUnavailable, Retryable: true}
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		payload, _ := io.ReadAll(io.LimitReader(response.Body, 64*1024))
		return Result{}, &clients.ClientError{Code: "worker_transport_error", Message: strings.TrimSpace(string(payload)), StatusCode: response.StatusCode, Retryable: response.StatusCode >= 500}
	}
	var output map[string]any
	var wire *clients.WireResponse
	scanner := bufio.NewScanner(response.Body)
	scanner.Buffer(make([]byte, 4096), 1024*1024)
	for scanner.Scan() {
		var frame InternalExecutionFrame
		if err := json.Unmarshal(scanner.Bytes(), &frame); err != nil {
			return Result{}, &clients.ClientError{Code: "worker_protocol_error", Message: "invalid worker execution frame", StatusCode: http.StatusBadGateway, Retryable: false}
		}
		switch frame.Type {
		case "delta":
			if onDelta != nil && frame.Delta != nil {
				if err := onDelta(*frame.Delta); err != nil {
					return Result{}, err
				}
			}
		case "result":
			output, wire = frame.Output, frame.Wire
		case "error":
			return Result{}, &clients.ClientError{
				Code: frame.Code, Message: frame.Message, StatusCode: frame.Status,
				Retryable: frame.Retryable, Details: frame.Details,
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return Result{}, &clients.ClientError{Code: "worker_transport_error", Message: err.Error(), StatusCode: http.StatusBadGateway, Retryable: true}
	}
	finished, err := s.store.GetTask(ctx, task.ID)
	if err != nil {
		return Result{}, err
	}
	return Result{Task: finished, Output: output, Wire: wire}, nil
}

func (s *Service) ExecuteInternal(ctx context.Context, authorization string, input InternalExecutionRequest, onDelta clients.StreamDelta) (Result, error) {
	if s.cfg.RunsPublicHTTP() || !s.cfg.RunsAsyncExecutionWorker() {
		return Result{}, errors.New("internal execution is only available on worker processes")
	}
	token := strings.TrimSpace(strings.TrimPrefix(authorization, "Worker "))
	if token == authorization || token == "" {
		return Result{}, errors.New("worker execution authorization is required")
	}
	signer := executionpool.TokenSigner{Secret: []byte(s.cfg.WorkerExecutionSecret)}
	claims, err := signer.Verify(token, workerExecutionAudience)
	if err != nil {
		return Result{}, err
	}
	localWorkerID := firstNonEmptyString(s.cfg.WorkerID, s.workerInstanceID)
	if claims.TaskID != input.TaskID || claims.PoolID != s.cfg.ExecutionPoolID || claims.WorkerID != localWorkerID {
		return Result{}, errors.New("worker execution claims do not match this worker")
	}
	lease, err := s.coordinationStore.ConsumeWorkerExecutionLease(ctx, input.LeaseID, claims.Nonce)
	if err != nil {
		return Result{}, err
	}
	defer func() {
		_ = s.coordinationStore.ReleaseWorkerExecutionLease(context.WithoutCancel(ctx), lease.LeaseID)
	}()
	if lease.TaskID != input.TaskID || lease.PoolID != s.cfg.ExecutionPoolID || lease.WorkerID != localWorkerID || lease.InstanceID != s.workerInstanceID {
		return Result{}, errors.New("worker execution lease does not match this worker")
	}
	loadLease, admitted := s.tryStartWorkerTask()
	if !admitted {
		return Result{}, &clients.ClientError{Code: "worker_capacity_unavailable", Message: "worker has no safe execution capacity", StatusCode: http.StatusServiceUnavailable, Retryable: true}
	}
	defer loadLease.Release()
	ctx = context.WithValue(ctx, workerLoadLeaseContextKey{}, loadLease)
	task, err := s.store.GetTask(ctx, input.TaskID)
	if err != nil {
		return Result{}, err
	}
	return s.executeLocal(ctx, task, authUserFromTask(task), onDelta)
}

func (s *Service) workerExecutionHTTPClient() (*http.Client, error) {
	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.Proxy = nil
	if strings.TrimSpace(s.cfg.WorkerExecutionCAFile) != "" {
		caBytes, err := os.ReadFile(s.cfg.WorkerExecutionCAFile)
		if err != nil {
			return nil, err
		}
		roots := x509.NewCertPool()
		if !roots.AppendCertsFromPEM(caBytes) {
			return nil, errors.New("worker execution CA file does not contain a certificate")
		}
		transport.TLSClientConfig = &tls.Config{RootCAs: roots, MinVersion: tls.VersionTLS12}
		if strings.TrimSpace(s.cfg.WorkerExecutionCertFile) != "" || strings.TrimSpace(s.cfg.WorkerExecutionKeyFile) != "" {
			certificate, err := tls.LoadX509KeyPair(s.cfg.WorkerExecutionCertFile, s.cfg.WorkerExecutionKeyFile)
			if err != nil {
				return nil, err
			}
			transport.TLSClientConfig.Certificates = []tls.Certificate{certificate}
		}
	}
	return &http.Client{Transport: transport}, nil
}

func internalExecutionStatus(err error) int {
	var clientErr *clients.ClientError
	if errors.As(err, &clientErr) && clientErr.StatusCode > 0 {
		return clientErr.StatusCode
	}
	return http.StatusInternalServerError
}

func InternalExecutionFrameForError(err error) InternalExecutionFrame {
	var clientErr *clients.ClientError
	if errors.As(err, &clientErr) {
		return InternalExecutionFrame{
			Type: "error", Code: clientErr.Code, Message: clientErr.Message,
			Status: internalExecutionStatus(err), Retryable: clientErr.Retryable, Details: clientErr.Details,
		}
	}
	return InternalExecutionFrame{Type: "error", Code: "worker_execution_failed", Message: fmt.Sprint(err), Status: http.StatusInternalServerError}
}