package usecase

import (
	"context"
	"time"

	"github.com/google/uuid"

	"github.com/emil28092005/SciMesh/coordinator/internal/domain"
	"github.com/emil28092005/SciMesh/coordinator/internal/workloads"
)

// Job operations: the submitter-facing lifecycle of a whole submission.
//
//	CreateJob    register a job and fan it out into tasks
//	GetJobStatus aggregate progress
//	ListResults  completed manifests, ordered for the stitcher
//	StitchJob    merge partial results into the final artifact

// --- CreateJob -----------------------------------------------------------

type CreateJob struct {
	jobs  JobRepository
	tasks TaskRepository
	tx    TxManager
	clock Clock
}

func NewCreateJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock Clock) *CreateJob {
	return &CreateJob{jobs: jobs, tasks: tasks, tx: tx, clock: clock}
}

// Execute builds the job and its tasks, then writes them in one transaction.
// The all-or-none guarantee comes from TxManager: a half-created job would
// leave chunks no worker could ever complete.
func (uc *CreateJob) Execute(ctx context.Context, in CreateJobInput) (*domain.Job, error) {
	if in.Workload == "similarity-graph" || in.Workload == "similarity_graph" {
		// CTX-10 must plan triangular block pairs; ordinary independent input
		// chunks would silently omit every cross-chunk molecular pair.
		return nil, domain.ErrInvalidInput
	}
	if (in.Workload == "similarity-search" || in.Workload == "similarity_search") &&
		len(in.Chunks) > 1 && in.Parameters["query_id"] != nil {
		// Resolving once against the source dataset belongs to CTX-07. Letting
		// each shard resolve it would make most tasks fail or use inconsistent data.
		return nil, domain.ErrInvalidInput
	}
	chunks := make([]domain.ChunkSpec, 0, len(in.Chunks))
	for _, c := range in.Chunks {
		chunks = append(chunks, domain.ChunkSpec(c))
	}

	job, tasks, err := domain.NewJobWithTasks(in.Workload, in.InputURI, in.Parameters, chunks, uc.clock.Now())
	if err != nil {
		return nil, err
	}
	job.OwnerID = ownerFromContext(ctx)

	err = uc.tx.WithinTx(ctx, func(ctx context.Context) error {
		if err := uc.jobs.Insert(ctx, job); err != nil {
			return err
		}
		return uc.tasks.InsertBatch(ctx, tasks)
	})
	if err != nil {
		return nil, err
	}
	return job, nil
}

// --- GetJobStatus --------------------------------------------------------

type GetJobStatus struct {
	jobs  JobRepository
	tasks TaskRepository
}

// --- CancelJob -----------------------------------------------------------

type CancelJob struct {
	jobs  JobRepository
	tasks TaskRepository
	tx    TxManager
	clock Clock
}

func NewCancelJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock Clock) *CancelJob {
	return &CancelJob{jobs: jobs, tasks: tasks, tx: tx, clock: clock}
}

// Execute stops a job atomically. Completed and finally failed tasks are kept
// as historical evidence; all other tasks are cancelled, including leased and
// running ones. A repeated cancel of an already cancelled job is idempotent.
func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error) {
	now := uc.clock.Now()
	var cancelled int64
	err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
		job, err := uc.jobs.Get(ctx, jobID)
		if err != nil {
			return err
		}
		if err := authorizeJobAccess(ctx, job); err != nil {
			return err
		}
		if job.Status == domain.JobCancelled {
			return nil
		}
		if job.Status == domain.JobReducing || job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
			return domain.ErrJobNotCancellable
		}
		// The lease reaper can be the transition that exhausted the final task.
		// Check the authoritative task histogram as well as the cached job status,
		// so a stale status can never turn a failed/completed job into cancelled.
		counts, err := uc.tasks.CountByStatus(ctx, jobID)
		if err != nil {
			return err
		}
		derived := progressFrom(*job, counts).DeriveStatus()
		if derived == domain.JobReducing || derived == domain.JobCompleted || derived == domain.JobFailed {
			return domain.ErrJobNotCancellable
		}
		cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now)
		if err != nil {
			return err
		}
		return uc.jobs.UpdateStatus(ctx, jobID, domain.JobCancelled, &now)
	})
	return cancelled, err
}

func NewGetJobStatus(jobs JobRepository, tasks TaskRepository) *GetJobStatus {
	return &GetJobStatus{jobs: jobs, tasks: tasks}
}

