package batch

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"sub2api-cn-relay-manager/internal/probe"
	"sub2api-cn-relay-manager/internal/store/sqlite"
)

type BatchImportEntry struct {
	BaseURL         string
	APIKey          string
	RequestedModels []string
}

type BatchImportRunRequest struct {
	RunID             string
	Mode              string
	AccessMode        string
	HostID            string
	HostBaseURL       string
	SubscriptionUsers []string
	SubscriptionDays  int
	ProbeAPIKey       string
	Entries           []BatchImportEntry
}

type BatchImportRunResult struct {
	RunID   string
	ItemIDs []string
}

type RunStateStore interface {
	Create(ctx context.Context, run sqlite.ImportRun) error
	Update(ctx context.Context, run sqlite.ImportRun) error
}

type ItemStateStore interface {
	Upsert(ctx context.Context, item sqlite.ImportRunItem) error
}

type ReuseLookupInput struct {
	HostID                 string
	ProviderID             string
	BaseURL                string
	APIKeyFingerprint      string
	CanonicalModelFamilies []string
}

type ReuseLookupResult struct {
	ProviderMatched          bool
	ExistingProviderID       string
	ExistingAccessStatus     AccessStatus
	ExistingCanonicalFamilys []string
	MatchedAccountID         int64
	MatchedAccountState      MatchedAccountState
	ExistingModelMapping     map[string]string
	LegacyBatchID            *int64
}

type ProvisionRequest struct {
	RunID             string
	ItemID            string
	Entry             BatchImportEntry
	ProviderID        string
	ResolvedModel     string
	RoutingStrategy   ImportRoutingStrategy
	CapabilityProfile *probe.CapabilityProfile
}

type ProvisionResult struct {
	LegacyBatchID    *int64
	LegacyProviderID string
}

type PatchProvisionRequest struct {
	ProviderID string
	Contract   ChannelPatchContract
}

type BatchProvisioner interface {
	Provision(ctx context.Context, req ProvisionRequest) (ProvisionResult, error)
	Patch(ctx context.Context, req PatchProvisionRequest) error
}

type BatchImportService struct {
	RunStore          RunStateStore
	ItemStore         ItemStateStore
	ProbeModels       func(ctx context.Context, baseURL, apiKey string) (*probe.ModelsResult, error)
	ProbeCapabilities func(ctx context.Context, baseURL, apiKey string, rawModels []string) (*probe.CapabilityProfile, error)
	InspectReuse      func(ctx context.Context, input ReuseLookupInput) (ReuseLookupResult, error)
	Provisioner       BatchProvisioner
}

