"""Concurrent batch validation of STAC files using multiprocessing."""

import concurrent.futures
import json
import logging
import multiprocessing
from itertools import islice
from typing import Any, Dict, Iterable, List, Optional, Tuple

from .utilities import _build_cached_validator
from .validate import StacValidate

logger = logging.getLogger(__name__)


def _warm_schema_cache(stac_version: str = "1.0.0") -> None:
    """
    Pre-warm the schema cache before forking worker processes.

    This function loads and compiles core STAC schemas in the main process.
    When ProcessPoolExecutor forks child processes, they inherit the parent's
    memory via Copy-on-Write (CoW), so all workers get a fully-warmed cache
    without making any network requests.

    Args:
        stac_version: STAC version to use for schema selection (e.g., "1.0.0" or "1.1.0")

    Core schemas warmed:
    - STAC Item
    - STAC Collection
    - STAC Catalog
    """
    try:
        import os

        # Use local schema files bundled with the package
        # Support both v1.0.0 and v1.1.0
        schema_dir = os.path.join(
            os.path.dirname(__file__), "local_schemas", stac_version
        )

        core_schemas = [
            os.path.join(schema_dir, "item.json"),
            os.path.join(schema_dir, "collection.json"),
            os.path.join(schema_dir, "catalog.json"),
        ]

        from .utilities import fetch_and_parse_schema

        # Load and cache each schema
        for schema_path in core_schemas:
            try:
                schema = fetch_and_parse_schema(schema_path)
                # Convert to JSON string and cache via _build_cached_validator
                schema_json = json.dumps(schema, sort_keys=True, separators=(",", ":"))
                _build_cached_validator(schema_json)
            except Exception:
                # If a schema fails to load, continue with others
                pass

    except Exception:
        # If cache warming fails entirely, continue without it
        pass


def get_optimal_worker_count(max_workers: Optional[int] = None) -> int:
    """
    Get the optimal number of worker processes.

    Detects available CPU cores, accounting for containerized environments.
    Falls back to os.sched_getaffinity() on Linux for Docker/container support.

    Args:
        max_workers: Maximum number of workers to use. Options:
            - None or 0: Use all available cores (default)
            - Positive int: Use that many cores (capped at available)
            - Negative int: Use all cores minus that many (e.g., -1 reserves 1 core for OS)

    Returns:
        Optimal number of worker processes

    Note:
        In containerized environments (Docker, ECS), this uses os.sched_getaffinity()
        on Linux to detect the actual cores allocated to the container, rather than
        the host machine's total cores.
    """
    import os as os_module

    # Try to get container-aware core count on Linux
    try:
        if hasattr(os_module, "sched_getaffinity"):
            # Linux: respects Docker/container CPU limits
            total_cores = len(os_module.sched_getaffinity(0))
        else:
            # Fallback for non-Linux systems
            total_cores = multiprocessing.cpu_count()
    except Exception:
        # Fallback if anything goes wrong
        total_cores = multiprocessing.cpu_count()

    if max_workers is None or max_workers == 0:
        # Use all available cores
        return total_cores
    elif max_workers < 0:
        # Use all cores minus the specified amount (useful for reserving cores for OS)
        return max(1, total_cores + max_workers)
    else:
        # Use the specified number, but cap at available cores
        return min(max_workers, total_cores)


def _validate_single_file(file_path: str) -> Tuple[str, bool, List[str]]:
    """
    Worker function that runs on an individual CPU core.
    Validates a single STAC file and returns results.

    Each worker process has its own isolated schema cache for optimal performance.

    Args:
        file_path: Path to the STAC JSON file to validate.

    Returns:
        Tuple of (file_path, is_valid, list_of_errors)
    """
    errors = []

    try:
        # Use StacValidate for comprehensive validation (includes version check, core, and extensions)
        # Each worker process has its own isolated schema cache (no need to reconfigure)
        validator = StacValidate(file_path)
        validator.run()

        # Collect errors from validation
        # validator.message is a list of result objects
        if validator.message:
            try:
                messages = validator.message
                if isinstance(messages, list) and len(messages) > 0:
                    msg_obj = messages[0]
                    if not msg_obj.get("valid_stac", False):
                        # Try multiple error field names
                        if "errors" in msg_obj:
                            errors.extend(msg_obj["errors"])
                        elif "error_message" in msg_obj:
                            errors.append(msg_obj["error_message"])
                        else:
                            errors.append(
                                f"{msg_obj.get('error_type', 'ValidationError')}"
                            )
            except (KeyError, IndexError, TypeError):
                # If we can't parse, just use the raw message
                if validator.message:
                    errors.append(str(validator.message))

    except Exception as e:
        # Catch any exception including StacValidate instantiation failures
        # This prevents UnboundLocalError if validator creation fails
        return file_path, False, [f"Critical error processing file: {str(e)}"]

    is_valid = len(errors) == 0
    return file_path, is_valid, errors


