#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
GitCode Hub SDK - 类似huggingface_hub的SDK接口

这个模块提供了类似huggingface_hub的编程接口，
让用户可以通过Python代码直接使用GitCode平台的功能。
"""

import os
from typing import Optional, Union, List, Dict, Any
from pathlib import Path

# 设置OpenMind Hub的API端点为GitCode
os.environ["OPENMIND_HUB_ENDPOINT"] = "https://hub.gitcode.com"
# 设置缓存目录
cache_dir = os.path.expanduser("~/.cache/gitcode")
os.makedirs(cache_dir, exist_ok=True)
os.environ["HF_HOME"] = cache_dir

from openmind_hub import snapshot_download as om_snapshot_download
from openmind_hub import om_hub_download, upload_folder as om_upload_folder, create_repo

# 导入数据集相关模块
try:
    from openmind import OmDataset  # type: ignore
    from openmind.utils.hub import OM_DATASETS_CACHE  # type: ignore
    DATASET_SUPPORT = True
except ImportError:
    DATASET_SUPPORT = False
    OmDataset = None
    OM_DATASETS_CACHE = None

try:
    from .config import config
except ImportError:
    from config import config


def _normalize_repo_id(repo_id: str) -> str:
    """标准化仓库ID，处理三层格式转换"""
    parts = repo_id.split('/')
    
    # 如果是三层格式（如 wuyw/Qwen3-Reranker/0.6B-test）
    # 转换为特殊格式（如 wuyw%2FQwen3-Reranker/0.6B-test）
    if len(parts) >= 3:
        # 只编码第一个斜杠，保留后面的斜杠
        first_part = parts[0]
        second_part = parts[1]
        remaining_parts = parts[2:]
        
        # 构建新格式：第一部分%2F第二部分/其余部分
        normalized = first_part + '%2F' + second_part
        if remaining_parts:
            normalized += '/' + '/'.join(remaining_parts)
        
        return normalized
    
    # 二层或单层格式直接返回
    return repo_id


def _get_token() -> Optional[str]:
    """获取保存的认证token"""
    try:
        credentials = config.get_credentials()
        if credentials and 'token' in credentials:
            return credentials['token']
    except Exception:
        pass
    return None


def snapshot_download(
    repo_id: str,
    revision: Optional[str] = None,
    cache_dir: Optional[Union[str, Path]] = None,
    local_dir: Optional[Union[str, Path]] = None,
    local_dir_use_symlinks: Union[bool, str] = "auto",
    library_name: Optional[str] = None,
    library_version: Optional[str] = None,
    user_agent: Optional[Union[str, Dict[str, str]]] = None,
    proxies: Optional[Dict[str, str]] = None,
    etag_timeout: float = 10,
    resume_download: bool = False,
    force_download: bool = False,
    token: Optional[Union[str, bool]] = None,
    local_files_only: bool = False,
    allow_patterns: Optional[Union[List[str], str]] = None,
    ignore_patterns: Optional[Union[List[str], str]] = None,
    max_workers: int = 8,
    tqdm_class: Optional[Any] = None,
    disable_tqdm: bool = False
) -> str:
    """
    从GitCode Hub下载整个仓库的快照到本地目录
    
    参数:
        repo_id (str): 仓库ID，格式为 "username/repo-name" 或 "username/namespace/repo-name"
        revision (str, 可选): 指定版本/分支/标签，默认为 "main"
        cache_dir (str 或 Path, 可选): 缓存目录路径
        local_dir (str 或 Path, 可选): 下载到的本地目录路径
        local_dir_use_symlinks (bool 或 str, 可选): 是否使用符号链接，默认为 "auto"
        library_name (str, 可选): 调用库的名称
        library_version (str, 可选): 调用库的版本
        user_agent (str 或 dict, 可选): 用户代理
        proxies (dict, 可选): 代理配置
        etag_timeout (float, 可选): ETag超时时间，默认10秒
        resume_download (bool, 可选): 是否断点续传，默认False
        force_download (bool, 可选): 是否强制重新下载，默认False
        token (str 或 bool, 可选): 认证token，如果为None则尝试使用保存的token
        local_files_only (bool, 可选): 仅使用本地文件，默认False
        allow_patterns (list 或 str, 可选): 允许下载的文件模式
        ignore_patterns (list 或 str, 可选): 忽略的文件模式
        max_workers (int, 可选): 最大并发数，默认8
        tqdm_class (可选): 进度条类
        disable_tqdm (bool, 可选): 是否禁用进度条，默认False
    
    返回:
        str: 下载完成的本地目录路径
    
    示例:
        >>> from gitcode_hub import snapshot_download
        >>> local_path = snapshot_download("wuyw/Qwen3-Reranker-0.6B-test", local_dir="./models")
    """
    # 标准化仓库ID（处理三层格式）
    normalized_repo_id = _normalize_repo_id(repo_id)
    
    # 如果没有提供token，尝试使用保存的token
    if token is None:
        token = _get_token()
    
    # 构建参数字典
    kwargs = {
        'repo_id': normalized_repo_id,
    }
    
    # 添加可选参数
    if revision is not None:
        kwargs['revision'] = revision
    if cache_dir is not None:
        kwargs['cache_dir'] = str(cache_dir)
    if local_dir is not None:
        kwargs['local_dir'] = str(local_dir)
    if token is not None:
        kwargs['token'] = token
    
    # 其他参数
    kwargs.update({
        'local_dir_use_symlinks': local_dir_use_symlinks,
        'library_name': library_name or "gitcode_hub",
        'library_version': library_version,
        'user_agent': user_agent,
        'proxies': proxies,
        'etag_timeout': etag_timeout,
        'resume_download': resume_download,
        'force_download': force_download,
        'local_files_only': local_files_only,
        'allow_patterns': allow_patterns,
        'ignore_patterns': ignore_patterns,
        'max_workers': max_workers,
        'tqdm_class': tqdm_class,
        'disable_tqdm': disable_tqdm,
    })
    
    try:
        # 首先尝试公开下载（适用于公开仓库）
        if not local_files_only and token is None:
            try:
                result = om_snapshot_download(**{k: v for k, v in kwargs.items() if v is not None})
                return result
            except Exception as e:
                error_msg = str(e)
                # 如果是认证问题，尝试使用token
                if ("403" in error_msg or "FORBIDDEN" in error_msg or 
                    "no scopes" in error_msg or "unauthorized" in error_msg.lower()):
                    saved_token = _get_token()
                    if saved_token:
                        kwargs['token'] = saved_token
                    else:
                        raise Exception(f"仓库访问需要认证，但未找到保存的token。请先使用 'gitcode login' 登录，或在调用时提供token参数。")
                else:
                    raise e
        
        # 使用token下载
        result = om_snapshot_download(**{k: v for k, v in kwargs.items() if v is not None})
        return result
        
    except Exception as e:
        error_msg = str(e)
        if "401" in error_msg or "403" in error_msg:
            raise Exception(f"认证失败：{error_msg}。请检查token是否正确，或使用 'gitcode login' 重新登录。")
        elif "404" in error_msg:
            raise Exception(f"仓库不存在：{repo_id}。请检查仓库名称是否正确。")
        else:
            raise Exception(f"下载失败：{error_msg}")


def hub_download_url(
    repo_id: str,
    filename: str,
    revision: Optional[str] = None,
    repo_type: Optional[str] = None,
) -> str:
    """
    获取GitCode Hub上文件的下载URL
    
    参数:
        repo_id (str): 仓库ID
        filename (str): 文件名
        revision (str, 可选): 版本/分支/标签
        repo_type (str, 可选): 仓库类型
    
    返回:
        str: 文件的下载URL
    """
    normalized_repo_id = _normalize_repo_id(repo_id)
    base_url = "https://hub.gitcode.com"
    
    if revision is None:
        revision = "main"
    
    if repo_type is None:
        repo_type = "model"
    
    # 构建URL
    url = f"{base_url}/{normalized_repo_id}/resolve/{revision}/{filename}"
    return url


def download_file(
    repo_id: str,
    filename: str,
    local_dir: Optional[Union[str, Path]] = None,
    revision: Optional[str] = None,
    token: Optional[str] = None,
    force_download: bool = False,
) -> str:
    """
    从GitCode Hub下载单个文件
    
    参数:
        repo_id (str): 仓库ID
        filename (str): 文件名
        local_dir (str 或 Path, 可选): 本地目录
        revision (str, 可选): 版本/分支/标签
        token (str, 可选): 认证token
        force_download (bool, 可选): 是否强制重新下载
    
    返回:
        str: 下载的文件路径
    """
    # 标准化仓库ID
    normalized_repo_id = _normalize_repo_id(repo_id)
    
    # 如果没有提供token，尝试使用保存的token
    if token is None:
        token = _get_token()
    
    # 构建参数
    kwargs = {
        'repo_id': normalized_repo_id,
        'filename': filename,
    }
    
    if local_dir is not None:
        kwargs['local_dir'] = str(local_dir)
    if revision is not None:
        kwargs['revision'] = revision
    if token is not None:
        kwargs['token'] = token
    if force_download:
        kwargs['force_download'] = force_download
    
    try:
        # 首先尝试公开下载
        if token is None:
            try:
                result = om_hub_download(**{k: v for k, v in kwargs.items() if v is not None})
                return result
            except Exception as e:
                error_msg = str(e)
                if ("403" in error_msg or "FORBIDDEN" in error_msg or 
                    "no scopes" in error_msg or "unauthorized" in error_msg.lower()):
                    saved_token = _get_token()
                    if saved_token:
                        kwargs['token'] = saved_token
                    else:
                        raise Exception(f"文件访问需要认证，但未找到保存的token。请先使用 'gitcode login' 登录。")
                else:
                    raise e
        
        # 使用token下载
        result = om_hub_download(**{k: v for k, v in kwargs.items() if v is not None})
        return result
        
    except Exception as e:
        error_msg = str(e)
        if "401" in error_msg or "403" in error_msg:
            raise Exception(f"认证失败：{error_msg}")
        elif "404" in error_msg:
            raise Exception(f"文件不存在：{repo_id}/{filename}")
        else:
            raise Exception(f"下载失败：{error_msg}")


def upload_folder(
    folder_path: Union[str, Path],
    repo_id: str,
    token: Optional[str] = None,
    repo_type: Optional[str] = None,
    revision: Optional[str] = None,
    commit_message: Optional[str] = None,
    commit_description: Optional[str] = None,
    path_in_repo: str = "./",
    ignore_patterns: Optional[List[str]] = None,
) -> str:
    """
    上传文件夹到GitCode Hub
    
    参数:
        folder_path (str 或 Path): 本地文件夹路径
        repo_id (str): 仓库ID
        token (str, 可选): 认证token
        repo_type (str, 可选): 仓库类型
        revision (str, 可选): 分支名
        commit_message (str, 可选): 提交消息
        commit_description (str, 可选): 提交描述
        path_in_repo (str, 可选): 在仓库中的路径，默认为根目录
        ignore_patterns (List[str], 可选): 要忽略的文件模式
    
    返回:
        str: 提交的URL或ID
    
    示例:
        >>> upload_folder("./my-model/", "username/repo-name")
        >>> upload_folder("./data/", "username/repo", path_in_repo="datasets/")
    """
    # 标准化仓库ID
    normalized_repo_id = _normalize_repo_id(repo_id)
    
    # 转换为Path对象
    folder_path = Path(folder_path)
    
    if not folder_path.exists():
        raise FileNotFoundError(f"文件夹不存在: {folder_path}")
    
    if not folder_path.is_dir():
        raise NotADirectoryError(f"路径不是目录: {folder_path}")
    
    # 如果没有提供token，尝试使用保存的token
    if token is None:
        token = _get_token()
        if token is None:
            raise Exception("上传需要认证token，请先使用 'gitcode login' 登录，或提供token参数。")
    
    # 直接使用原始目录，或者创建临时目录来重新组织结构
    import tempfile
    import shutil
    
    if path_in_repo == "./" or path_in_repo == "." or path_in_repo == "":
        # 如果要上传到根目录，直接使用源文件夹
        upload_path = str(folder_path)
    else:
        # 如果要上传到特定路径，需要重新组织目录结构
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            
            # 创建目标路径
            target_path = temp_path / path_in_repo.strip('./')
            target_path.parent.mkdir(parents=True, exist_ok=True)
            
            # 复制整个目录树
            shutil.copytree(folder_path, target_path, dirs_exist_ok=True)
            
            upload_path = str(temp_path)
    
    try:
        # 使用openmind_hub的upload_folder上传
        commit_msg = commit_message or f"Upload folder {folder_path.name}"
        result = om_upload_folder(
            repo_id=normalized_repo_id,
            folder_path=upload_path,
            token=token,
            commit_message=commit_msg
        )
        
        return result
        
    except Exception as e:
        error_msg = str(e)
        if "401" in error_msg or "403" in error_msg:
            raise Exception(f"认证失败：{error_msg}")
        elif "404" in error_msg:
            raise Exception(f"仓库不存在：{repo_id}")
        else:
            raise Exception(f"上传失败：{error_msg}")


def create_repository(
    repo_id: str,
    token: Optional[str] = None,
    private: bool = False,
    repo_type: str = "model",
    exist_ok: bool = False,
    space_sdk: Optional[str] = None,
    space_hardware: Optional[str] = None,
    space_storage: Optional[str] = None,
    space_sleep_time: Optional[int] = None,
    space_secrets: Optional[List[Dict]] = None,
    space_variables: Optional[List[Dict]] = None,
) -> str:
    """
    在GitCode Hub上创建仓库
    
    参数:
        repo_id (str): 仓库ID
        token (str, 可选): 认证token
        private (bool, 可选): 是否为私有仓库，默认False
        repo_type (str, 可选): 仓库类型，默认"model"
        exist_ok (bool, 可选): 如果仓库已存在是否报错，默认False
        其他参数: 主要用于Space类型仓库
    
    返回:
        str: 仓库的URL
    """
    # 标准化仓库ID
    normalized_repo_id = _normalize_repo_id(repo_id)
    
    # 如果没有提供token，尝试使用保存的token
    if token is None:
        token = _get_token()
        if token is None:
            raise Exception("创建仓库需要认证token，请先使用 'gitcode login' 登录，或提供token参数。")
    
    try:
        result = create_repo(
            repo_id=normalized_repo_id,
            token=token,
            private=private,
            repo_type=repo_type,
            exist_ok=exist_ok,
        )
        return result
        
    except Exception as e:
        error_msg = str(e)
        if "401" in error_msg or "403" in error_msg:
            raise Exception(f"认证失败：{error_msg}")
        elif "409" in error_msg or "exists" in error_msg.lower():
            if exist_ok:
                return f"https://hub.gitcode.com/{repo_id}"
            else:
                raise Exception(f"仓库已存在：{repo_id}")
        else:
            raise Exception(f"创建仓库失败：{error_msg}")


def load_dataset(
    path: str,
    name: Optional[str] = None,
    data_dir: Optional[str] = None,
    data_files: Optional[Union[str, List[str], Dict]] = None,
    split: Optional[str] = None,
    cache_dir: Optional[Union[str, Path]] = None,
    token: Optional[str] = None,
    streaming: bool = False,
    revision: Optional[str] = None,
    **kwargs
) -> Any:
    """
    从GitCode Hub加载数据集
    
    参数:
        path (str): 数据集路径，格式为 "username/dataset-name"
        name (str, 可选): 数据集名称
        data_dir (str, 可选): 数据目录路径
        data_files (str, List[str], Dict, 可选): 数据文件
        split (str, 可选): 数据集分割，如 "train", "test", "validation"
        cache_dir (str 或 Path, 可选): 缓存目录路径，默认使用OM_DATASETS_CACHE
        token (str, 可选): 认证token，如果为None则尝试使用保存的token
        streaming (bool, 可选): 是否以流式方式加载数据集，默认False
        revision (str, 可选): 指定版本/分支/标签，默认为 "main"
        **kwargs: 其他传递给OmDataset.load_dataset的参数
    
    返回:
        Dataset或DatasetDict: 加载的数据集对象
    
    示例:
        >>> from gitcode_hub import load_dataset
        >>> dataset = load_dataset("wuyw/my-dataset", split="train")
        >>> dataset = load_dataset("wuyw/my-dataset", streaming=True)
    
    异常:
        ImportError: 如果openmind库未安装
        Exception: 各种下载和认证错误
    """
    # 检查数据集支持
    if not DATASET_SUPPORT:
        raise ImportError("数据集功能需要安装openmind库。请运行: pip install openmind")
    
    # 标准化数据集路径（处理三层格式）
    normalized_path = _normalize_repo_id(path)
    
    # 如果没有提供token，尝试使用保存的token
    if token is None:
        token = _get_token()
    
    # 设置缓存目录
    if cache_dir is None:
        cache_dir = OM_DATASETS_CACHE
    
    # 构建参数字典
    load_kwargs = {
        'path': normalized_path,
        'cache_dir': str(cache_dir),
        'streaming': streaming,
    }
    
    # 添加可选参数
    if name is not None:
        load_kwargs['name'] = name
    if data_dir is not None:
        load_kwargs['data_dir'] = data_dir
    if data_files is not None:
        load_kwargs['data_files'] = data_files
    if split is not None:
        load_kwargs['split'] = split
    if token is not None:
        load_kwargs['token'] = token
    if revision is not None:
        load_kwargs['revision'] = revision
    
    # 合并额外的关键字参数
    load_kwargs.update(kwargs)
    
    try:
        # 首先尝试公开访问（适用于公开数据集）
        if token is None:
            try:
                dataset = OmDataset.load_dataset(**load_kwargs)
                return dataset
            except Exception as e:
                error_msg = str(e)
                # 如果是认证问题，尝试使用保存的token
                if ("403" in error_msg or "FORBIDDEN" in error_msg or 
                    "no scopes" in error_msg or "unauthorized" in error_msg.lower() or
                    "401" in error_msg):
                    saved_token = _get_token()
                    if saved_token:
                        load_kwargs['token'] = saved_token
                    else:
                        raise Exception(f"数据集访问需要认证，但未找到保存的token。请先使用 'gitcode login' 登录，或在调用时提供token参数。原始错误：{error_msg}")
                else:
                    raise e
        
        # 使用token加载数据集
        dataset = OmDataset.load_dataset(**load_kwargs)
        return dataset
        
    except Exception as e:
        error_msg = str(e)
        if "401" in error_msg or "403" in error_msg:
            raise Exception(f"认证失败：{error_msg}。请检查token是否正确，或使用 'gitcode login' 重新登录。")
        elif "404" in error_msg:
            raise Exception(f"数据集不存在：{path}。请检查数据集路径是否正确。")
        elif "ImportError" in error_msg or "No module named" in error_msg:
            raise ImportError("数据集功能需要安装openmind库。请运行: pip install openmind")
        else:
            raise Exception(f"加载数据集失败：{error_msg}")


# 为了兼容性，导出常用函数
__all__ = [
    'snapshot_download',
    'hub_download_url', 
    'download_file',
    'upload_folder',
    'create_repository',
    'load_dataset',
]