import json
from pathlib import Path
from typing import Any

import psutil
from common_logging import get_logger, log_errors

logger = get_logger(__name__)


@log_errors(logger)
def get_memory_usage() -> float:
    return psutil.virtual_memory().used / 1024 ** 3

@log_errors(logger)
def update_status(job_dir: Path, status: str, progress: float=0.0, metrics: dict[str, Any]=None, error: str=None):
    status_file = job_dir / 'status.json'
    data = {'status': status, 'progress': progress, 'metrics': metrics or {}, 'error': error}
    status_file.write_text(json.dumps(data, indent=2))

@log_errors(logger)
def log_message(job_dir: Path, message: str):
    log_file = job_dir / 'training.log'
    with log_file.open('a') as f:
        f.write(f'{message}\n')