import asyncio
import json
import logging
import os
import time
from pathlib import Path
from typing import Dict, List, Optional

import huggingface_hub

logger = logging.getLogger(__name__)

_FORBIDDEN_PATHS = frozenset({
    "/bin", "/boot", "/dev", "/etc", "/lib", "/lib64",
    "/proc", "/run", "/sbin", "/sys", "/usr", "/var",
})


def get_default_models_dir() -> Path:
    env = os.environ.get("OPENARC_MODELS_DIR", "").strip()
    if env:
        return Path(env).expanduser().resolve()
    return (Path.home() / ".cache" / "openarc" / "models").resolve()


def default_download_path(model_name: str) -> str:
    """default target dir for a hf repo id, matching what GET /openarc/models scans"""
    return str(get_default_models_dir() / model_name.replace("/", "__"))


def validate_download_path(path: str) -> str:
    """Validate and resolve a download path. Raises ValueError if unsafe."""
    resolved = Path(path).resolve()

    if resolved.parent == resolved:
        raise ValueError(f"Cannot download to filesystem root: {resolved}")

    if str(resolved) in _FORBIDDEN_PATHS:
        raise ValueError(f"Cannot download to system directory: {resolved}")

    resolved.mkdir(parents=True, exist_ok=True)

    if not resolved.is_dir():
        raise ValueError(f"Path is not a directory: {resolved}")

    if not os.access(resolved, os.W_OK):
        raise ValueError(f"Path is not writable: {resolved}")

    return str(resolved)


class DownloadTask:
    def __init__(self, model_name: str, path: Optional[str] = None):
        self.model_name = model_name
        self.path = path
        self.status = "pending"  # pending, downloading, paused, cancelled, completed, error
        self.task: Optional[asyncio.Task] = None
        self.progress = 0.0
        self.total_bytes = 0
        self.downloaded_bytes = 0
        self.total_files = 0
        self.completed_files = 0
        self.error_msg = ""
        self.started_at = time.time()
        self.completed_at: Optional[float] = None
        self.download_speed = 0

        # async control signals
        self._cancel = asyncio.Event()
        self._pause = asyncio.Event()
        self._pause.set()  # not paused by default


class Downloader:
    CLEANUP_AFTER_SECONDS = 3600

    def __init__(self):
        self.tasks: Dict[str, DownloadTask] = {}
        self._lock = asyncio.Lock()

    async def start(self, model_name: str) -> bool:
        """Start a new download. Returns False if one is already active."""
        path = validate_download_path(default_download_path(model_name))

        async with self._lock:
            self._cleanup_stale()

            existing = self.tasks.get(model_name)
            if existing and existing.status in ("downloading", "paused"):
                return False

            task = DownloadTask(model_name, path)
            self.tasks[model_name] = task
            task.status = "downloading"
            task.task = asyncio.create_task(self._download_worker(task))
            return True

    async def pause(self, model_name: str) -> bool:
        """Pause an active download. The worker blocks between files."""
        async with self._lock:
            task = self.tasks.get(model_name)
            if not task or task.status != "downloading":
                return False
            task.status = "paused"
            task._pause.clear()
            return True

    async def resume(self, model_name: str) -> bool:
        """Resume a paused download."""
        async with self._lock:
            task = self.tasks.get(model_name)
            if not task or task.status != "paused":
                return False
            task.status = "downloading"
            task._pause.set()
            return True

    async def cancel(self, model_name: str) -> bool:
        """Cancel an active or paused download."""
        async with self._lock:
            task = self.tasks.get(model_name)
            if not task or task.status in ("completed", "cancelled"):
                return False
            task.status = "cancelled"
            task.completed_at = time.time()
            task._cancel.set()
            task._pause.set()  # unblock worker so it can exit
            if task.task:
                task.task.cancel()
            return True

    def list_tasks(self) -> List[Dict]:
        self._cleanup_stale()
        return [
            {
                "model_name": t.model_name,
                "total_size": t.total_bytes,
                "downloaded_size": t.downloaded_bytes,
                "status": t.status,
                "progress": int(t.progress),
                "download_speed": t.download_speed,
                "path": t.path or "",
                "started_at": str(t.started_at),
                "finished_at": str(t.completed_at) if t.completed_at is not None else None,
            }
            for t in self.tasks.values()
        ]

    def _cleanup_stale(self):
        """Remove terminal tasks older than CLEANUP_AFTER_SECONDS."""
        now = time.time()
        stale = [
            name
            for name, t in self.tasks.items()
            if t.status in ("completed", "cancelled", "error")
            and t.completed_at
            and (now - t.completed_at) > self.CLEANUP_AFTER_SECONDS
        ]
        for name in stale:
            del self.tasks[name]

    async def _download_worker(self, task: DownloadTask):
        try:
            await self._download_files(task)
            if not task._cancel.is_set():
                await asyncio.to_thread(self._write_hf_metadata, task)
                task.status = "completed"
                task.progress = 100.0
                task.completed_at = time.time()
                logger.info(f"Download completed for {task.model_name}")
        except asyncio.CancelledError:
            logger.info(f"Download {task.status} for {task.model_name}")
            raise
        except Exception as e:
            task.status = "error"
            task.error_msg = str(e)
            task.completed_at = time.time()
            logger.error(f"Download failed for {task.model_name}: {e}")

    @staticmethod
    def _write_hf_metadata(task: DownloadTask):
        """stamp openarc.json with the source hf author/repo so the listing can show provenance"""
        if not task.path:
            return
        author, _, repo = task.model_name.partition("/")
        if not repo:
            author, repo = "", author
        config_path = Path(task.path) / "openarc.json"
        config: Dict = {}
        if config_path.exists():
            try:
                with open(config_path, "r", encoding="utf-8") as f:
                    config = json.load(f) or {}
            except Exception:
                config = {}
        config["hf_author"] = author
        config["hf_repo"] = repo
        try:
            with open(config_path, "w", encoding="utf-8") as f:
                json.dump(config, f, indent=4)
        except Exception as e:
            logger.warning(f"failed to write openarc.json for {task.model_name}: {e}")

    async def _download_files(self, task: DownloadTask):
        """Async download loop. Each file is fetched in a thread so the event loop stays free,
        but pause/cancel work between files without blocking the loop.
        """
        api = huggingface_hub.HfApi()

        repo = await asyncio.to_thread(api.repo_info, task.model_name, files_metadata=True)
        siblings = repo.siblings or []

        files = [(f.rfilename, getattr(f, "size", None) or 0) for f in siblings]
        task.total_bytes = sum(size for _, size in files)
        task.total_files = len(files)

        if task._cancel.is_set():
            return

        for filename, file_size in files:
            if task._cancel.is_set():
                return

            # block here while paused
            await task._pause.wait()

            if task._cancel.is_set():
                return

            # run blocking download in a thread so the event loop stays free
            file_start = time.time()
            await asyncio.to_thread(
                huggingface_hub.hf_hub_download,
                repo_id=task.model_name,
                filename=filename,
                local_dir=task.path,
            )
            elapsed = time.time() - file_start
            if elapsed > 0:
                task.download_speed = int(file_size / elapsed)

            task.completed_files += 1
            task.downloaded_bytes += file_size
            if task.total_bytes > 0:
                task.progress = (task.downloaded_bytes / task.total_bytes) * 100.0


global_downloader = Downloader()