src/crayon/resources.py

"""
Crayon Resources Module.
Manages atomic building and streaming for Vocabulary Profiles.
"""
import os
import json
import shutil
import logging
import csv
from pathlib import Path
from typing import Iterator, List, Optional
from itertools import chain

from .core.profiles import VocabProfile, PROFILES

# Configure module logger
logger = logging.getLogger(__name__)

# Optional imports
try:
    import requests
    _REQUESTS_AVAILABLE = True
except ImportError:
    _REQUESTS_AVAILABLE = False

try:
    from datasets import load_dataset
    _HF_AVAILABLE = True
except ImportError:
    _HF_AVAILABLE = False


# ============================================================================
# Profile Streaming and Caching
# ============================================================================

# Cache Configuration
CACHE_DIR = Path.home() / ".cache" / "xerv" / "crayon" / "profiles"

def get_profile_path(profile: VocabProfile) -> Path:
    """Returns versioned path: ~/.cache/.../vocab_science_v1.json"""
    return CACHE_DIR / f"vocab_{profile.name}_{profile.version}.json"

def yield_profile_stream(profile: VocabProfile, prefer_local_only: bool = False) -> Iterator[str]:
    """
    Resilient Streamer: Iterates through sources. 
    1. Checks for local sample/bootstrap corpus first.
    2. Streams from Hugging Face if available (unless prefer_local_only=True).
    """
    # 1. Local Bootstrap Corpus (Seamless Offline Fallback)
    # Checks for resources/science_corpus.txt, resources/code_corpus.txt, etc.
    # The convention is resources/{profile_name}_corpus.txt
    local_corpus_path = RESOURCE_DIR / f"{profile.name}_corpus.txt"
    has_local = False
    
    if local_corpus_path.exists():
        logger.info(f"[Sources] Found local bootstrap corpus: {local_corpus_path}")
        has_local = True
        try:
            with open(local_corpus_path, 'r', encoding='utf-8') as f:
                for line in f:
                    if line.strip():
                        yield line.strip()
        except Exception as e:
            logger.warning(f"Failed to read local corpus {local_corpus_path}: {e}")
            
    # Also support specific overrides
    if profile.name == "lite":
        # Lite profile always includes Shakespeare & RainDrop from local if present
        yield from yield_local_resources()
        has_local = True

    # If we want to force local usage and we found local data, skip remote
    if prefer_local_only and has_local:
        logger.info(f"[Mode] Skipping remote sources for {profile.name} (Local-Only Build)")
        return

    # 2. Hugging Face Sources
    if not _HF_AVAILABLE:
        logger.info("HuggingFace 'datasets' not installed. Skipping remote sources.")
        return

    for ds_name, split, cols in profile.sources:
        try:
            logger.info(f"[Stream] Connecting to {ds_name}...")
            
            # Special handling for wikitext which requires a config name
            load_args = [ds_name]
            if ds_name == "wikitext":
                load_args.append("wikitext-103-v1")
                
            # Try loading with trust_remote_code=True first
            try:
                ds = load_dataset(*load_args, split=split, streaming=True, trust_remote_code=True)
            except Exception:
                # Fallback without trust_remote_code (some datasets forbid it)
                ds = load_dataset(*load_args, split=split, streaming=True, trust_remote_code=False)
            
            # Safety Cap: Process max 100k rows per source to prevent infinite hangs
            sample_count = 0
            for row in ds:
                if sample_count >= 100000: 
                    break 
                
                for col in cols:
                    val = row.get(col)
                    if isinstance(val, str): 
                        yield val
                    elif isinstance(val, list): 
                        # Handle list of strings (e.g. sentences)
                        yield " ".join(str(x) for x in val)
                
                sample_count += 1
                    
        except Exception as e:
            logger.warning(f"[Stream Warning] Failed to stream {ds_name}: {e}. Skipping source.")

def build_and_cache_profile(profile_name: str, prefer_local_only: bool = False) -> Path:
    """
    The Production Builder.
    1. Validates profile.
    2. Streams data (Zero-Disk).
    3. Trains entropy model.
    4. ATOMIC WRITE (Write tmp -> Rename) to prevent corruption.
    """
    # Lazy import to prevent circular dependency
    from .training import train_vocabulary 
    
    profile = PROFILES.get(profile_name)
    if not profile:
        raise ValueError(f"Unknown profile: '{profile_name}'. Available: {list(PROFILES.keys())}")

    target_path = get_profile_path(profile)
    
    # Fast Path: Return if already exists
    if target_path.exists():
        return target_path

    logger.info(f"--- BUILDING PROFILE: {profile.name.upper()} ---")
    logger.info(f"Target Size: {profile.target_size} | Sources: {len(profile.sources)}")
    
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    
    # 1. Train
    stream = yield_profile_stream(profile, prefer_local_only=prefer_local_only)
    
    # If HF is not available or stream yields nothing, we might crash training.
    # But train_vocabulary handles iterators.
    vocab_list = train_vocabulary(
        stream, 
        target_size=profile.target_size,
        min_frequency=profile.min_frequency
    )
    
    # 2. Atomic Write Pattern
    temp_path = target_path.with_suffix(".tmp")
    try:
        with open(temp_path, 'w', encoding='utf-8') as f:
            json.dump(vocab_list, f, indent=2)
        
        # Instant rename (Atomic)
        shutil.move(str(temp_path), str(target_path))
        logger.info(f"[Success] Saved profile to: {target_path}")
        
    except Exception as e:
        if temp_path.exists(): 
            os.remove(temp_path)
        raise RuntimeError(f"Failed to save profile: {e}")
        
    return target_path


# ============================================================================
# Local Resource Iterators (Legacy / Fallback support)
# ============================================================================

RESOURCE_DIR = Path(__file__).parent / "resources"

def yield_local_resources(max_grad_entries: int = 5000) -> Iterator[str]:
    """
    Yields text from local resource files if they exist.
    """
    if not RESOURCE_DIR.exists():
        return

    # 1. Shakespeare
    shakespeare_path = RESOURCE_DIR / "input.txt"
    if shakespeare_path.exists():
        logger.info(f"Using local Shakespeare: {shakespeare_path}")
        try:
            with open(shakespeare_path, 'r', encoding='utf-8') as f:
                for line in f:
                    if line.strip():
                        yield line.strip()
        except Exception as e:
            logger.warning(f"Error reading local Shakespeare: {e}")

def get_default_corpus_iterator(
    include_shakespeare: bool = True,
    include_hf_sources: bool = True, # Ignored in legacy shim
    include_builtin: bool = True,
    max_hf_samples: Optional[int] = None
) -> Iterator[str]:
    """
    Legacy shim: Returns an iterator over 'lite' profile resources or local.
    """
    # Prefer local resources first
    local_iter = yield_local_resources()
    
    # If no local resources, try to stream 'lite' profile if HF available
    if _HF_AVAILABLE:
        lite_profile = PROFILES.get("lite")
        if lite_profile:
            return chain(local_iter, yield_profile_stream(lite_profile))
            
    return local_iter

def check_resource_availability() -> dict:
    """Check which data sources are available."""
    local_files = [f.name for f in RESOURCE_DIR.iterdir()] if RESOURCE_DIR.exists() else []
    
    return {
        "requests_available": _REQUESTS_AVAILABLE,
        "huggingface_available": _HF_AVAILABLE,
        "local_resources_dir": str(RESOURCE_DIR),
        "local_files": local_files,
        "builtin_available": True
    }