# task_status_manager.py
from enum import Enum
import random
import string
from threading import Lock
from typing import List, Dict

class TaskStatus(Enum):
    PENDING = "PENDING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"
    PROCESSING= "PROCESSING"
    ACTIVE = "ACTIVE"

# Dictionary to store task status based on content_id
_contents_status = {}
# Dictionary to store random folder names based on content_id
_content_folders = {}
# Dictornary to store task counters... 
_task_counter = {}
_task_counter_lock = Lock()
_document_data = {}

def set_content_status(content_id, status):
    """Set the status for a given content ID."""
    _contents_status[content_id] = status

def get_content_status(content_id):
    """Retrieve the status for a given content ID."""
    return _contents_status.get(content_id, TaskStatus.PENDING.value)


def set_folder_for_content(content_id, length: int = 6) -> str:
    """Set a random folder for a given content ID and return it."""
    characters = string.ascii_letters + string.digits
    folder_name = ''.join(random.choice(characters) for i in range(length))
    _content_folders[content_id] = folder_name
    return folder_name


def get_folder_for_content(content_id) -> str:
    """Retrieve the folder for a given content ID."""
    return _content_folders.get(content_id)

def remove_folder_for_content(content_id) -> None:
    """Remove the folder entry for a given content ID."""
    _content_folders.pop(content_id, None)

def task_count_add(content_id):
    with _task_counter_lock:
        if content_id not in _task_counter:
            _task_counter[content_id] = 0
        _task_counter[content_id] += 1
        return _task_counter[content_id]

def task_count_completed(content_id):
    with _task_counter_lock:
        if content_id not in _task_counter:
            _task_counter[content_id] = 0
        else:
            _task_counter[content_id] -= 1
            _task_counter[content_id] = max(0, _task_counter[content_id])  # Ensure it doesn't go below zero
        return _task_counter[content_id]

def get_task_count(content_id):
    with _task_counter_lock:
        return _task_counter.get(content_id, 0)

def store_document_data(content_id, document):
    """Store document data for a given content ID."""
    if content_id not in _document_data:
        _document_data[content_id] = []
    _document_data[content_id].append(document)

def retrieve_document_data(content_id) -> List[Dict]:
    """Retrieve document data for a given content ID."""
    return _document_data.get(content_id, [])

def clear_document_data(content_id):
    """Clear stored document data for a given content ID."""
    if content_id in _document_data:
        del _document_data[content_id]
