tm_back/internal/company/repository.go

package company

import (
	"context"
	"errors"
	"regexp"
	"strings"
	"time"
	"tm/pkg/logger"
	orm "tm/pkg/mongo"
	"tm/pkg/response"

	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
)

// Repository defines the interface for company data operations
type Repository interface {
	Create(ctx context.Context, company *Company) error
	Update(ctx context.Context, company *Company) error
	GetByID(ctx context.Context, id string) (*Company, error)
	GetByIDs(ctx context.Context, ids []string) ([]Company, error)
	GetByName(ctx context.Context, name string) (*Company, error)
	GetByRegistrationNumber(ctx context.Context, registrationNumber string) (*Company, error)
	GetByTaxID(ctx context.Context, taxID string) (*Company, error)
	Search(ctx context.Context, form *SearchForm, pagination *response.Pagination) (*orm.PaginatedResult[Company], error)
	ListIDs(ctx context.Context) ([]string, error)
	Delete(ctx context.Context, id string) error
	UpdateStatus(ctx context.Context, id string, status CompanyStatus) error
	UpdateVerification(ctx context.Context, id string, isVerified, isCompliant bool, complianceNotes *string) error
	RemoveDocumentFileID(ctx context.Context, fileID string) error
}

// companyRepository implements the Repository interface using the MongoDB ORM
type companyRepository struct {
	ormRepo    orm.Repository[Company]
	collection *mongo.Collection
	logger     logger.Logger
}

// NewCompanyRepository creates a new company repository
func NewRepository(mongoManager *orm.ConnectionManager, logger logger.Logger) Repository {
	// Create indexes using the ORM's index management
	indexes := []orm.Index{
		*orm.NewIndex("registration_number_idx", bson.D{{Key: "registration_number", Value: 1}}),
	}

	// Create indexes
	err := mongoManager.CreateIndexes("companies", indexes)
	if err != nil {
		logger.Warn("Failed to create company indexes", map[string]interface{}{
			"error": err.Error(),
		})
	}

	collection := mongoManager.GetCollection("companies")

	// Create ORM repository
	ormRepo := orm.NewRepository[Company](collection, logger)

	return &companyRepository{
		ormRepo:    ormRepo,
		collection: collection,
		logger:     logger,
	}
}

func mapCompanySortField(sortBy string) string {
	switch sortBy {
	case "country":
		return "address.country"
	case "state":
		return "address.state"
	default:
		return sortBy
	}
}

// Create creates a new company
func (r *companyRepository) Create(ctx context.Context, company *Company) error {
	// Set created/updated timestamps using Unix timestamps
	now := time.Now().Unix()
	company.SetCreatedAt(now)

	// Use ORM to create company
	err := r.ormRepo.Create(ctx, company)
	if err != nil {
		r.logger.Error("Failed to create company", map[string]interface{}{
			"error":               err.Error(),
			"name":                company.Name,
			"registration_number": company.RegistrationNumber,
		})
		return err
	}

	r.logger.Info("Company created successfully", map[string]interface{}{
		"company_id": company.ID,
		"name":       company.Name,
		"type":       company.Type,
	})

	return nil
}

// GetByID retrieves a company by ID
func (r *companyRepository) GetByID(ctx context.Context, id string) (*Company, error) {
	company, err := r.ormRepo.FindByID(ctx, id)
	if err != nil {
		if errors.Is(err, orm.ErrDocumentNotFound) {
			return nil, errors.New("company not found")
		}
		r.logger.Error("Failed to get company by ID", map[string]interface{}{
			"error":      err.Error(),
			"company_id": id,
		})
		return nil, err
	}

	return company, nil
}

// GetByIDs retrieves companies by IDs
func (r *companyRepository) GetByIDs(ctx context.Context, ids []string) ([]Company, error) {
	objectIDs := make([]bson.ObjectID, len(ids))
	for i, id := range ids {
		objectID, err := bson.ObjectIDFromHex(id)
		if err != nil {
			return nil, err
		}
		objectIDs[i] = objectID
	}

	companies, err := r.ormRepo.FindAll(ctx, bson.M{"_id": bson.M{"$in": objectIDs}}, orm.NewPaginationBuilder().Build())
	if err != nil {
		return nil, err
	}

	return companies.Items, nil
}

// GetByName retrieves a company by name
func (r *companyRepository) GetByName(ctx context.Context, name string) (*Company, error) {
	company, err := r.ormRepo.FindOne(ctx, bson.M{"name": name})
	if err != nil {
		return nil, err
	}
	return company, nil
}

// GetByRegistrationNumber retrieves a company by registration number
func (r *companyRepository) GetByRegistrationNumber(ctx context.Context, registrationNumber string) (*Company, error) {
	company, err := r.ormRepo.FindOne(ctx, bson.M{"registration_number": registrationNumber})
	if err != nil {
		return nil, err
	}
	return company, nil
}