func (uc *GetJobStatus) Execute(ctx context.Context, jobID uuid.UUID) (domain.JobProgress, error) {
	job, err := uc.jobs.Get(ctx, jobID)
	if err != nil {
		return domain.JobProgress{}, err
	}
	if err := authorizeJobAccess(ctx, job); err != nil {
		return domain.JobProgress{}, err
	}
	counts, err := uc.tasks.CountByStatus(ctx, jobID)
	if err != nil {
		return domain.JobProgress{}, err
	}
	return progressFrom(*job, counts), nil
}

// --- ListResults ---------------------------------------------------------

type ListResults struct {
	tasks TaskRepository
}

func NewListResults(tasks TaskRepository) *ListResults {
	return &ListResults{tasks: tasks}
}

// Execute preserves chunk_index order: the stitcher merges these into one
// artifact, and a non-deterministic order would make the final result depend on
// which worker happened to finish first.
func (uc *ListResults) Execute(ctx context.Context, jobID uuid.UUID) ([]domain.ResultManifest, error) {
	tasks, err := uc.tasks.ListCompleted(ctx, jobID)
	if err != nil {
		return nil, err
	}

	manifests := make([]domain.ResultManifest, 0, len(tasks))
	for _, t := range tasks {
		if t.ResultArtifactID == nil {
			continue // a completed task always references its result; skip defensively
		}
		manifests = append(manifests, domain.ResultManifest{
			TaskID:           t.ID,
			ChunkIndex:       t.ChunkIndex,
			ResultArtifactID: *t.ResultArtifactID,
			Metrics:          t.Metrics,
		})
	}
	return manifests, nil
}

// --- StitchJob -----------------------------------------------------------

// StitchJob merges every chunk's partial result into the job's final artifact.
// For similarity search that means concatenating each worker's local top-k,
// sorting by similarity, and keeping the global top-k — the distributed result
// must match what a single local run would produce.
type StitchJob struct {
	results *ListResults
}

func NewStitchJob(results *ListResults) *StitchJob {
	return &StitchJob{results: results}
}

// Execute returns the URI of the assembled artifact.
//
// TODO(phase 6): fetch each manifest's CSV, merge, and persist the result.
func (uc *StitchJob) Execute(ctx context.Context, jobID uuid.UUID) (string, error) {
	if _, err := uc.results.Execute(ctx, jobID); err != nil {
		return "", err
	}
	return "", ErrNotImplemented
}

// --- shared helpers ------------------------------------------------------

// progressFrom turns a status histogram into the domain's progress view.
func progressFrom(job domain.Job, counts map[domain.TaskStatus]int) domain.JobProgress {
	p := domain.JobProgress{
		Job:     job,
		Pending: counts[domain.TaskPending],
		// Leased and running are both "in flight" for progress purposes.
		Leased:    counts[domain.TaskLeased] + counts[domain.TaskRunning],
		Done:      counts[domain.TaskCompleted],
		Failed:    counts[domain.TaskFailed],
		Cancelled: counts[domain.TaskCancelled],
	}
	for _, n := range counts {
		p.Total += n
	}
	return p
}

// syncJobStatus recomputes a job's status from its task counts and persists it.
// Shared by CompleteTask and FailTask so both close a job by the same rule —
// the rule itself lives in domain.JobProgress.DeriveStatus.
func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository,
	catalog *workloads.Catalog, jobID uuid.UUID, now time.Time) error {

	counts, err := tasks.CountByStatus(ctx, jobID)
	if err != nil {
		return err
	}
	job, err := jobs.Get(ctx, jobID)
	if err != nil {
		return err
	}
	status := progressFrom(*job, counts).DeriveStatus()
	// All worker shards being complete means scientific reduction is ready, not
	// that the job's final artifact already exists. Known catalog workloads
	// transition to reducing so ReduceJob can produce the final artifact; the
	// reducer then completes the job with the result. Unknown (URI-based) jobs
	// complete without a coordinator-owned final artifact.
	if status == domain.JobCompleted && catalog != nil && catalog.Reduction(job.Workload) != "" {
		status = domain.JobReducing
	}

	var completedAt *time.Time
	if status == domain.JobFailed {
		completedAt = &now
	}
	return jobs.UpdateStatus(ctx, jobID, status, completedAt)
}

func syncExpiredJobStatuses(ctx context.Context, jobs JobRepository, tasks TaskRepository,
	catalog *workloads.Catalog, jobIDs []uuid.UUID, now time.Time) error {
	seen := make(map[uuid.UUID]struct{}, len(jobIDs))
	for _, jobID := range jobIDs {
		if _, duplicate := seen[jobID]; duplicate {
			continue
		}
		seen[jobID] = struct{}{}
		if err := syncJobStatus(ctx, jobs, tasks, catalog, jobID, now); err != nil {
			return err
		}
	}
	return nil
}