def _validate_dict(
    item_dict: Dict[str, Any], source_path: str
) -> Tuple[str, bool, Dict[str, Any], str]:
    """
    Worker function that validates a STAC item dictionary directly (no temp files).

    This function receives pickled dictionaries from the main process, avoiding
    expensive disk I/O. Each worker process has its own isolated schema cache.

    Args:
        item_dict: Dictionary representation of a STAC item to validate.
        source_path: Display path for result reporting (e.g., "file.json[0]").

    Returns:
        Tuple of (source_path, is_valid, message_dict, item_id)
        where message_dict is the complete validation result from StacValidate
    """
    item_id = item_dict.get("id", "unknown")
    message = {}

    try:
        # Validate the dictionary directly without writing to disk
        # StacValidate can accept a dict via validate_dict() method
        validator = StacValidate()
        validator.stac_content = item_dict

        # Set STAC version from item
        validator.version = item_dict.get("stac_version", "1.0.0")

        # Run validation on the dictionary
        validator.validate_dict(item_dict)

        # Get the validation result message
        if validator.message:
            messages = validator.message
            if isinstance(messages, list) and len(messages) > 0:
                message = messages[0]
            else:
                message = messages if isinstance(messages, dict) else {}

            # Set path to the display path (not a real file)
            message["path"] = source_path

    except Exception as e:
        # Catch any exception during validation
        message = {
            "path": source_path,
            "valid_stac": False,
            "error_type": "Exception",
            "error_message": f"Critical error processing item: {str(e)}",
            "failed_schema": "",
        }

    is_valid = message.get("valid_stac", False)
    return source_path, is_valid, message, item_id


def validate_concurrently(
    file_paths: List[str],
    max_workers: Optional[int] = None,
    show_progress: bool = True,
    feature_collection: bool = False,
    batch_size: int = 2000,
) -> List[Dict[str, Any]]:
    """
    Validates a list of STAC files concurrently using available CPU cores.

    Uses ProcessPoolExecutor to bypass Python's GIL by creating separate Python
    processes for each worker. Pre-warms the schema cache in the main process
    before forking, so all workers inherit a fully-populated cache via Copy-on-Write.

    Args:
        file_paths: List of paths to STAC JSON files or FeatureCollections.
        max_workers: Number of CPU cores to use. Options:
            - None or 0: Use all available cores (default)
            - Positive int: Use that many cores (capped at available)
            - Negative int: Use all cores minus that many (e.g., -1 = all cores - 1)
        show_progress: Whether to display a progress bar (requires tqdm).
        feature_collection: If True, treat files as FeatureCollections and validate each feature.
        batch_size: Number of items to process at a time to bound memory usage (default: 2000).

    Returns:
        List of result dictionaries with keys: path, valid_stac, errors
    """
    # If feature_collection mode, extract features and delegate to validate_dicts
    # for memory-safe chunked processing
    if feature_collection:
        error_results: List[Dict[str, Any]] = []

        def item_iter() -> Iterable[Dict[str, Any]]:
            """
            Lazily iterate over all items (features or standalone objects)
            from the provided file_paths, annotating each with source
            metadata. Any file-level read/parse errors are recorded in
            error_results and skipped.
            """
            for file_path in file_paths:
                try:
                    with open(file_path, "r") as f:
                        data = json.load(f)

                    # Check if it's a FeatureCollection
                    if (
                        isinstance(data, dict)
                        and data.get("type") == "FeatureCollection"
                    ):
                        features = data.get("features", [])
                        for idx, feature in enumerate(features):
                            # Attach source info for result display
                            feature["__source_file__"] = file_path
                            feature["__feature_index__"] = idx
                            yield feature
                    else:
                        # Regular file, treat as single item
                        data["__source_file__"] = file_path
                        data["__feature_index__"] = None
                        yield data

                except Exception as e:
                    # Accumulate error for this file and continue processing others
                    error_results.append(
                        {
                            "path": file_path,
                            "valid_stac": False,
                            "errors": [f"Failed to read file: {str(e)}"],
                        }
                    )

        # Delegate to validate_dicts for chunked, memory-safe processing
        # validate_dicts will detect STAC version from first item and warm cache
        validation_results = validate_dicts(
            item_iter(),
            max_workers=max_workers,
            show_progress=show_progress,
            chunk_size=batch_size,
        )

        # Combine error results with validation results
        return error_results + validation_results

    # Standard file path validation (no FeatureCollection expansion)
    results: List[Dict[str, Any]] = []

    try:
        from tqdm import tqdm

        has_tqdm = True
    except ImportError:
        has_tqdm = False
        show_progress = False

    optimal_workers = get_optimal_worker_count(max_workers)

    if not file_paths:
        return results

    try:
        with concurrent.futures.ProcessPoolExecutor(
            max_workers=optimal_workers
        ) as executor:

            # Submit all tasks to the pool
            future_to_file = {
                executor.submit(_validate_single_file, path): path
                for path in file_paths
            }

            # Wrap the as_completed iterator with tqdm for a nice progress bar
            iterator = concurrent.futures.as_completed(future_to_file)
            if show_progress and has_tqdm:
                iterator = tqdm(  # type: ignore
                    iterator,
                    total=len(file_paths),
                    desc="Validating STAC Items",
                )

            # Collect results as they finish
            for future in iterator:
                file_path, is_valid, errors = future.result()

                result = {
                    "path": file_path,
                    "valid_stac": is_valid,
                }

                if errors:
                    result["errors"] = errors

                results.append(result)

    finally:
        pass

    return results


