package lake

import (
	"context"
	"errors"
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"
)

type Operation string

const (
	OperationSetupInfo        Operation = "setup_info"
	OperationTestDataSource   Operation = "test_datasource"
	OperationCreateDataSource Operation = "create_datasource"
	OperationGetDataSource    Operation = "get_datasource"
	OperationUpdateDataSource Operation = "update_datasource"
	OperationDeleteDataSource Operation = "delete_datasource"
	OperationCreateTask       Operation = "create_task"
	OperationGetTask          Operation = "get_task"
	OperationUpdateTask       Operation = "update_task"
	OperationStartTask        Operation = "start_task"
	OperationStopTask         Operation = "stop_task"
	OperationDropTask         Operation = "drop_task"
	OperationListTaskRuns     Operation = "list_task_runs"
	OperationGetTaskRun       Operation = "get_task_run"
	OperationObserveTask      Operation = "observe_task"
)

// FakeClient is a concurrency-safe, in-memory Client for model, workflow, and
// worker tests. FailNext supports deterministic fault injection.
type FakeClient struct {
	mu               sync.Mutex
	clock            func() time.Time
	nextDataSourceID int32
	nextTaskID       int32
	dataSources      map[int32]DataSource
	tasks            map[int32]Task
	runs             map[int32][]TaskRun
	failures         map[Operation][]error
	calls            []Operation
	setupInfo        *SetupInfo
}

var _ Client = (*FakeClient)(nil)

func NewFakeClient() *FakeClient {
	return &FakeClient{
		clock:            time.Now,
		nextDataSourceID: 1,
		nextTaskID:       1,
		dataSources:      make(map[int32]DataSource),
		tasks:            make(map[int32]Task),
		runs:             make(map[int32][]TaskRun),
		failures:         make(map[Operation][]error),
	}
}

func (f *FakeClient) SetClock(clock func() time.Time) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if clock == nil {
		f.clock = time.Now
		return
	}
	f.clock = clock
}

func (f *FakeClient) FailNext(operation Operation, err error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.failures[operation] = append(f.failures[operation], err)
}

func (f *FakeClient) Calls() []Operation {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]Operation(nil), f.calls...)
}

func (f *FakeClient) PutDataSource(dataSource DataSource) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.dataSources[dataSource.ID] = cloneDataSource(dataSource)
	if dataSource.ID >= f.nextDataSourceID {
		f.nextDataSourceID = dataSource.ID + 1
	}
}

func (f *FakeClient) PutTask(task Task) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.tasks[task.ID] = cloneTask(task)
	if task.ID >= f.nextTaskID {
		f.nextTaskID = task.ID + 1
	}
}

func (f *FakeClient) PutTaskRun(run TaskRun) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.runs[run.TaskID] = append(f.runs[run.TaskID], cloneTaskRun(run))
}

// SetupInfo answers with the two identities the real Lake reports, so a caller
// that only handles one of them fails here rather than in production.
func (f *FakeClient) SetupInfo(_ context.Context, invocation Invocation) (SetupInfo, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationSetupInfo); err != nil {
		return SetupInfo{}, err
	}
	if err := invocation.validate(); err != nil {
		return SetupInfo{}, err
	}
	if f.setupInfo != nil {
		return *f.setupInfo, nil
	}
	// The HTTP client resolves the organization during the token exchange when
	// the caller does not name it, and Lake fills externalID from that. Mirror
	// it, or a caller that relies on the resolution sees an empty external ID
	// here and a populated one in production.
	externalID := invocation.OrgSlug
	if externalID == "" {
		externalID = invocation.TiDBOrgID
	}
	return SetupInfo{
		ExternalID: externalID,
		TrustPrincipals: []TrustPrincipal{
			{Name: "setup and validation role", ARN: "arn:aws:iam::123456789012:role/fake-webapi", Required: true},
			{Name: "data loading role", ARN: "arn:aws:iam::123456789012:role/fake-tenant", Required: true},
		},
	}, nil
}

// PutSetupInfo overrides what SetupInfo answers.
func (f *FakeClient) PutSetupInfo(info SetupInfo) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.setupInfo = &info
}

func (f *FakeClient) TestDataSource(_ context.Context, invocation Invocation, spec DataSourceSpec) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationTestDataSource); err != nil {
		return err
	}
	return f.validateDataSource(invocation, spec)
}

