package projectmemory

import (
	"context"
	"errors"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/pgvector/pgvector-go"

	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
	projectmemorysqlc "github.com/yan1h/agent-coding-workflow/internal/projectmemory/sqlc"
)

// Repository is the projectmemory persistence contract.
type Repository interface {
	Insert(ctx context.Context, in WriteInput, embeddingModel *string) (*Entry, error)
	SetEmbedding(ctx context.Context, id uuid.UUID, vec []float32, model string) error
	Get(ctx context.Context, id uuid.UUID) (*Entry, error)
	ListByProject(ctx context.Context, projectID uuid.UUID, kind string) ([]Entry, error)
	SearchKeyword(ctx context.Context, projectID uuid.UUID, kind, text string, topK int) ([]Entry, error)
	SearchKNN(ctx context.Context, projectID uuid.UUID, vec []float32, kind string, topK int) ([]Entry, error)
	Delete(ctx context.Context, id uuid.UUID) error
}

// PgRepository is the production Repository on a pgxpool.Pool.
type PgRepository struct {
	pool *pgxpool.Pool
	q    *projectmemorysqlc.Queries
}

// NewPostgresRepository builds a PgRepository.
func NewPostgresRepository(pool *pgxpool.Pool) *PgRepository {
	return &PgRepository{pool: pool, q: projectmemorysqlc.New(pool)}
}

var _ Repository = (*PgRepository)(nil)

func pgUUID(u uuid.UUID) pgtype.UUID { return pgtype.UUID{Bytes: u, Valid: true} }

func pgUUIDPtr(u *uuid.UUID) pgtype.UUID {
	if u == nil {
		return pgtype.UUID{}
	}
	return pgtype.UUID{Bytes: *u, Valid: true}
}

func uuidPtr(p pgtype.UUID) *uuid.UUID {
	if !p.Valid {
		return nil
	}
	id := uuid.UUID(p.Bytes)
	return &id
}

// entryFields is the common projection shared by every sqlc row type in this
// package; a tiny adapter keeps a single conversion path.
type entryFields struct {
	ID             pgtype.UUID
	ProjectID      pgtype.UUID
	WorkspaceID    pgtype.UUID
	Kind           string
	Title          string
	Body           string
	Tags           []string
	Source         string
	EmbeddingModel *string
	CreatedBy      pgtype.UUID
	CreatedAt      pgtype.Timestamptz
	UpdatedAt      pgtype.Timestamptz
}

func toEntry(f entryFields) Entry {
	e := Entry{
		ID:             uuid.UUID(f.ID.Bytes),
		ProjectID:      uuid.UUID(f.ProjectID.Bytes),
		WorkspaceID:    uuidPtr(f.WorkspaceID),
		Kind:           f.Kind,
		Title:          f.Title,
		Body:           f.Body,
		Tags:           f.Tags,
		Source:         f.Source,
		EmbeddingModel: f.EmbeddingModel,
		CreatedBy:      uuidPtr(f.CreatedBy),
	}
	if e.Tags == nil {
		e.Tags = []string{}
	}
	if f.CreatedAt.Valid {
		e.CreatedAt = f.CreatedAt.Time
	}
	if f.UpdatedAt.Valid {
		e.UpdatedAt = f.UpdatedAt.Time
	}
	return e
}

func (r *PgRepository) Insert(ctx context.Context, in WriteInput, embeddingModel *string) (*Entry, error) {
	tags := in.Tags
	if tags == nil {
		// TEXT[] NOT NULL: never write a nil slice (would violate the constraint).
		tags = []string{}
	}
	source := in.Source
	if source == "" {
		source = SourceAgent
	}
	row, err := r.q.InsertEntry(ctx, projectmemorysqlc.InsertEntryParams{
		ID:             pgUUID(uuid.New()),
		ProjectID:      pgUUID(in.ProjectID),
		WorkspaceID:    pgUUIDPtr(in.WorkspaceID),
		Kind:           in.Kind,
		Title:          in.Title,
		Body:           in.Body,
		Tags:           tags,
		Source:         source,
		EmbeddingModel: embeddingModel,
		CreatedBy:      pgUUIDPtr(in.CreatedBy),
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "insert project memory")
	}
	e := toEntry(entryFields(row))
	return &e, nil
}

