"""Reference data management for worker container.

Downloads and caches reference data from S3 for usage stage processing.
"""

import os
from pathlib import Path
import boto3
from botocore.exceptions import ClientError


def ensure_reference_data(s3_bucket: str) -> bool:
    """
    Ensure reference data is available locally, download from S3 if needed.

    Args:
        s3_bucket: S3 bucket name

    Returns:
        True if reference data is available, False otherwise
    """
    reference_dir = Path.home() / ".czbg" / "cache" / "reference_data"
    reference_path = reference_dir / "reference.parquet"

    # Check if already exists
    if reference_path.exists():
        print(f"✓ Reference data already exists: {reference_path}")
        return True

    # Download from S3
    print(f"Downloading reference data from S3...")
    reference_dir.mkdir(parents=True, exist_ok=True)

    s3_key = "czbillgen/reference_data/current.parquet"

    try:
        s3_client = boto3.client('s3')
        s3_client.download_file(s3_bucket, s3_key, str(reference_path))
        print(f"✓ Downloaded reference data: {reference_path}")
        return True

    except ClientError as e:
        print(f"✗ Failed to download reference data: {e}")
        return False
    except Exception as e:
        print(f"✗ Unexpected error downloading reference data: {e}")
        return False


def get_reference_path() -> Path:
    """Get the local reference data path."""
    return Path.home() / ".czbg" / "cache" / "reference_data" / "reference.parquet"