func (f *FakeClient) CreateDataSource(_ context.Context, invocation Invocation, spec DataSourceSpec) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationCreateDataSource); err != nil {
		return err
	}
	if err := f.validateDataSource(invocation, spec); err != nil {
		return err
	}
	for _, existing := range f.dataSources {
		if existing.Name == spec.Name {
			return &APIError{StatusCode: http.StatusConflict, Code: ErrorCodeAlreadyExists}
		}
	}
	now := f.clock().UTC()
	id := f.nextDataSourceID
	f.nextDataSourceID++
	f.dataSources[id] = DataSource{
		ID: id, Name: spec.Name,
		S3: cloneS3(spec.S3), CreatedAt: now, UpdatedAt: now,
	}
	return nil
}

func (f *FakeClient) GetDataSource(_ context.Context, invocation Invocation, name string) (DataSource, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationGetDataSource); err != nil {
		return DataSource{}, err
	}
	if err := invocation.validate(); err != nil {
		return DataSource{}, err
	}
	for _, dataSource := range f.dataSources {
		if dataSource.Name == name {
			return publicDataSource(dataSource), nil
		}
	}
	return DataSource{}, notFound("datasource " + name)
}

func (f *FakeClient) UpdateDataSource(_ context.Context, invocation Invocation, name string, spec DataSourceSpec) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationUpdateDataSource); err != nil {
		return err
	}
	if err := f.validateDataSource(invocation, spec); err != nil {
		return err
	}
	dataSourceID, ok := f.dataSourceIDByName(name)
	if !ok {
		return notFound("datasource " + name)
	}
	dataSource := f.dataSources[dataSourceID]
	dataSource.Name = spec.Name
	dataSource.S3 = cloneS3(spec.S3)
	dataSource.UpdatedAt = f.clock().UTC()
	f.dataSources[dataSourceID] = dataSource
	return nil
}

func (f *FakeClient) DeleteDataSource(_ context.Context, invocation Invocation, name string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationDeleteDataSource); err != nil {
		return err
	}
	if err := invocation.validate(); err != nil {
		return err
	}
	if dataSourceID, ok := f.dataSourceIDByName(name); ok {
		delete(f.dataSources, dataSourceID)
	}
	return nil
}

func (f *FakeClient) dataSourceIDByName(name string) (int32, bool) {
	for id, dataSource := range f.dataSources {
		if dataSource.Name == name {
			return id, true
		}
	}
	return 0, false
}

func (f *FakeClient) taskIDByName(name string) (int32, bool) {
	for id, task := range f.tasks {
		if task.Name == name {
			return id, true
		}
	}
	return 0, false
}

func notFound(what string) error {
	return &APIError{StatusCode: http.StatusNotFound, Code: ErrorCodeNotFound, Message: what + " does not exist"}
}

func (f *FakeClient) CreateTask(_ context.Context, invocation Invocation, spec TaskSpec) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationCreateTask); err != nil {
		return err
	}
	if err := f.validateTask(invocation, spec); err != nil {
		return err
	}
	dataSourceID, ok := f.dataSourceIDByName(spec.DataSourceName)
	if !ok {
		return notFound("datasource " + spec.DataSourceName)
	}
	for _, existing := range f.tasks {
		if existing.Name == spec.Name {
			return &APIError{StatusCode: http.StatusConflict, Code: ErrorCodeAlreadyExists}
		}
	}
	now := f.clock().UTC()
	id := f.nextTaskID
	f.nextTaskID++
	f.tasks[id] = Task{
		ID: id, Name: spec.Name, DataSourceID: dataSourceID, Warehouse: spec.Warehouse,
		ContinuousIngestion: spec.ContinuousIngestion, Status: TaskStatusStopped,
		TiDB: cloneTiDBTaskSpec(spec.TiDB), CreatedAt: now, UpdatedAt: now,
	}
	return nil
}

func (f *FakeClient) GetTask(_ context.Context, invocation Invocation, name string) (Task, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationGetTask); err != nil {
		return Task{}, err
	}
	if err := invocation.validate(); err != nil {
		return Task{}, err
	}
	taskID, ok := f.taskIDByName(name)
	if !ok {
		return Task{}, notFound("task " + name)
	}
	return cloneTask(f.tasks[taskID]), nil
}