// GetByTaxID retrieves a company by tax ID
func (r *companyRepository) GetByTaxID(ctx context.Context, taxID string) (*Company, error) {
	company, err := r.ormRepo.FindOne(ctx, bson.M{"tax_id": taxID})
	if err != nil {
		return nil, err
	}
	return company, nil
}

// Update updates a company
func (r *companyRepository) Update(ctx context.Context, company *Company) error {
	// Set updated timestamp using Unix timestamp
	company.SetUpdatedAt(time.Now().Unix())

	clearDocuments := company.DocumentFileIDs != nil && len(company.DocumentFileIDs) == 0
	clearLinks := company.Links != nil && len(company.Links) == 0

	// Use ORM to update company
	err := r.ormRepo.Update(ctx, company)
	if err != nil {
		r.logger.Error("Failed to update company", map[string]interface{}{
			"error":      err.Error(),
			"company_id": company.ID,
		})
		return err
	}

	// BSON omitempty can skip empty slices in $set, so persist explicit empty arrays.
	if clearDocuments || clearLinks {
		setFields := bson.M{
			"updated_at": company.UpdatedAt,
		}
		if clearDocuments {
			setFields["document_file_ids"] = []string{}
		}
		if clearLinks {
			setFields["links"] = []CompanyLink{}
		}

		_, err = r.collection.UpdateOne(ctx, bson.M{"_id": company.ID}, bson.M{
			"$set": setFields,
		})
		if err != nil {
			r.logger.Error("Failed to clear company collection fields", map[string]interface{}{
				"error":      err.Error(),
				"company_id": company.ID,
			})
			return err
		}
	}

	r.logger.Info("Company updated successfully", map[string]interface{}{
		"company_id": company.ID,
		"name":       company.Name,
	})

	return nil
}

// Delete deletes a company
func (r *companyRepository) Delete(ctx context.Context, id string) error {
	// Get company first
	_, err := r.GetByID(ctx, id)
	if err != nil {
		return err
	}

	// Update in database
	err = r.ormRepo.Delete(ctx, id)
	if err != nil {
		r.logger.Error("Failed to delete company", map[string]interface{}{
			"error":      err.Error(),
			"company_id": id,
		})
		return err
	}

	r.logger.Info("Company deleted successfully", map[string]interface{}{
		"company_id": id,
	})

	return nil
}

// Search retrieves companies with search and filters
func (r *companyRepository) Search(ctx context.Context, form *SearchForm, pagination *response.Pagination) (*orm.PaginatedResult[Company], error) {
	// Build filter
	filter := bson.M{}

	if form.Search != nil {
		filter["$text"] = bson.M{"$search": *form.Search}
	}

	if form.Name != nil {
		name := strings.TrimSpace(*form.Name)
		if name != "" {
			filter["name"] = bson.M{
				"$regex":   regexp.QuoteMeta(name),
				"$options": "i",
			}
		}
	}

	if form.Email != nil {
		email := strings.TrimSpace(*form.Email)
		if email != "" {
			filter["email"] = bson.M{
				"$regex":   regexp.QuoteMeta(email),
				"$options": "i",
			}
		}
	}

	if form.Phone != nil {
		phone := strings.TrimSpace(*form.Phone)
		if phone != "" {
			filter["phone"] = bson.M{
				"$regex":   regexp.QuoteMeta(phone),
				"$options": "i",
			}
		}
	}

	if form.Type != nil {
		filter["type"] = *form.Type
	}

	if form.Status != nil {
		filter["status"] = *form.Status
	}

	if form.Industry != nil {
		filter["industry"] = *form.Industry
	}

	if form.IsVerified != nil {
		filter["is_verified"] = *form.IsVerified
	}

	if form.IsCompliant != nil {
		filter["is_compliant"] = *form.IsCompliant
	}

	if form.Language != nil {
		filter["language"] = *form.Language
	}

	if form.Currency != nil {
		filter["currency"] = *form.Currency
	}

	// Tag filters
	if len(form.CPVCodes) > 0 {
		filter["tags.cpv_codes"] = bson.M{"$in": form.CPVCodes}
	}

	if len(form.Categories) > 0 {
		filter["tags.categories"] = bson.M{"$in": form.Categories}
	}

	if len(form.Keywords) > 0 {
		filter["tags.keywords"] = bson.M{"$in": form.Keywords}
	}

	if len(form.Specializations) > 0 {
		filter["tags.specializations"] = bson.M{"$in": form.Specializations}
	}

	// Range filters
	if form.EmployeeCount != nil {
		filter["employee_count"] = *form.EmployeeCount
	} else if form.EmployeeCountMin != nil || form.EmployeeCountMax != nil {
		rangeFilter := bson.M{}
		if form.EmployeeCountMin != nil {
			rangeFilter["$gte"] = *form.EmployeeCountMin
		}
		if form.EmployeeCountMax != nil {
			rangeFilter["$lte"] = *form.EmployeeCountMax
		}
		filter["employee_count"] = rangeFilter
	}

	if form.AnnualRevenueMin != nil || form.AnnualRevenueMax != nil {
		rangeFilter := bson.M{}
		if form.AnnualRevenueMin != nil {
			rangeFilter["$gte"] = *form.AnnualRevenueMin
		}
		if form.AnnualRevenueMax != nil {
			rangeFilter["$lte"] = *form.AnnualRevenueMax
		}
		filter["annual_revenue"] = rangeFilter
	}

	if form.FoundedYearMin != nil || form.FoundedYearMax != nil {
		rangeFilter := bson.M{}
		if form.FoundedYearMin != nil {
			rangeFilter["$gte"] = *form.FoundedYearMin
		}
		if form.FoundedYearMax != nil {
			rangeFilter["$lte"] = *form.FoundedYearMax
		}
		filter["founded_year"] = rangeFilter
	}

	mongoPagination, err := orm.BuildListPagination(
		pagination.Limit,
		pagination.Offset,
		pagination.Cursor,
		mapCompanySortField(pagination.SortBy),
		pagination.SortOrder,
		filter,
		orm.ListPaginationOptions{},
	)
	if err != nil {
		return nil, err
	}

	result, err := r.ormRepo.FindAll(ctx, filter, mongoPagination)
	if err != nil {
		r.logger.Error("Failed to search companies", map[string]interface{}{
			"error": err.Error(),
		})
		return nil, err
	}

	return result, nil
}

