"""Code templates for generated loaders.

Contains Python code templates for dataset loaders.
"""

from __future__ import annotations

# Main loader template
LOADER_TEMPLATE = '''"""Generated loader for {dataset_id}.

This file was generated by `warpdata init {dataset_id}`.
Run `python {filename}` to verify the dataset is accessible.

Usage:
    from {module_name} import get_dataset, stream_batches

    # Get dataset handle
    ds = get_dataset()

    # Stream batches for training
    for batch in stream_batches(batch_size=1024):
        # batch is a dict with column names as keys
        ...
"""

from __future__ import annotations

from typing import Iterator, Any

import warpdatasets as wd

# Dataset configuration
DATASET_ID = "{dataset_id}"
DEFAULT_TABLE = "{default_table}"
DEFAULT_MODE = "{default_mode}"
DEFAULT_PREFETCH = "{default_prefetch}"


def get_dataset(
    version: str | None = None,
    mode: str = DEFAULT_MODE,
    prefetch: str = DEFAULT_PREFETCH,
    cache_dir: str | None = None,
) -> wd.Dataset:
    """Load the dataset.

    Args:
        version: Version hash or None for latest
        mode: Access mode - "auto", "remote", "hybrid", or "local"
        prefetch: Prefetch mode - "off", "auto", or "aggressive"
        cache_dir: Override cache directory

    Returns:
        Dataset handle
    """
    return wd.dataset(
        DATASET_ID,
        version=version,
        mode=mode,
        prefetch=prefetch,
        cache_dir=cache_dir,
    )


def stream_batches(
    batch_size: int = 50000,
    shard: str = "auto",
    columns: list[str] | None = None,
    limit: int | None = None,
    wrap_refs: bool = {wrap_refs},
    table: str = DEFAULT_TABLE,
) -> Iterator[dict[str, Any]]:
    """Stream batches for training.

    Args:
        batch_size: Number of rows per batch
        shard: Sharding config - "auto" or "rank,world"
        columns: Columns to include (None for all)
        limit: Maximum rows to stream
        wrap_refs: Wrap ref columns as Ref objects
        table: Table name

    Yields:
        Dictionaries with column values
    """
    ds = get_dataset()
    t = ds.table(table)
    yield from t.batch_dicts(
        batch_size=batch_size,
        shard=shard,
        columns=columns,
        limit=limit,
        wrap_refs=wrap_refs,
    )

{helpers}

if __name__ == "__main__":
    import sys

    print(f"Loading dataset: {{DATASET_ID}}")
    print()

    try:
        ds = get_dataset()
        info = ds.info()

        print(f"Dataset: {{info['id']}}")
        print(f"Version: {{info['version']}}")
        print()

        print("Tables:")
        for name, tinfo in info["tables"].items():
            rows = f"{{tinfo['row_count']:,}}" if tinfo["row_count"] else "?"
            print(f"  {{name}}: {{tinfo['shards']}} shard(s), {{rows}} rows")
        print()

        if info["artifacts"]:
            print(f"Artifacts: {{', '.join(info['artifacts'])}}")
            print()

        if info["bindings"]:
            print(f"Bindings: {{info['bindings']}}")
            print()

        # Show schema
        table = ds.table(DEFAULT_TABLE)
        print(f"Schema for {{DEFAULT_TABLE}}:")
        for col, dtype in table.schema().items():
            print(f"  {{col}}: {{dtype}}")
        print()

        # Stream a small batch to verify
        print("Streaming test batch...")
        it = stream_batches(batch_size=128, limit=256, wrap_refs={wrap_refs})
        batch = next(iter(it))
        print(f"Batch columns: {{list(batch.keys())}}")
        print(f"Batch size: {{len(next(iter(batch.values())))}}")
        print()

        print("SUCCESS: Dataset is accessible and streamable.")

    except Exception as e:
        print(f"ERROR: {{e}}", file=sys.stderr)
        print()
        print("Run 'warpdata doctor {{DATASET_ID}}' to diagnose issues.")
        sys.exit(1)
'''


# Helper function templates
IMAGE_DECODE_HELPER = '''
def decode_images(batch: dict, key: str = "{column}") -> list:
    """Decode image refs to PIL Images.

    Args:
        batch: Batch dictionary
        key: Column key containing image refs

    Returns:
        List of PIL Images (or None for missing refs)
    """
    refs = batch.get(key, [])
    return [ref.as_pil() if ref is not None else None for ref in refs]
'''


AUDIO_DECODE_HELPER = '''
def decode_audio(batch: dict, key: str = "{column}") -> list:
    """Decode audio refs to numpy arrays.

    Args:
        batch: Batch dictionary
        key: Column key containing audio refs

    Returns:
        List of (waveform, sample_rate) tuples
    """
    refs = batch.get(key, [])
    results = []
    for ref in refs:
        if ref is not None:
            import soundfile as sf
            import io
            data, sr = sf.read(io.BytesIO(ref.read_bytes()))
            results.append((data, sr))
        else:
            results.append(None)
    return results
'''


FILE_DECODE_HELPER = '''
def read_files(batch: dict, key: str = "{column}") -> list[bytes]:
    """Read file refs as bytes.

    Args:
        batch: Batch dictionary
        key: Column key containing file refs

    Returns:
        List of bytes (or None for missing refs)
    """
    refs = batch.get(key, [])
    return [ref.read_bytes() if ref is not None else None for ref in refs]
'''