func (s BatchImportService) StartRun(ctx context.Context, req BatchImportRunRequest) (BatchImportRunResult, error) {
	if s.RunStore == nil {
		return BatchImportRunResult{}, fmt.Errorf("run store is required")
	}
	if s.ItemStore == nil {
		return BatchImportRunResult{}, fmt.Errorf("item store is required")
	}
	if s.ProbeModels == nil {
		return BatchImportRunResult{}, fmt.Errorf("model probe is required")
	}
	if s.ProbeCapabilities == nil {
		return BatchImportRunResult{}, fmt.Errorf("capability probe is required")
	}

	runID := strings.TrimSpace(req.RunID)
	if runID == "" {
		runID = fmt.Sprintf("run-%d", time.Now().UnixNano())
	}
	if len(req.Entries) == 0 {
		return BatchImportRunResult{}, fmt.Errorf("entries are required")
	}

	if err := s.RunStore.Create(ctx, sqlite.ImportRun{
		RunID:                 runID,
		HostID:                strings.TrimSpace(req.HostID),
		Mode:                  strings.TrimSpace(req.Mode),
		AccessMode:            strings.TrimSpace(req.AccessMode),
		SubscriptionUsersJSON: mustMarshalJSON(req.SubscriptionUsers, "[]"),
		SubscriptionDays:      req.SubscriptionDays,
		ProbeAPIKey:           strings.TrimSpace(req.ProbeAPIKey),
		State:                 string(RunStateRunning),
		TotalItems:            len(req.Entries),
	}); err != nil {
		return BatchImportRunResult{}, err
	}

	result := BatchImportRunResult{
		RunID:   runID,
		ItemIDs: make([]string, 0, len(req.Entries)),
	}

	for idx, entry := range req.Entries {
		itemID := fmt.Sprintf("%s-item-%d", runID, idx+1)
		result.ItemIDs = append(result.ItemIDs, itemID)

		providerID := NormalizeProviderID(entry.BaseURL)
		fingerprint := fingerprintAPIKey(entry.APIKey)
		initialItem := sqlite.ImportRunItem{
			ItemID:              itemID,
			RunID:               runID,
			BaseURL:             strings.TrimSpace(entry.BaseURL),
			ProviderID:          providerID,
			APIKeyFingerprint:   fingerprint,
			CurrentStage:        string(ItemStageProbe),
			ConfirmationStatus:  string(ConfirmationPending),
			AccessStatus:        string(AccessStatusUnknown),
			MatchedAccountState: string(MatchedAccountStateNone),
			AccountResolution:   string(AccountResolutionCreated),
		}
		if err := s.ItemStore.Upsert(ctx, initialItem); err != nil {
			return BatchImportRunResult{}, err
		}

		modelsResult, err := s.ProbeModels(ctx, entry.BaseURL, entry.APIKey)
		if err != nil {
			if failErr := s.failRun(ctx, req, initialItem, ItemStageProbe, err); failErr != nil {
				return BatchImportRunResult{}, failErr
			}
			return result, nil
		}
		rawModels := append([]string(nil), modelsResult.RawModels...)
		capabilityProfile, err := s.ProbeCapabilities(ctx, entry.BaseURL, entry.APIKey, rawModels)
		if err != nil {
			if failErr := s.failRun(ctx, req, initialItem, ItemStageProbe, err); failErr != nil {
				return BatchImportRunResult{}, failErr
			}
			return result, nil
		}
		routingStrategy := BuildImportRoutingStrategy(capabilityProfile)
		resolvedSmokeModel, recommendedModels, err := probe.ResolveSmokeModel(entry.RequestedModels, rawModels, capabilityProfile)
		if err != nil {
			if failErr := s.failRun(ctx, req, initialItem, ItemStageProbe, err); failErr != nil {
				return BatchImportRunResult{}, failErr
			}
			return result, nil
		}

		canonicalFamilies := uniqueCanonicalFamilies(rawModels)
		reuseLookup := ReuseLookupResult{}
		if s.InspectReuse != nil {
			reuseLookup, err = s.InspectReuse(ctx, ReuseLookupInput{
				HostID:                 strings.TrimSpace(req.HostID),
				ProviderID:             providerID,
				BaseURL:                entry.BaseURL,
				APIKeyFingerprint:      fingerprint,
				CanonicalModelFamilies: canonicalFamilies,
			})
			if err != nil {
				if failErr := s.failRun(ctx, req, initialItem, ItemStageProbe, err); failErr != nil {
					return BatchImportRunResult{}, failErr
				}
				return result, nil
			}
		}

		reuseDecision := DecideReuse(ReuseInput{
			ProviderMatched:          reuseLookup.ProviderMatched,
			ProviderID:               providerID,
			CanonicalModelFamilies:   canonicalFamilies,
			MatchedAccountID:         reuseLookup.MatchedAccountID,
			MatchedAccountState:      reuseLookup.MatchedAccountState,
			ExistingProviderID:       reuseLookup.ExistingProviderID,
			ExistingAccessStatus:     reuseLookup.ExistingAccessStatus,
			ExistingCanonicalFamilys: reuseLookup.ExistingCanonicalFamilys,
		})

		finalItem := sqlite.ImportRunItem{
			ItemID:                itemID,
			RunID:                 runID,
			BaseURL:               strings.TrimSpace(entry.BaseURL),
			ProviderID:            providerID,
			APIKeyFingerprint:     fingerprint,
			RequestedModelsJSON:   mustMarshalJSON(entry.RequestedModels, "[]"),
			RawModelsJSON:         mustMarshalJSON(rawModels, "[]"),
			NormalizedModelsJSON:  mustMarshalJSON(uniqueNormalizedModels(rawModels), "[]"),
			CanonicalFamiliesJSON: mustMarshalJSON(canonicalFamilies, "[]"),
			RecommendedModelsJSON: mustMarshalJSON(recommendedModels, "[]"),
			ResolvedSmokeModel:    resolvedSmokeModel,
			CapabilityProfileJSON: mustMarshalJSON(capabilityProfile, "{}"),
			CurrentStage:          string(ItemStageConfirm),
			ConfirmationStatus:    string(ConfirmationPending),
			AccessStatus:          string(AccessStatusUnknown),
			MatchedAccountState:   string(reuseDecision.MatchedAccountState),
			AccountResolution:     string(reuseDecision.AccountResolution),
			ProvisionReused:       reuseDecision.ProvisionReused,
			ReusedFromProviderID:  reuseDecision.ReusedFromProviderID,
			ReusedFromAccountID:   int64PtrIfSet(reuseDecision.ReusedFromAccountID),
			LegacyBatchID:         reuseLookup.LegacyBatchID,
			LegacyProviderID:      strings.TrimSpace(reuseLookup.ExistingProviderID),
		}

		if reuseDecision.ProvisionReused {
			patchContract := ModelMappingDelta(reuseLookup.ExistingModelMapping, probe.BuildAliasTable(rawModels))
			if shouldPatchAliases(reuseLookup.ExistingModelMapping, patchContract.ModelMapping) {
				if s.Provisioner == nil {
					if failErr := s.failRun(ctx, req, initialItem, ItemStageProvision, fmt.Errorf("provisioner is required for patch-only flow")); failErr != nil {
						return BatchImportRunResult{}, failErr
					}
					return result, nil
				}
				if err := s.Provisioner.Patch(ctx, PatchProvisionRequest{
					ProviderID: reuseDecision.ReusedFromProviderID,
					Contract:   patchContract,
				}); err != nil {
					if failErr := s.failRun(ctx, req, initialItem, ItemStageProvision, err); failErr != nil {
						return BatchImportRunResult{}, failErr
					}
					return result, nil
				}
			}
		} else {
			if s.Provisioner == nil {
				if failErr := s.failRun(ctx, req, initialItem, ItemStageProvision, fmt.Errorf("provisioner is required")); failErr != nil {
					return BatchImportRunResult{}, failErr
				}
				return result, nil
			}
			provisionResult, err := s.Provisioner.Provision(ctx, ProvisionRequest{
				RunID:             runID,
				ItemID:            itemID,
				Entry:             entry,
				ProviderID:        providerID,
				ResolvedModel:     resolvedSmokeModel,
				RoutingStrategy:   routingStrategy,
				CapabilityProfile: capabilityProfile,
			})
			if err != nil {
				if failErr := s.failRun(ctx, req, initialItem, ItemStageProvision, err); failErr != nil {
					return BatchImportRunResult{}, failErr
				}
				return result, nil
			}
			finalItem.LegacyBatchID = provisionResult.LegacyBatchID
			finalItem.LegacyProviderID = strings.TrimSpace(provisionResult.LegacyProviderID)
		}

		if err := s.ItemStore.Upsert(ctx, finalItem); err != nil {
			return BatchImportRunResult{}, err
		}
	}

	return result, nil
}