// SetEmbedding writes the vector + model with a hand-written pgx UPDATE (the
// embedding column is excluded from sqlc). updated_at is bumped.
func (r *PgRepository) SetEmbedding(ctx context.Context, id uuid.UUID, vec []float32, model string) error {
	_, err := r.pool.Exec(ctx,
		`UPDATE project_memory SET embedding = $2::vector, embedding_model = $3, updated_at = now() WHERE id = $1`,
		pgUUID(id), pgvector.NewVector(vec), model)
	if err != nil {
		return errs.Wrap(err, errs.CodeInternal, "set memory embedding")
	}
	return nil
}

func (r *PgRepository) Get(ctx context.Context, id uuid.UUID) (*Entry, error) {
	row, err := r.q.GetEntry(ctx, pgUUID(id))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, errs.New(errs.CodeNotFound, "memory entry not found")
		}
		return nil, errs.Wrap(err, errs.CodeInternal, "get memory entry")
	}
	e := toEntry(entryFields(row))
	return &e, nil
}

func (r *PgRepository) ListByProject(ctx context.Context, projectID uuid.UUID, kind string) ([]Entry, error) {
	rows, err := r.q.ListByProject(ctx, projectmemorysqlc.ListByProjectParams{
		ProjectID: pgUUID(projectID),
		Column2:   kind,
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "list project memory")
	}
	out := make([]Entry, 0, len(rows))
	for _, row := range rows {
		out = append(out, toEntry(entryFields(row)))
	}
	return out, nil
}

func (r *PgRepository) SearchKeyword(ctx context.Context, projectID uuid.UUID, kind, text string, topK int) ([]Entry, error) {
	if topK <= 0 {
		topK = 10
	}
	rows, err := r.q.SearchByKeyword(ctx, projectmemorysqlc.SearchByKeywordParams{
		ProjectID: pgUUID(projectID),
		Column2:   kind,
		Column3:   text,
		Limit:     int32(topK),
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "search memory keyword")
	}
	out := make([]Entry, 0, len(rows))
	for _, row := range rows {
		out = append(out, toEntry(entryFields(row)))
	}
	return out, nil
}

// SearchKNN runs cosine-KNN over entries that have an embedding, scoped to the
// project (and optional kind). Hand-written because the vector column is excluded
// from sqlc.
func (r *PgRepository) SearchKNN(ctx context.Context, projectID uuid.UUID, vec []float32, kind string, topK int) ([]Entry, error) {
	if topK <= 0 {
		topK = 10
	}
	q := `
SELECT id, project_id, workspace_id, kind, title, body, tags, source,
       embedding_model, created_by, created_at, updated_at
FROM project_memory
WHERE project_id = $1
  AND embedding IS NOT NULL
  AND ($3::text = '' OR kind = $3)
ORDER BY embedding <=> $2::vector
LIMIT $4`
	rows, err := r.pool.Query(ctx, q, pgUUID(projectID), pgvector.NewVector(vec), kind, int32(topK))
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "search memory knn")
	}
	defer rows.Close()
	var out []Entry
	for rows.Next() {
		var f entryFields
		if err := rows.Scan(&f.ID, &f.ProjectID, &f.WorkspaceID, &f.Kind, &f.Title,
			&f.Body, &f.Tags, &f.Source, &f.EmbeddingModel, &f.CreatedBy,
			&f.CreatedAt, &f.UpdatedAt); err != nil {
			return nil, errs.Wrap(err, errs.CodeInternal, "scan memory knn row")
		}
		out = append(out, toEntry(f))
	}
	if err := rows.Err(); err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "memory knn rows")
	}
	return out, nil
}

func (r *PgRepository) Delete(ctx context.Context, id uuid.UUID) error {
	if err := r.q.DeleteEntry(ctx, pgUUID(id)); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "delete memory entry")
	}
	return nil
}