func (f *FakeClient) UpdateTask(_ context.Context, invocation Invocation, name string, spec TaskSpec) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationUpdateTask); err != nil {
		return err
	}
	if err := f.validateTask(invocation, spec); err != nil {
		return err
	}
	taskID, ok := f.taskIDByName(name)
	if !ok {
		return notFound("task " + name)
	}
	// Resolve the datasource the spec names, exactly as the HTTP client does.
	// Reusing the one the task already points at would let a test move a task
	// between datasources and see it succeed, which production cannot do.
	dataSourceID, ok := f.dataSourceIDByName(spec.DataSourceName)
	if !ok {
		return notFound("datasource " + spec.DataSourceName)
	}
	task := f.tasks[taskID]
	task.Name = spec.Name
	task.DataSourceID = dataSourceID
	task.Warehouse = spec.Warehouse
	task.ContinuousIngestion = spec.ContinuousIngestion
	task.TiDB = cloneTiDBTaskSpec(spec.TiDB)
	task.UpdatedAt = f.clock().UTC()
	f.tasks[taskID] = task
	return nil
}

func (f *FakeClient) StartTask(_ context.Context, invocation Invocation, name string) error {
	return f.setTaskStatus(invocation, name, OperationStartTask, TaskStatusRunning, false)
}

func (f *FakeClient) StopTask(_ context.Context, invocation Invocation, name string) error {
	return f.setTaskStatus(invocation, name, OperationStopTask, TaskStatusStopped, false)
}

func (f *FakeClient) DropTask(_ context.Context, invocation Invocation, name string) error {
	return f.setTaskStatus(invocation, name, OperationDropTask, TaskStatusUnknown, true)
}

func (f *FakeClient) setTaskStatus(invocation Invocation, name string, operation Operation, status TaskStatus, drop bool) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(operation); err != nil {
		return err
	}
	if err := invocation.validate(); err != nil {
		return err
	}
	taskID, ok := f.taskIDByName(name)
	if !ok {
		return notFound("task " + name)
	}
	task := f.tasks[taskID]
	if drop {
		delete(f.tasks, taskID)
		delete(f.runs, taskID)
		return nil
	}
	task.Status = status
	task.UpdatedAt = f.clock().UTC()
	f.tasks[taskID] = task
	return nil
}

func (f *FakeClient) ListTaskRuns(_ context.Context, invocation Invocation, name string, options ListTaskRunsOptions) (TaskRunPage, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationListTaskRuns); err != nil {
		return TaskRunPage{}, err
	}
	if err := invocation.validate(); err != nil {
		return TaskRunPage{}, err
	}
	taskID, ok := f.taskIDByName(name)
	if !ok {
		return TaskRunPage{}, notFound("task " + name)
	}
	runs := append([]TaskRun(nil), f.runs[taskID]...)
	sort.SliceStable(runs, func(i, j int) bool { return runs[i].CreatedAt.After(runs[j].CreatedAt) })
	if options.Page <= 0 {
		options.Page = 1
	}
	if options.PerPage <= 0 {
		options.PerPage = 20
	}
	start := int((options.Page - 1) * options.PerPage)
	if start >= len(runs) {
		return TaskRunPage{Total: int32(len(runs))}, nil
	}
	end := start + int(options.PerPage)
	if end > len(runs) {
		end = len(runs)
	}
	page := make([]TaskRun, 0, end-start)
	for _, run := range runs[start:end] {
		page = append(page, cloneTaskRun(run))
	}
	return TaskRunPage{Runs: page, Total: int32(len(runs))}, nil
}

func (f *FakeClient) GetTaskRun(_ context.Context, invocation Invocation, executeID string) (TaskRun, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationGetTaskRun); err != nil {
		return TaskRun{}, err
	}
	if err := invocation.validate(); err != nil {
		return TaskRun{}, err
	}
	for _, runs := range f.runs {
		for _, run := range runs {
			if run.ExecuteID == executeID {
				return cloneTaskRun(run), nil
			}
		}
	}
	return TaskRun{}, &APIError{StatusCode: http.StatusNotFound, Code: ErrorCodeNotFound}
}

