"""
Utility functions for FHEnom AI operations.
"""

import hashlib
import json
import os
import secrets
from pathlib import Path
from typing import Dict, Any, Optional, List
import logging

logger = logging.getLogger(__name__)


def calculate_file_hash(file_path: str, algorithm: str = "sha256") -> str:
    """
    Calculate hash of a file.
    
    Args:
        file_path: Path to file
        algorithm: Hash algorithm (md5, sha1, sha256, etc.)
    
    Returns:
        Hexadecimal hash string
    
    Example:
        ```python
        from fhenomai.utils import calculate_file_hash
        
        hash_val = calculate_file_hash("model.safetensors")
        print(f"SHA256: {hash_val}")
        ```
    """
    hash_obj = hashlib.new(algorithm)
    
    with open(file_path, 'rb') as f:
        while chunk := f.read(8192):
            hash_obj.update(chunk)
    
    return hash_obj.hexdigest()


def calculate_directory_hash(
    directory: str,
    algorithm: str = "sha256",
    exclude_patterns: Optional[List[str]] = None
) -> str:
    """
    Calculate combined hash of all files in a directory.
    
    Args:
        directory: Directory path
        algorithm: Hash algorithm
        exclude_patterns: Patterns to exclude
    
    Returns:
        Hexadecimal hash string
    """
    if exclude_patterns is None:
        exclude_patterns = ['.git', '__pycache__', '.DS_Store']
    
    hash_obj = hashlib.new(algorithm)
    
    # Sort files for consistent ordering
    files = sorted(Path(directory).rglob('*'))
    
    for file_path in files:
        if not file_path.is_file():
            continue
        
        # Check exclusions
        if any(pattern in str(file_path) for pattern in exclude_patterns):
            continue
        
        # Add relative path to hash
        rel_path = file_path.relative_to(directory)
        hash_obj.update(str(rel_path).encode())
        
        # Add file content to hash
        with open(file_path, 'rb') as f:
            while chunk := f.read(8192):
                hash_obj.update(chunk)
    
    return hash_obj.hexdigest()


def format_bytes(num_bytes: int) -> str:
    """
    Format bytes as human-readable string.
    
    Args:
        num_bytes: Number of bytes
    
    Returns:
        Formatted string (e.g., "1.5 GB")
    
    Example:
        ```python
        from fhenomai.utils import format_bytes
        
        print(format_bytes(1536000000))  # "1.43 GB"
        ```
    """
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if num_bytes < 1024.0:
            return f"{num_bytes:.2f} {unit}"
        num_bytes /= 1024.0
    return f"{num_bytes:.2f} PB"


def get_directory_size(directory: str) -> int:
    """
    Calculate total size of a directory in bytes.
    
    Args:
        directory: Directory path
    
    Returns:
        Total size in bytes
    """
    total = 0
    for file_path in Path(directory).rglob('*'):
        if file_path.is_file():
            total += file_path.stat().st_size
    return total


