# Copyright (c) Alibaba, Inc. and its affiliates.

import contextlib
import errno
import hashlib
import os
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Generator, List, Optional, Union

from filelock import BaseFileLock, FileLock, SoftFileLock, Timeout

from modelscope.hub.constants import (DEFAULT_MODELSCOPE_DOMAIN,
                                      DEFAULT_MODELSCOPE_GROUP,
                                      DEFAULT_MODELSCOPE_INTL_DOMAIN,
                                      MODEL_ID_SEPARATOR, MODELSCOPE_DOMAIN,
                                      MODELSCOPE_SDK_DEBUG,
                                      MODELSCOPE_URL_SCHEME)
from modelscope.utils.logger import get_logger

logger = get_logger()


def model_id_to_group_owner_name(model_id):
    if MODEL_ID_SEPARATOR in model_id:
        group_or_owner = model_id.split(MODEL_ID_SEPARATOR)[0]
        name = model_id.split(MODEL_ID_SEPARATOR)[1]
    else:
        group_or_owner = DEFAULT_MODELSCOPE_GROUP
        name = model_id
    return group_or_owner, name


def is_env_true(var_name):
    value = os.environ.get(var_name, '').strip().lower()
    return value == 'true'


def get_domain(cn_site=True):
    if MODELSCOPE_DOMAIN in os.environ and os.getenv(MODELSCOPE_DOMAIN):
        return os.getenv(MODELSCOPE_DOMAIN)
    if cn_site:
        return DEFAULT_MODELSCOPE_DOMAIN
    else:
        return DEFAULT_MODELSCOPE_INTL_DOMAIN


def convert_patterns(raw_input: Union[str, List[str]]):
    output = None
    if isinstance(raw_input, str):
        output = list()
        if ',' in raw_input:
            output = [s.strip() for s in raw_input.split(',')]
        else:
            output.append(raw_input.strip())
    elif isinstance(raw_input, list):
        output = list()
        for s in raw_input:
            if isinstance(s, str):
                if ',' in s:
                    output.extend([ss.strip() for ss in s.split(',')])
                else:
                    output.append(s.strip())
    return output


def extract_root_from_patterns(
    allow_file_pattern: Optional[List[str]] = None,
    allow_patterns: Optional[List[str]] = None,
) -> Optional[str]:
    """Extract common directory prefix from include patterns for server-side filtering.

    Only processes allow/include patterns (ignore/exclude is irrelevant for prefix).
    Returns None if no meaningful prefix can be extracted.

    Algorithm:
    1. Merge allow_file_pattern and allow_patterns into one list
    2. For each pattern, find position of first wildcard char (* ? [)
    3. Extract text before that position, then take directory part (up to last '/')
    4. Find longest common directory prefix (path-segment-aware)
    5. Return None if no valid common prefix

    Examples:
        ['TacExo/*'] -> 'TacExo'
        ['data/train/*.parquet'] -> 'data/train'
        ['data/train/*', 'data/valid/*'] -> 'data'
        ['*.safetensors'] -> None
        ['TacExo/*', 'OtherDir/*'] -> None (no common prefix)
        ['data/*/train.csv'] -> 'data'
    """
    # Merge both pattern lists
    patterns = []
    if allow_file_pattern:
        patterns.extend(allow_file_pattern)
    if allow_patterns:
        patterns.extend(allow_patterns)

    if not patterns:
        return None

    extracted_dirs = []
    for pattern in patterns:
        # Find position of first wildcard character
        first_wildcard = len(pattern)
        for wc in ('*', '?', '['):
            pos = pattern.find(wc)
            if pos != -1:
                first_wildcard = min(first_wildcard, pos)

        # Get text before wildcard
        prefix = pattern[:first_wildcard]

        # Extract directory part (up to last '/')
        last_sep = prefix.rfind('/')
        if last_sep > 0:
            dir_part = prefix[:last_sep]
            # Validate: no wildcards should remain in dir_part
            if not any(c in dir_part for c in ('*', '?', '[')):
                extracted_dirs.append(dir_part)
        # If no '/' found or only at position 0, this pattern has no directory prefix

    if not extracted_dirs:
        return None

    # Find longest common directory prefix (path-segment-aware)
    if len(extracted_dirs) == 1:
        return extracted_dirs[0]

    # Split all paths into segments and find common prefix segments
    split_paths = [d.split('/') for d in extracted_dirs]
    common_segments = []
    for segments in zip(*split_paths):
        if len(set(segments)) == 1:
            common_segments.append(segments[0])
        else:
            break

    if not common_segments:
        return None

    return '/'.join(common_segments)


