import os
import sys
import shutil
from modelscope.hub.snapshot_download import snapshot_download
from dotenv import load_dotenv

load_dotenv()

# 使用持久化的缓存目录（挂载到宿主机）
CACHE_DIR = os.getenv('MODEL_CACHE_DIR', '/app/model_cache')
os.makedirs(CACHE_DIR, exist_ok=True)


def discover_models():
    """从环境变量中发现所有模型"""
    models = {}
    for key, value in os.environ.items():
        if key.endswith('_MODEL_ID'):
            prefix = key[:-9]
            service_name = os.getenv(f'{prefix}_SERVICE_NAME', f'{prefix.lower()}_classifier')
            service_folder = service_name.replace('_classifier', '')

            models[prefix.lower()] = {
                'prefix': prefix,
                'model_id': value,
                'port': os.getenv(f'{prefix}_PORT', '5002'),
                'service_name': service_name,
                'model_dir': os.getenv(f'{prefix}_MODEL_DIR', f'/app/services/bert_{service_folder}/model'),
                'service_folder': service_folder,
                'cache_dir': os.path.join(CACHE_DIR, prefix.lower())
            }
    return models


def check_model_complete(model_dir):
    """检查模型是否完整"""
    if not os.path.exists(model_dir):
        return False

    try:
        files = os.listdir(model_dir)
        if not files:
            return False

        has_weights = any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in files)
        has_config = any(f == 'config.json' for f in files)

        if has_weights and has_config:
            return True
        return False

    except Exception:
        return False


def find_model_files(cache_dir):
    """在下载的缓存目录中查找实际的模型文件"""
    # 检查是否直接有模型文件
    direct_files = [f for f in os.listdir(cache_dir) if
                    not f.startswith('.') and os.path.isfile(os.path.join(cache_dir, f))]
    if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in direct_files):
        return cache_dir

    # 检查 models/ 子目录（ModelScope 的标准结构）
    models_dir = os.path.join(cache_dir, 'models')
    if os.path.exists(models_dir):
        for root, dirs, files in os.walk(models_dir):
            if 'snapshots' in root:
                for f in files:
                    if f.endswith(('.safetensors', '.bin', '.pt', '.pth')):
                        return root
            if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in files):
                return root

    # 检查是否存在单层子目录
    subdirs = [d for d in os.listdir(cache_dir) if os.path.isdir(os.path.join(cache_dir, d)) and not d.startswith('.')]
    for subdir in subdirs:
        subpath = os.path.join(cache_dir, subdir)
        if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in os.listdir(subpath)):
            return subpath

    # 递归查找任何包含模型文件的目录
    for root, dirs, files in os.walk(cache_dir):
        if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in files):
            return root

    return None


def copy_model_from_cache(cache_dir, target_dir):
    """从缓存目录复制模型到目标目录"""
    src_dir = find_model_files(cache_dir)

    if src_dir is None:
        print(f"  ✗ Could not find model files in cache: {cache_dir}")
        return False

    print(f"  ✓ Found cached model files in: {src_dir}")

    os.makedirs(target_dir, exist_ok=True)

    copied_count = 0
    for file in os.listdir(src_dir):
        src_file = os.path.join(src_dir, file)
        if os.path.isfile(src_file):
            dest_file = os.path.join(target_dir, file)
            shutil.copy2(src_file, dest_file)
            copied_count += 1

    print(f"  ✓ Copied {copied_count} files to {target_dir}")
    return copied_count > 0