func (f *FakeClient) ObserveTask(_ context.Context, invocation Invocation, name string) (TaskObservation, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if err := f.begin(OperationObserveTask); err != nil {
		return TaskObservation{}, err
	}
	if err := invocation.validate(); err != nil {
		return TaskObservation{}, err
	}
	taskID, ok := f.taskIDByName(name)
	if !ok {
		return TaskObservation{}, notFound("task " + name)
	}
	task := f.tasks[taskID]
	observation := TaskObservation{
		TaskName: task.Name, TaskStatus: task.Status, LastResolvedTS: task.LastProgress.LastResolvedTS,
		Progress: task.LastProgress, ObservedAt: f.clock().UTC(),
	}
	if runs := newestRuns(f.runs[taskID]); len(runs) > 0 {
		latest := latestRun(runs)
		observation.LatestRunID = latest.ExecuteID
		observation.LatestRunStatus = latest.Status
		observation.LatestRunMessage = latest.Message
		// Merge exactly as the HTTP client does. A fake that reads only the
		// latest run would let a SnapshotAndCdc test pass while production
		// discards the dumpling run's fullSyncCompleted and backupTSO.
		mergeRunCounters(&observation.Progress, latest.Progress)
		for _, run := range runs {
			mergeTaskFacts(&observation.Progress, run.Progress)
		}
		if observation.Progress.LastResolvedTS > observation.LastResolvedTS {
			observation.LastResolvedTS = observation.Progress.LastResolvedTS
		}
	}
	observation.LastResolvedAt = tsoTime(observation.LastResolvedTS)
	return observation, nil
}

func (f *FakeClient) begin(operation Operation) error {
	f.calls = append(f.calls, operation)
	failures := f.failures[operation]
	if len(failures) == 0 {
		return nil
	}
	err := failures[0]
	f.failures[operation] = failures[1:]
	return err
}

func (f *FakeClient) validateDataSource(invocation Invocation, spec DataSourceSpec) error {
	if err := invocation.validate(); err != nil {
		return err
	}
	if strings.TrimSpace(spec.Name) == "" || strings.TrimSpace(spec.S3.Bucket) == "" || strings.TrimSpace(spec.S3.Region) == "" {
		return errors.New("datasource name, S3 bucket, and S3 region are required")
	}
	if strings.TrimSpace(spec.S3.RoleARN) == "" && (spec.S3.AccessKey.Empty() || spec.S3.SecretKey.Empty()) {
		return errors.New("either role ARN or S3 access key and secret key are required")
	}
	return nil
}

func (f *FakeClient) validateTask(invocation Invocation, spec TaskSpec) error {
	if err := invocation.validate(); err != nil {
		return err
	}
	// Mirror the wire validation: a task with no rule would silently widen to
	// every database under the prefix.
	if err := validateTableRules(spec.TiDB.TableRules); err != nil {
		return err
	}
	if strings.TrimSpace(spec.Name) == "" || strings.TrimSpace(spec.DataSourceName) == "" ||
		strings.TrimSpace(spec.Warehouse) == "" {
		return errors.New("task name, datasource name, and warehouse are required")
	}
	return validateTaskSpec(spec.TiDB)
}

func cloneDataSource(value DataSource) DataSource {
	value.S3 = cloneS3(value.S3)
	return value
}

func publicDataSource(value DataSource) DataSource {
	value = cloneDataSource(value)
	value.S3.AccessKey = Secret{}
	value.S3.SecretKey = Secret{}
	value.S3.ExternalID = Secret{}
	// The queue is not a secret and Lake returns it, so the fake returns it
	// too. Stripping it here would let a callers' spec comparison pass against
	// the fake while never converging against the real service.
	return value
}

func cloneS3(value S3Credentials) S3Credentials {
	return value
}

func cloneTask(value Task) Task {
	value.TiDB = cloneTiDBTaskSpec(value.TiDB)
	return value
}

func cloneTiDBTaskSpec(value TiDBTaskSpec) TiDBTaskSpec {
	value.TableRules = append(TableRules(nil), value.TableRules...)
	value.PollIntervalSeconds = cloneInt32(value.PollIntervalSeconds)
	value.BatchFileCount = cloneInt32(value.BatchFileCount)
	value.MergeIntervalSeconds = cloneInt32(value.MergeIntervalSeconds)
	value.TableParallelism = cloneInt32(value.TableParallelism)
	value.DumplingCSVSkipHeader = cloneInt32(value.DumplingCSVSkipHeader)
	return value
}

func cloneInt32(value *int32) *int32 {
	if value == nil {
		return nil
	}
	cloned := *value
	return &cloned
}

func cloneTaskRun(value TaskRun) TaskRun {
	return value
}