# during model download, the '.' would be converted to '___' to produce
# actual physical (masked) directory for storage
def get_model_masked_directory(directory, model_id):
    if sys.platform.startswith('win'):
        parts = directory.rsplit('\\', 2)
    else:
        parts = directory.rsplit('/', 2)
    # this is the actual directory the model files are located.
    masked_directory = os.path.join(parts[0], model_id.replace('.', '___'))
    return masked_directory


def convert_readable_size(size_bytes: int) -> str:
    import math
    if size_bytes == 0:
        return '0B'
    size_name = ('B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB')
    i = int(math.floor(math.log(size_bytes, 1024)))
    p = math.pow(1024, i)
    s = round(size_bytes / p, 2)
    return f'{s} {size_name[i]}'


def get_folder_size(folder_path: str) -> int:
    total_size = 0
    for path in Path(folder_path).rglob('*'):
        if path.is_file():
            total_size += path.stat().st_size
    return total_size


# return a readable string that describe size of for a given folder (MB, GB etc.)
def get_readable_folder_size(folder_path: str) -> str:
    return convert_readable_size(get_folder_size(folder_path=folder_path))


def get_cache_dir(model_id: Optional[str] = None):
    """cache dir precedence:
        function parameter > environment > ~/.cache/modelscope/hub
    Args:
        model_id (str, optional): The model id.
    Returns:
        str: the model_id dir if model_id not None, otherwise cache root dir.
    """
    default_cache_dir = Path.home().joinpath('.cache', 'modelscope')
    base_path = os.getenv('MODELSCOPE_CACHE',
                          os.path.join(default_cache_dir, 'hub'))
    return base_path if model_id is None else os.path.join(
        base_path, model_id + '/')


def _modelscope_hub_cache_root() -> Path:
    """Cache root used by ``modelscope_hub`` downloads (not SDK ``.../hub``)."""
    env = os.environ.get('MODELSCOPE_CACHE')
    if env:
        return Path(env).expanduser()
    return Path.home() / '.cache' / 'modelscope'


def find_reusable_legacy_repo_dir(
    repo_id: str,
    repo_type: str = 'model',
    cache_dir: Optional[Union[str, Path]] = None,
) -> Optional[str]:
    """Find old on-disk cache layouts that ``modelscope_hub`` download misses.

    ``modelscope_hub`` reuses ``{cache}/{type}s/{owner}/{safe_name}/`` (dots in
    ``name`` replaced by ``___``) and writes to
    ``{cache}/{type}s/{owner}--{name}/snapshots/{rev}/``. Older SDKs also
    stored repos at:

    - ``{cache}/{owner}/{name}/`` (flat, when ``MODELSCOPE_CACHE`` was set)
    - ``{cache}/hub/{owner}/{name}/`` (pre-``models/`` restructuring)
    - ``{cache}/{type}s/{owner}/{name}/`` (unsafed name; hub only checks
      ``safe_name``)

    Returns a non-empty legacy path only when the layouts hub already handles
    are absent, so callers can pass it as ``local_dir`` and avoid re-download.
    """
    if not repo_id or '/' not in repo_id:
        return None

    base = Path(cache_dir).expanduser() if cache_dir is not None else \
        _modelscope_hub_cache_root()
    segment = f'{repo_type}s' if not repo_type.endswith('s') else repo_type
    owner, name = repo_id.split('/', 1)
    safe_name = name.replace('.', '___')
    safe_id = repo_id.replace('/', '--')

    # Layouts already handled by modelscope_hub — do not override.
    hub_known = [
        base / segment / safe_id,
        base / segment / owner / safe_name,
    ]
    for path in hub_known:
        if _non_empty_dir(path):
            return None

    # Layouts hub download does not probe today.
    legacy_candidates = [
        base / owner / name,
        base / owner / safe_name,
        base / 'hub' / owner / name,
        base / 'hub' / owner / safe_name,
        base / segment / owner / name,
    ]
    for path in legacy_candidates:
        if _non_empty_dir(path):
            logger.info('Found legacy cache at %s for %s, reusing.', path,
                        repo_id)
            return str(path)
    return None