def download_model_with_fallback(name, config):
    """下载单个模型，带备选方案"""
    print("=" * 50)
    print(f"Processing {name} model...")
    print(f"Model ID: {config['model_id']}")
    print(f"Target: {config['model_dir']}")
    print(f"Cache: {config['cache_dir']}")
    print("=" * 50)

    try:
        os.makedirs(config['model_dir'], exist_ok=True)
        os.makedirs(config['cache_dir'], exist_ok=True)

        # 1. 检查目标目录是否已有完整模型
        if check_model_complete(config['model_dir']):
            print(f"✅ Model already exists in target, skipping")
            return True

        # 2. 检查缓存目录是否有模型
        if check_model_complete(config['cache_dir']):
            print(f"✅ Found complete model in local cache")
            print(f"   Copying from cache to target...")

            if copy_model_from_cache(config['cache_dir'], config['model_dir']):
                if check_model_complete(config['model_dir']):
                    print(f"✅ Model copied successfully from cache")
                    return True
                else:
                    print(f"⚠️  Model copied but appears incomplete, will re-download")

        # 3. 缓存中没有，需要下载
        print(f"📥 Downloading model from ModelScope...")

        download_success = False

        # 尝试方式1：直接下载
        try:
            snapshot_download(
                model_id=config['model_id'],
                cache_dir=config['cache_dir'],
                revision="master",
                ignore_file_pattern=[".git", ".gitattributes"]
            )
            print(f"✅ Download completed")
            download_success = True
        except Exception as e:
            print(f"⚠️  Primary download failed: {e}")

            # 尝试方式2：使用 HubApi（如果有 token）
            token = os.getenv('MODELSCOPE_TOKEN', '')
            if token:
                print("   Trying alternative download with HubApi...")
                try:
                    from modelscope.hub.api import HubApi
                    api = HubApi()
                    api.login(token=token)

                    snapshot_download(
                        model_id=config['model_id'],
                        cache_dir=config['cache_dir'],
                        revision="master",
                        ignore_file_pattern=[".git", ".gitattributes"]
                    )
                    print(f"✅ Alternative download completed")
                    download_success = True
                except Exception as alt_error:
                    print(f"✗ Alternative download also failed: {alt_error}")
                    raise
            else:
                print("✗ No MODELSCOPE_TOKEN set, cannot use alternative method")
                raise

        if not download_success:
            return False

        # 验证并复制
        if not check_model_complete(config['cache_dir']):
            print(f"⚠️  Downloaded model appears incomplete in cache")
            print(f"   Attempting to copy anyway...")

        print(f"Copying from cache to target...")
        if copy_model_from_cache(config['cache_dir'], config['model_dir']):
            if check_model_complete(config['model_dir']):
                print(f"✅ Model deployed successfully!")
                return True
            else:
                print(f"⚠️  Model deployed but appears incomplete")
                return True
        else:
            print(f"✗ Failed to copy model to target")
            return False

    except Exception as e:
        print(f"✗ Failed to process {name} model: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    print("=" * 50)
    print("Model Download Script")
    print("=" * 50)
    print(f"Cache directory: {CACHE_DIR}")

    force_download = os.getenv('FORCE_MODEL_DOWNLOAD', 'false').lower() == 'true'
    if force_download:
        print("⚠️  FORCE_MODEL_DOWNLOAD=true - will re-download all models")

    print("\nDiscovering models from .env...")
    models = discover_models()

    if not models:
        print("⚠️  No models found in .env file!")
        sys.exit(1)

    print(f"\nFound {len(models)} model(s):")
    all_exist = True
    for name, config in models.items():
        target_exists = check_model_complete(config['model_dir'])
        cache_exists = check_model_complete(config['cache_dir'])

        print(f"  {name}: {config['model_id']}")
        print(f"    Target: {'✅' if target_exists else '❌'} {config['model_dir']}")
        print(f"    Cache:  {'✅' if cache_exists else '❌'} {config['cache_dir']}")

        if not target_exists and not cache_exists:
            all_exist = False
    print("")

    if all_exist and not force_download:
        print("✅ All models exist!")
        print("   Skipping downloads.")
        sys.exit(0)

    if force_download:
        print("⚠️  Force download enabled, clearing caches...")
        for name, config in models.items():
            if os.path.exists(config['cache_dir']):
                print(f"  Clearing cache for {name}")
                shutil.rmtree(config['cache_dir'])
                os.makedirs(config['cache_dir'], exist_ok=True)

    print("\n" + "=" * 50)
    print("Processing models...")
    print("=" * 50)

    success_count = 0
    for name, config in models.items():
        if download_model_with_fallback(name, config):
            success_count += 1
        print("")

    print("=" * 50)
    print(f"Processing completed: {success_count}/{len(models)} models")
    print("=" * 50)

    if success_count < len(models):
        sys.exit(1)


if __name__ == "__main__":
    main()