func (s BatchImportService) failRun(ctx context.Context, req BatchImportRunRequest, item sqlite.ImportRunItem, stage ItemStage, cause error) error {
	item.CurrentStage = string(ItemStageDone)
	item.ConfirmationStatus = string(ConfirmationFailed)
	item.AccessStatus = string(AccessStatusBroken)
	item.LastErrorStage = string(stage)
	item.LastError = strings.TrimSpace(cause.Error())
	item.LeaseOwner = ""
	item.LeaseUntil = ""
	item.NextRetryAt = ""
	if err := s.ItemStore.Upsert(ctx, item); err != nil {
		return err
	}

	return s.RunStore.Update(ctx, sqlite.ImportRun{
		RunID:                 strings.TrimSpace(req.RunID),
		HostID:                strings.TrimSpace(req.HostID),
		Mode:                  strings.TrimSpace(req.Mode),
		AccessMode:            strings.TrimSpace(req.AccessMode),
		SubscriptionUsersJSON: mustMarshalJSON(req.SubscriptionUsers, "[]"),
		SubscriptionDays:      req.SubscriptionDays,
		ProbeAPIKey:           strings.TrimSpace(req.ProbeAPIKey),
		State:                 string(RunStateFailed),
		TotalItems:            len(req.Entries),
		CompletedItems:        1,
		BrokenItems:           1,
		FinishedAt:            time.Now().UTC().Format(time.RFC3339),
	})
}

func uniqueCanonicalFamilies(rawModels []string) []string {
	seen := make(map[string]struct{}, len(rawModels))
	families := make([]string, 0, len(rawModels))
	for _, rawModel := range rawModels {
		family := probe.CanonicalModelFamily(rawModel)
		if family == "" {
			continue
		}
		if _, ok := seen[family]; ok {
			continue
		}
		seen[family] = struct{}{}
		families = append(families, family)
	}
	return families
}

func uniqueNormalizedModels(rawModels []string) []string {
	seen := make(map[string]struct{}, len(rawModels))
	models := make([]string, 0, len(rawModels))
	for _, rawModel := range rawModels {
		normalized := probe.NormalizeModelID(rawModel)
		if normalized == "" {
			continue
		}
		if _, ok := seen[normalized]; ok {
			continue
		}
		seen[normalized] = struct{}{}
		models = append(models, normalized)
	}
	return models
}

func mustMarshalJSON(value any, fallback string) string {
	payload, err := json.Marshal(value)
	if err != nil {
		return fallback
	}
	return string(payload)
}

func fingerprintAPIKey(apiKey string) string {
	trimmed := strings.TrimSpace(apiKey)
	if trimmed == "" {
		return ""
	}
	sum := sha256.Sum256([]byte(trimmed))
	return fmt.Sprintf("sha256:%x", sum[:8])
}

func int64PtrIfSet(value int64) *int64 {
	if value == 0 {
		return nil
	}
	result := value
	return &result
}

func shouldPatchAliases(existing map[string]string, next map[string]string) bool {
	if len(existing) == 0 {
		return false
	}
	for key, value := range next {
		if existingValue, ok := existing[key]; !ok || existingValue != value {
			return true
		}
	}
	return false
}