def _non_empty_dir(path: Path) -> bool:
    if not path.is_dir():
        return False
    try:
        return any(path.iterdir())
    except OSError:
        return False


def get_release_datetime():
    if MODELSCOPE_SDK_DEBUG in os.environ:
        rt = int(round(datetime.now().timestamp()))
    else:
        from modelscope import version
        rt = int(
            round(
                datetime.strptime(version.__release_datetime__,
                                  '%Y-%m-%d %H:%M:%S').timestamp()))
    return rt


def get_endpoint(cn_site=True):
    return MODELSCOPE_URL_SCHEME + get_domain(cn_site)


def resolve_endpoint(cli_endpoint: Optional[str] = None,
                     cn_site: bool = True) -> str:
    """Resolve the ModelScope API endpoint with automatic scheme completion.

    Priority (highest to lowest):
        1. ``cli_endpoint`` (explicit CLI --endpoint argument)
        2. Environment variable ``MODELSCOPE_DOMAIN``
        3. Built-in default (https://www.modelscope.cn)

    Scheme auto-completion:
        If the resolved value does not start with ``http://`` or ``https://``,
        ``https://`` is prepended automatically so that callers may pass bare
        domain names such as ``modelscope.ai``.

    Args:
        cli_endpoint: Value from the CLI ``--endpoint`` flag.  When *None*,
            the function falls back to :func:`get_endpoint`.
        cn_site: Forwarded to :func:`get_endpoint` when *cli_endpoint* is
            *None*.  ``True`` selects the Chinese site, ``False`` the
            international site.

    Returns:
        A fully-qualified endpoint URL, e.g. ``https://www.modelscope.cn``.
    """
    if cli_endpoint is None:
        return get_endpoint(cn_site=cn_site)
    endpoint = cli_endpoint.strip().rstrip('/')
    if not endpoint:
        return get_endpoint(cn_site=cn_site)
    if not endpoint.startswith('http://') and not endpoint.startswith(
            'https://'):
        endpoint = MODELSCOPE_URL_SCHEME + endpoint
    return endpoint


def compute_hash(file_path):
    # 16MB buffer for large file hash computation
    BUFFER_SIZE = 1024 * 1024 * 16
    sha256_hash = hashlib.sha256()
    with open(file_path, 'rb') as f:
        while True:
            data = f.read(BUFFER_SIZE)
            if not data:
                break
            sha256_hash.update(data)
    return sha256_hash.hexdigest()


def file_integrity_validation(file_path, expected_sha256) -> bool:
    """Validate the file hash is expected, if not, delete the file

    Args:
        file_path (str): The file to validate
        expected_sha256 (str): The expected sha256 hash

    Returns:
        bool: True if the file is valid, False otherwise
    """
    file_sha256 = compute_hash(file_path)
    if not file_sha256 == expected_sha256:
        os.remove(file_path)
        msg = 'File %s integrity check failed, expected sha256 signature is %s, actual is %s, the download may be incomplete, please try again.' % (  # noqa E501
            file_path, expected_sha256, file_sha256)
        logger.error(msg)
        return False

    return True


def add_content_to_file(repo,
                        file_name: str,
                        patterns: List[str],
                        commit_message: Optional[str] = None,
                        ignore_push_error=False) -> None:
    if isinstance(patterns, str):
        patterns = [patterns]
    if commit_message is None:
        commit_message = f'Add `{patterns[0]}` patterns to {file_name}'

    # Get current file content
    repo_dir = repo.model_dir
    file_path = os.path.join(repo_dir, file_name)
    if os.path.exists(file_path):
        with open(file_path, 'r', encoding='utf-8') as f:
            current_content = f.read()
    else:
        current_content = ''
    # Add the patterns to file
    content = current_content
    for pattern in patterns:
        if pattern not in content:
            if len(content) > 0 and not content.endswith('\n'):
                content += '\n'
            content += f'{pattern}\n'

    # Write the file if it has changed
    if content != current_content:
        with open(file_path, 'w', encoding='utf-8') as f:
            logger.debug(f'Writing {file_name} file. Content: {content}')
            f.write(content)
    try:
        repo.push(commit_message)
    except Exception as e:
        if ignore_push_error:
            pass
        else:
            raise e


_TIMESINCE_CHUNKS = (
    # Label, divider, max value
    ('second', 1, 60),
    ('minute', 60, 60),
    ('hour', 60 * 60, 24),
    ('day', 60 * 60 * 24, 6),
    ('week', 60 * 60 * 24 * 7, 6),
    ('month', 60 * 60 * 24 * 30, 11),
    ('year', 60 * 60 * 24 * 365, None),
)


