import re

with open('kifi-api/src/main/java/com/kifi/api/service/ProjectService.java', 'r') as f:
    content = f.read()

# Add import
content = content.replace(
    'import com.kifi.api.entity.project.ProjectTask;',
    'import com.kifi.api.entity.project.ProjectTask;\nimport com.kifi.api.entity.project.ProjectTaskComment;\nimport com.kifi.api.repository.ProjectTaskCommentRepository;\nimport io.r2dbc.postgresql.codec.Json;'
)

# Add repository field
content = content.replace(
    'private final ProjectTaskRepository projectTaskRepository;',
    'private final ProjectTaskRepository projectTaskRepository;\n    private final ProjectTaskCommentRepository projectTaskCommentRepository;'
)

# Add comment methods
methods = """
    public Flux<ProjectTaskComment> getCommentsForTask(Long taskId) {
        return projectTaskCommentRepository.findByTaskIdOrderByCreatedAtAsc(taskId);
    }

    public Mono<ProjectTaskComment> addComment(Long taskId, Long userId, String content, String imagesJson) {
        ProjectTaskComment comment = new ProjectTaskComment();
        comment.setTaskId(taskId);
        comment.setUserId(userId);
        comment.setContent(content);
        if (imagesJson != null && !imagesJson.isEmpty()) {
            comment.setImages(Json.of(imagesJson));
        }
        comment.setCreatedAt(LocalDateTime.now());
        return projectTaskCommentRepository.save(comment);
    }
}
"""

content = content.replace('}\n', methods, 1) # Note: this will replace the first `}\n` which might not be EOF if not careful, let's use regex for end of file

content = re.sub(r'}\s*$', methods.strip() + '\n}', content)

with open('kifi-api/src/main/java/com/kifi/api/service/ProjectService.java', 'w') as f:
    f.write(content)