"""
Download database and vector index files from HF bucket at container startup.

Files are fetched from: dhammawatthumpra/tipitaka-storage
"""

import os
import sys
from pathlib import Path
from dotenv import load_dotenv

# Add current directory to path so we can import app.config
sys.path.append(str(Path(__file__).resolve().parent))
from app.config import get_settings

settings = get_settings()

# Define paths from settings
DB_PATH = Path(settings.DATABASE_PATH)
QDRANT_DIR = Path(settings.QDRANT_PATH)
SNAPSHOT_DIR = Path(settings.SNAPSHOT_DIR)
MODELS_DIR = Path(__file__).resolve().parent / "models"
TV_PATH = Path(settings.DATA_DIR) / "tipitaka_chunks.tvim"

BUCKET_ID = "dhammawatthumpra/tipitaka-storage"
RERANK_REPO_ID = "jinaai/jina-reranker-v2-base-multilingual"

# Files to check/download
BUCKET_FILES = [
    ("tipitaka_mcu.db", str(DB_PATH)),
    ("qdrant_storage/meta.json", str(QDRANT_DIR / "meta.json")),
    ("qdrant_storage/collection/tipitaka_chunks/storage.sqlite", str(QDRANT_DIR / "collection" / "tipitaka_chunks" / "storage.sqlite")),
    ("qdrant_storage/collection/tipitaka_scripture/storage.sqlite", str(QDRANT_DIR / "collection" / "tipitaka_scripture" / "storage.sqlite")),
]

# Essential Reranker files (ONNX)
RERANK_FILES = [
    "config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "onnx/model_int8.onnx",
]


def check_file_exists(path: Path) -> bool:
    """Check if a file exists locally."""
    return path.exists()


def download_files() -> None:
    """Download missing files from HF bucket."""
    import shutil
    hf_token = os.getenv("HF_TOKEN", "")
    from huggingface_hub import hf_hub_download
    
    # Check and handle Turbovec index file
    if not check_file_exists(TV_PATH):
        packaged_tv = Path(__file__).resolve().parent / "data" / "tipitaka_chunks.tvim"
        if check_file_exists(packaged_tv):
            print(f"Copying packaged Turbovec index: {packaged_tv} -> {TV_PATH}")
            TV_PATH.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy(packaged_tv, TV_PATH)
        else:
            # Fallback: add to bucket files so it attempts downloading from HF Space bucket
            BUCKET_FILES.append(("tipitaka_chunks.tvim", str(TV_PATH)))
            
    # 1. Download Core Assets (DB, Snapshots, and optionally Turbovec index)
    pending_core = []
    for remote, local_path_str in BUCKET_FILES:
        local_path = Path(local_path_str)
        if not check_file_exists(local_path):
            pending_core.append((remote, local_path))

    if pending_core:
        if not hf_token:
            print("WARNING: HF_TOKEN not set — cannot download core assets.")
        else:
            print(f"Downloading {len(pending_core)} core assets from '{BUCKET_ID}'...")
            for remote, dest_path in pending_core:
                try:
                    print(f"Downloading {remote} -> {dest_path}")
                    dest_path.parent.mkdir(parents=True, exist_ok=True)
                    downloaded = hf_hub_download(
                        repo_id=BUCKET_ID,
                        filename=remote,
                        repo_type="dataset",
                        token=hf_token
                    )
                    shutil.copy(downloaded, dest_path)
                except Exception as e:
                    print(f"FAILED to download {remote}: {e}")

    # 2. Download Reranker Assets (if missing)
    rerank_dir = MODELS_DIR / "jina-v2-onnx"
    pending_rerank = []
    for filename in RERANK_FILES:
        local_path = rerank_dir / filename
        if not local_path.exists():
            pending_rerank.append(filename)
            
    if pending_rerank:
        print(f"Downloading {len(pending_rerank)} reranker assets from '{RERANK_REPO_ID}'...")
        for filename in pending_rerank:
            try:
                dest_path = rerank_dir / filename
                print(f"Downloading {filename} -> {dest_path}")
                dest_path.parent.mkdir(parents=True, exist_ok=True)
                downloaded = hf_hub_download(
                    repo_id=RERANK_REPO_ID,
                    filename=filename,
                    token=hf_token if hf_token else None
                )
                shutil.copy(downloaded, dest_path)
            except Exception as e:
                print(f"FAILED to download reranker file {filename}: {e}")


def verify_assets() -> bool:
    """Check that critical files exist."""
    missing = []
    if not DB_PATH.exists():
        missing.append("tipitaka_mcu.db")
    
    if not TV_PATH.exists():
        missing.append("tipitaka_chunks.tvim")
        
    if not (QDRANT_DIR / "meta.json").exists():
        missing.append("qdrant_storage/meta.json")
        
    if not (QDRANT_DIR / "collection" / "tipitaka_chunks" / "storage.sqlite").exists():
        missing.append("qdrant_storage/collection/tipitaka_chunks/storage.sqlite")
        
    if not (QDRANT_DIR / "collection" / "tipitaka_scripture" / "storage.sqlite").exists():
        missing.append("qdrant_storage/collection/tipitaka_scripture/storage.sqlite")

    rerank_model = MODELS_DIR / "jina-v2-onnx" / "onnx" / "model_int8.onnx"
    if not rerank_model.exists():
        missing.append("reranker_model (model_int8.onnx)")

    if missing:
        print(f"WARNING: Missing assets: {missing}")
        return False
    
    print("Success: All critical assets verified.")
    return True


if __name__ == "__main__":
    QDRANT_DIR.mkdir(parents=True, exist_ok=True)
    SNAPSHOT_DIR.mkdir(parents=True, exist_ok=True)
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    
    download_files()
    verify_assets()
    print("Asset management complete.")