def format_timesince(ts: float) -> str:
    """Format timestamp in seconds into a human-readable string, relative to now.
    """
    delta = time.time() - ts
    if delta < 20:
        return 'a few seconds ago'
    for label, divider, max_value in _TIMESINCE_CHUNKS:  # noqa: B007
        value = round(delta / divider)
        if max_value is not None and value <= max_value:
            break
    return f"{value} {label}{'s' if value > 1 else ''} ago"


def tabulate(rows: List[List[Union[str, int]]], headers: List[str]) -> str:
    """
    Inspired by:

    - stackoverflow.com/a/8356620/593036
    - stackoverflow.com/questions/9535954/printing-lists-as-tabular-data
    """
    col_widths = [max(len(str(x)) for x in col) for col in zip(*rows, headers)]
    row_format = ('{{:{}}} ' * len(headers)).format(*col_widths)
    lines = []
    lines.append(row_format.format(*headers))
    lines.append(row_format.format(*['-' * w for w in col_widths]))
    for row in rows:
        lines.append(row_format.format(*row))
    return '\n'.join(lines)


# Part of the code borrowed from the awesome work of huggingface_hub/transformers
def strtobool(val):
    val = val.lower()
    if val in {'y', 'yes', 't', 'true', 'on', '1'}:
        return 1
    if val in {'n', 'no', 'f', 'false', 'off', '0'}:
        return 0
    raise ValueError(f'invalid truth value {val!r}')


@contextlib.contextmanager
def weak_file_lock(lock_file: Union[str, Path],
                   *,
                   timeout: Optional[float] = None
                   ) -> Generator[BaseFileLock, None, None]:
    default_interval = 60
    lock = FileLock(lock_file, timeout=default_interval)
    start_time = time.time()

    while True:
        elapsed_time = time.time() - start_time
        if timeout is not None and elapsed_time >= timeout:
            raise Timeout(str(lock_file))

        try:
            lock.acquire(
                timeout=min(default_interval, timeout - elapsed_time)
                if timeout else default_interval)  # noqa
        except Timeout:
            logger.info(
                f'Still waiting to acquire lock on {lock_file} (elapsed: {time.time() - start_time:.1f} seconds)'
            )
        except NotImplementedError as e:
            if 'use SoftFileLock instead' in str(e):
                logger.warning(
                    'FileSystem does not appear to support flock. Falling back to SoftFileLock for %s',
                    lock_file)
                lock = SoftFileLock(lock_file, timeout=default_interval)
                continue
            raise
        except OSError as e:
            # Handle NFS stale file handle and similar issues
            if e.errno in (errno.ESTALE, errno.ENOENT, errno.EREMOTEIO):
                logger.warning(
                    'Encountered OSError (errno=%d) on %s, Falling back to SoftFileLock.',
                    e.errno, lock_file)
                lock = SoftFileLock(lock_file, timeout=default_interval)
                continue
            raise
        else:
            break

    try:
        yield lock
    finally:
        try:
            lock.release()
        except OSError:
            try:
                Path(lock_file).unlink()
            except OSError:
                pass