def format_duration(seconds: float) -> str:
    """
    Format duration in seconds as human-readable string.
    
    Args:
        seconds: Duration in seconds
    
    Returns:
        Formatted string (e.g., "2h 15m 30s")
    
    Example:
        ```python
        from fhenomai.utils import format_duration
        
        print(format_duration(8130))  # "2h 15m 30s"
        ```
    """
    if seconds < 60:
        return f"{seconds:.1f}s"
    
    minutes = int(seconds // 60)
    seconds = seconds % 60
    
    if minutes < 60:
        return f"{minutes}m {seconds:.0f}s"
    
    hours = minutes // 60
    minutes = minutes % 60
    
    return f"{hours}h {minutes}m {seconds:.0f}s"


def validate_model_directory(directory: str) -> bool:
    """
    Validate that a directory contains a valid model.
    
    Checks for required files like config.json and model files.
    
    Args:
        directory: Model directory path
    
    Returns:
        True if valid, False otherwise
    """
    dir_path = Path(directory)
    
    if not dir_path.exists():
        logger.error(f"Directory does not exist: {directory}")
        return False
    
    # Check for config.json
    config_file = dir_path / "config.json"
    if not config_file.exists():
        logger.error(f"Missing config.json in {directory}")
        return False
    
    # Check for model files (.safetensors or .bin)
    model_files = list(dir_path.glob("*.safetensors")) + list(dir_path.glob("*.bin"))
    if not model_files:
        logger.error(f"No model files found in {directory}")
        return False
    
    logger.info(f"Valid model directory: {directory}")
    return True


def load_json_file(file_path: str) -> Dict[str, Any]:
    """
    Load and parse a JSON file.
    
    Args:
        file_path: Path to JSON file
    
    Returns:
        Parsed JSON data
    """
    with open(file_path, 'r') as f:
        return json.load(f)


def save_json_file(data: Dict[str, Any], file_path: str, indent: int = 2):
    """
    Save data to a JSON file.
    
    Args:
        data: Data to save
        file_path: Output file path
        indent: JSON indentation
    """
    with open(file_path, 'w') as f:
        json.dump(data, f, indent=indent)


def parse_model_id(model_path: str) -> str:
    """
    Extract model ID from a model path.
    
    Args:
        model_path: Model path (local or remote)
    
    Returns:
        Model ID
    
    Example:
        ```python
        from fhenomai.utils import parse_model_id
        
        model_id = parse_model_id("/models/llama-3.2-3b")
        print(model_id)  # "llama-3.2-3b"
        ```
    """
    return Path(model_path).name


def build_model_config(
    model_type: str,
    vocab_size: Optional[int] = None,
    hidden_size: Optional[int] = None,
    **kwargs
) -> Dict[str, Any]:
    """
    Build a model configuration dictionary.
    
    Args:
        model_type: Model architecture type
        vocab_size: Vocabulary size
        hidden_size: Hidden layer size
        **kwargs: Additional config parameters
    
    Returns:
        Configuration dictionary
    """
    config = {
        "model_type": model_type,
        **kwargs
    }
    
    if vocab_size is not None:
        config["vocab_size"] = vocab_size
    
    if hidden_size is not None:
        config["hidden_size"] = hidden_size
    
    return config


def retry_with_backoff(
    func,
    max_retries: int = 3,
    initial_delay: float = 1.0,
    backoff_factor: float = 2.0,
    exceptions: tuple = (Exception,)
):
    """
    Retry a function with exponential backoff.
    
    Args:
        func: Function to retry
        max_retries: Maximum number of retries
        initial_delay: Initial delay in seconds
        backoff_factor: Multiplier for delay after each retry
        exceptions: Tuple of exceptions to catch
    
    Returns:
        Function result
    
    Example:
        ```python
        from fhenomai.utils import retry_with_backoff
        
        def flaky_function():
            # Might fail sometimes
            pass
        
        result = retry_with_backoff(flaky_function, max_retries=5)
        ```
    """
    import time
    
    delay = initial_delay
    last_exception = None
    
    for attempt in range(max_retries):
        try:
            return func()
        except exceptions as e:
            last_exception = e
            if attempt < max_retries - 1:
                logger.warning(f"Attempt {attempt + 1} failed: {e}. Retrying in {delay}s...")
                time.sleep(delay)
                delay *= backoff_factor
            else:
                logger.error(f"All {max_retries} attempts failed")
    
    raise last_exception


def merge_dicts(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
    """
    Recursively merge two dictionaries.
    
    Args:
        base: Base dictionary
        override: Dictionary with overrides
    
    Returns:
        Merged dictionary
    """
    result = base.copy()
    
    for key, value in override.items():
        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
            result[key] = merge_dicts(result[key], value)
        else:
            result[key] = value
    
    return result


def generate_nonce(size: int = 64) -> bytes:
    """
    Generate a cryptographically secure random nonce for TEE attestation.
    
    Args:
        size: Size of nonce in bytes (default: 64)
    
    Returns:
        Random bytes
    
    Example:
        ```python
        from fhenomai.utils import generate_nonce, format_nonce_hex
        
        # Generate 64-byte nonce
        nonce = generate_nonce()
        nonce_hex = format_nonce_hex(nonce)
        print(f"Nonce: {nonce_hex}")
        ```
    """
    return os.urandom(size)


def format_nonce_hex(nonce: bytes) -> str:
    """
    Format nonce bytes as hexadecimal string.
    
    Args:
        nonce: Nonce bytes
    
    Returns:
        Hexadecimal string representation
    
    Example:
        ```python
        from fhenomai.utils import generate_nonce, format_nonce_hex
        
        nonce = generate_nonce(64)
        nonce_hex = format_nonce_hex(nonce)
        # Returns 128-character hex string
        ```
    """
    return nonce.hex()


def parse_nonce_hex(nonce_hex: str) -> bytes:
    """
    Parse hexadecimal nonce string to bytes.
    
    Args:
        nonce_hex: Hexadecimal nonce string
    
    Returns:
        Nonce bytes
    
    Raises:
        ValueError: If hex string is invalid
    
    Example:
        ```python
        from fhenomai.utils import parse_nonce_hex
        
        nonce_hex = "a1b2c3..."  # 128 hex chars
        nonce = parse_nonce_hex(nonce_hex)
        ```
    """
    return bytes.fromhex(nonce_hex)