def _chunked_iterable(iterable: Iterable, size: int):
    """Yield successive n-sized chunks from an iterable."""
    it = iter(iterable)
    while True:
        chunk = list(islice(it, size))
        if not chunk:
            break
        yield chunk


def validate_dicts(
    items: Iterable[Dict[str, Any]],
    max_workers: Optional[int] = None,
    show_progress: bool = False,
    chunk_size: int = 2000,
) -> List[Dict[str, Any]]:
    """
    Validate an iterable of STAC item dictionaries concurrently.

    Passes dictionaries directly to worker processes via pickle (memory-based),
    avoiding expensive disk I/O. Pre-warms the schema cache in the main process
    before forking, so all workers inherit a fully-populated cache via Copy-on-Write.
    Processes items in bounded chunks to prevent out-of-memory errors.

    Args:
        items: Iterable of STAC item dictionaries to validate
        max_workers: Maximum number of worker processes. Options:
            - None or 0: Use all available cores (default)
            - Positive int: Use that many cores (capped at available)
            - Negative int: Use all cores minus that many
        show_progress: Whether to show progress bar (default: False). Set to True for CLI usage.
        chunk_size: Number of items to process at a time to bound memory usage.

    Returns:
        List of validation results with path, valid_stac, and errors (if any)
    """
    all_results = []

    try:
        from tqdm import tqdm

        has_tqdm = True
    except ImportError:
        has_tqdm = False
        show_progress = False

    optimal_workers = get_optimal_worker_count(max_workers)

    # Detect STAC version from first item and warm cache accordingly
    stac_version = "1.0.0"  # default
    items_list = list(items)
    if items_list:
        first_item = items_list[0]
        stac_version = first_item.get("stac_version", "1.0.0")

    # Pre-warm the schema cache in the main process before forking workers
    # This leverages Copy-on-Write so all workers inherit the warmed cache
    _warm_schema_cache(stac_version=stac_version)

    # Create a single executor for the entire run to reuse workers across chunks
    # This preserves the schema cache across chunks and eliminates fork/spawn overhead
    try:
        with concurrent.futures.ProcessPoolExecutor(
            max_workers=optimal_workers
        ) as executor:
            total_items = len(items_list)
            progress_bar = None

            # Initialize progress bar if available
            if show_progress and has_tqdm:
                progress_bar = tqdm(  # type: ignore
                    total=total_items,
                    desc="Validating Items",
                    unit="item",
                )

            # Process items in chunks: submit a chunk, collect results, then submit next
            # This bounds in-flight futures and memory usage while reusing workers
            for chunk in _chunked_iterable(items_list, chunk_size):
                if not chunk:
                    continue

                # Submit all items in this chunk
                future_to_item = {}
                for idx, item in enumerate(chunk):
                    # Extract source metadata if present
                    source_file = item.get("__source_file__")
                    feature_index = item.get("__feature_index__")

                    # Create display path for results
                    if source_file is not None:
                        display_path = (
                            f"{source_file}[{feature_index}]"
                            if feature_index is not None
                            else source_file
                        )
                    else:
                        display_path = f"item-{idx}"

                    # Create payload without internal metadata keys
                    payload = {
                        k: v
                        for k, v in item.items()
                        if k not in ("__source_file__", "__feature_index__")
                    }

                    # Submit to worker - dictionary is pickled automatically
                    future = executor.submit(_validate_dict, payload, display_path)
                    future_to_item[future] = (display_path, idx)

                # Collect results for this chunk as they finish
                for future in concurrent.futures.as_completed(future_to_item):
                    display_path, idx = future_to_item[future]
                    source_path, is_valid, message, item_id = future.result()

                    # Use the complete message from validation
                    result = message.copy() if message else {}

                    # Ensure required fields are present
                    if "path" not in result:
                        result["path"] = source_path
                    if "valid_stac" not in result:
                        result["valid_stac"] = is_valid

                    # Add item_id for reference
                    result["item_id"] = item_id

                    all_results.append(result)

                    # Update progress bar
                    if progress_bar:
                        progress_bar.update(1)

            # Close progress bar
            if progress_bar:
                progress_bar.close()

    except Exception as e:
        # If processing fails, record error for all items
        for idx, item in enumerate(items_list):
            source_file = item.get("__source_file__")
            feature_index = item.get("__feature_index__")

            if source_file is not None:
                display_path = (
                    f"{source_file}[{feature_index}]"
                    if feature_index is not None
                    else source_file
                )
            else:
                display_path = f"item-{idx}"

            all_results.append(
                {
                    "path": display_path,
                    "valid_stac": False,
                    "errors": [f"Chunk processing error: {str(e)}"],
                }
            )

    return all_results