def convert_timestamp(time_stamp: Union[int, str, datetime],
                      time_zone: str = 'Asia/Shanghai') -> Optional[datetime]:
    """Convert a UNIX/string timestamp to a timezone-aware datetime object.
    Args:
        time_stamp: UNIX timestamp (int), ISO string, or datetime object
        time_zone: Target timezone for non-UTC timestamps (default: 'Asia/Shanghai')
    Returns:
        Timezone-aware datetime object or None if input is None
    """
    import zoneinfo
    if not time_stamp:
        return None

    # Handle datetime objects first
    if isinstance(time_stamp, datetime):
        return time_stamp

    if isinstance(time_stamp, str):
        try:
            if time_stamp.endswith('Z'):
                # Normalize fractional seconds to 6 digits
                if '.' not in time_stamp:
                    # No fractional seconds (e.g., "2024-11-16T00:27:02Z")
                    time_stamp = time_stamp[:-1] + '.000000Z'
                else:
                    # Has fractional seconds (e.g., "2022-08-19T07:19:38.123456789Z")
                    base, fraction = time_stamp[:-1].split('.')
                    # Truncate or pad to 6 digits
                    fraction = fraction[:6].ljust(6, '0')
                    time_stamp = f'{base}.{fraction}Z'

                dt = datetime.strptime(time_stamp,
                                       '%Y-%m-%dT%H:%M:%S.%fZ').replace(
                                           tzinfo=zoneinfo.ZoneInfo('UTC'))
                if time_zone != 'UTC':
                    dt = dt.astimezone(zoneinfo.ZoneInfo(time_zone))
                return dt
            else:
                # Try parsing common ISO formats
                formats = [
                    '%Y-%m-%dT%H:%M:%S.%f',  # With microseconds
                    '%Y-%m-%dT%H:%M:%S',  # Without microseconds
                    '%Y-%m-%d %H:%M:%S.%f',  # Space separator with microseconds
                    '%Y-%m-%d %H:%M:%S',  # Space separator without microseconds
                ]
                for fmt in formats:
                    try:
                        return datetime.strptime(
                            time_stamp,
                            fmt).replace(tzinfo=zoneinfo.ZoneInfo(time_zone))
                    except ValueError:
                        continue

                raise ValueError(
                    f"Unsupported timestamp format: '{time_stamp}'")

        except ValueError as e:
            raise ValueError(
                f"Cannot parse '{time_stamp}' as a datetime. Expected formats: "
                f"'YYYY-MM-DDTHH:MM:SS[.ffffff]Z' (UTC) or 'YYYY-MM-DDTHH:MM:SS[.ffffff]' (local)"
            ) from e

    elif isinstance(time_stamp, int):
        try:
            # UNIX timestamps are always in UTC, then convert to target timezone
            return datetime.fromtimestamp(
                time_stamp, tz=zoneinfo.ZoneInfo('UTC')).astimezone(
                    zoneinfo.ZoneInfo(time_zone))
        except (ValueError, OSError) as e:
            raise ValueError(
                f"Cannot convert '{time_stamp}' to datetime. Ensure it's a valid UNIX timestamp."
            ) from e

    else:
        raise TypeError(
            f"Unsupported type '{type(time_stamp)}'. Expected int, str, or datetime."
        )


# Fallback MIME types for common media formats that may not be registered
# in the system's MIME database on certain platforms.
_FALLBACK_MIME_TYPES = {
    '.webp': 'image/webp',
    '.avif': 'image/avif',
    '.heic': 'image/heic',
    '.heif': 'image/heif',
    '.jpg': 'image/jpeg',
    '.jpeg': 'image/jpeg',
    '.png': 'image/png',
    '.gif': 'image/gif',
    '.bmp': 'image/bmp',
    '.svg': 'image/svg+xml',
    '.tiff': 'image/tiff',
    '.tif': 'image/tiff',
    '.ico': 'image/x-icon',
    '.mp4': 'video/mp4',
    '.webm': 'video/webm',
    '.avi': 'video/x-msvideo',
    '.mov': 'video/quicktime',
    '.mkv': 'video/x-matroska',
}


def encode_media_to_base64(media_file_path: str) -> str:
    """
    Encode image or video file to base64 string.

    Args:
        media_file_path (str): Path to the image or video file

    Returns:
        str: Base64 encoded string with data URL prefix

    Raises:
        FileNotFoundError: If image/video file doesn't exist
        ValueError: If file is not a valid format
    """
    import base64
    import mimetypes

    # Expand user path
    media_file_path = os.path.expanduser(media_file_path)

    if not os.path.exists(media_file_path):
        raise FileNotFoundError(f'Image file not found: {media_file_path}')

    if not os.path.isfile(media_file_path):
        raise ValueError(f'Path is not a file: {media_file_path}')

    # Get MIME type, with fallback for formats that may not be registered
    # in the system's MIME database on some platforms (e.g. Linux servers,
    # Docker containers).
    mime_type, _ = mimetypes.guess_type(media_file_path)
    if not mime_type:
        mime_type = _FALLBACK_MIME_TYPES.get(
            os.path.splitext(media_file_path)[1].lower())
    if not mime_type:
        raise ValueError(f'File is not a valid format: {media_file_path}')

    # Read and encode file
    with open(media_file_path, 'rb') as media_file:
        image_data = media_file.read()
        base64_data = base64.b64encode(image_data).decode('utf-8')

    # Return data URL format
    return f'data:{mime_type};base64,{base64_data}'