// ListIDs returns every company document ID.
func (r *companyRepository) ListIDs(ctx context.Context) ([]string, error) {
	const pageSize = 500
	ids := make([]string, 0)
	skip := 0

	for {
		result, err := r.ormRepo.FindAll(ctx, bson.M{}, orm.Pagination{
			Limit:      pageSize,
			Skip:       skip,
			SortField:  "_id",
			SortOrder:  1,
			SkipCount:  true,
			Projection: bson.M{"_id": 1},
		})
		if err != nil {
			r.logger.Error("Failed to list company IDs", map[string]interface{}{
				"error": err.Error(),
				"skip":  skip,
			})
			return nil, err
		}

		if len(result.Items) == 0 {
			break
		}

		for i := range result.Items {
			id := strings.TrimSpace(result.Items[i].GetID())
			if id != "" {
				ids = append(ids, id)
			}
		}

		if len(result.Items) < pageSize {
			break
		}
		skip += pageSize
	}

	return ids, nil
}

// UpdateStatus updates company status
func (r *companyRepository) UpdateStatus(ctx context.Context, id string, status CompanyStatus) error {
	// Get company first
	company, err := r.GetByID(ctx, id)
	if err != nil {
		return err
	}

	// Update status
	company.Status = status
	company.SetUpdatedAt(time.Now().Unix())

	// Update in database
	err = r.ormRepo.Update(ctx, company)
	if err != nil {
		r.logger.Error("Failed to update company status", map[string]interface{}{
			"error":      err.Error(),
			"company_id": id,
			"status":     status,
		})
		return err
	}

	r.logger.Info("Company status updated successfully", map[string]interface{}{
		"company_id": id,
		"status":     status,
	})

	return nil
}

// UpdateVerification updates company verification status
func (r *companyRepository) UpdateVerification(ctx context.Context, id string, isVerified, isCompliant bool, complianceNotes *string) error {
	// Get company first
	company, err := r.GetByID(ctx, id)
	if err != nil {
		return err
	}

	// Update verification fields
	company.IsVerified = isVerified
	company.IsCompliant = isCompliant
	company.ComplianceNotes = complianceNotes
	company.SetUpdatedAt(time.Now().Unix())

	// Update in database
	err = r.ormRepo.Update(ctx, company)
	if err != nil {
		r.logger.Error("Failed to update company verification", map[string]interface{}{
			"error":      err.Error(),
			"company_id": id,
		})
		return err
	}

	r.logger.Info("Company verification updated successfully", map[string]interface{}{
		"company_id":   id,
		"is_verified":  isVerified,
		"is_compliant": isCompliant,
	})

	return nil
}

// RemoveDocumentFileID removes a file ID from all companies that reference it.
func (r *companyRepository) RemoveDocumentFileID(ctx context.Context, fileID string) error {
	if fileID == "" {
		return nil
	}

	filter := bson.M{"document_file_ids": fileID}
	update := bson.M{
		"$pull": bson.M{"document_file_ids": fileID},
		"$set":  bson.M{"updated_at": time.Now().Unix()},
	}

	result, err := r.collection.UpdateMany(ctx, filter, update)
	if err != nil {
		r.logger.Error("Failed to remove document file ID from companies", map[string]interface{}{
			"error":   err.Error(),
			"file_id": fileID,
		})
		return err
	}

	if result.ModifiedCount > 0 {
		r.logger.Info("Removed document file ID from companies", map[string]interface{}{
			"file_id":        fileID,
			"modified_count": result.ModifiedCount,
		})
	}

	return nil
}