import shutil
import subprocess
import tempfile
from pathlib import Path

import shortuuid
from pydantic import BaseModel

from harbor.models.task.id import GitTaskId, LocalTaskId


class TaskDownloadConfig(BaseModel):
    git_commit_id: str | None = None
    source_path: Path
    target_path: Path


class TasksClient:
    CACHE_DIR = Path("~/.cache/sandbox/tasks").expanduser()

    def _copy_task_source_to_target(self, source_path: Path, target_path: Path) -> None:
        if target_path.exists():
            shutil.rmtree(target_path)

        shutil.copytree(source_path, target_path)

    def _download_tasks_from_git_url(
        self, git_url: str, task_download_configs: list[TaskDownloadConfig]
    ):
        head_task_download_configs = [
            task_download_config
            for task_download_config in task_download_configs
            if task_download_config.git_commit_id is None
        ]

        commit_task_download_configs = {}
        for task_download_config in task_download_configs:
            commit_id = task_download_config.git_commit_id
            if commit_id is not None:
                commit_task_download_configs.setdefault(commit_id, []).append(
                    task_download_config
                )

        with tempfile.TemporaryDirectory() as temp_dir:
            temp_dir = Path(temp_dir)

            subprocess.run(
                ["git", "clone", git_url, temp_dir],
                check=True,
                capture_output=True,
                cwd=temp_dir,
            )

            for task_download_config in head_task_download_configs:
                self._copy_task_source_to_target(
                    temp_dir / task_download_config.source_path,
                    task_download_config.target_path,
                )

            for (
                git_commit_id,
                task_download_configs,
            ) in commit_task_download_configs.items():
                subprocess.run(
                    ["git", "checkout", git_commit_id],
                    check=True,
                    capture_output=True,
                    cwd=temp_dir,
                )
                for task_download_config in task_download_configs:
                    self._copy_task_source_to_target(
                        temp_dir / task_download_config.source_path,
                        task_download_config.target_path,
                    )

    def download_tasks(
        self,
        task_ids: list[GitTaskId | LocalTaskId],
        overwrite: bool = False,
        output_dir: Path | None = None,
    ) -> list[Path]:
        output_dir = output_dir or self.CACHE_DIR

        local_task_ids = [
            task_id for task_id in task_ids if isinstance(task_id, LocalTaskId)
        ]

        for local_task_id in local_task_ids:
            if not local_task_id.path.exists():
                raise FileNotFoundError(f"Local task {local_task_id.path} not found")

        git_task_ids = [
            task_id for task_id in task_ids if isinstance(task_id, GitTaskId)
        ]

        target_paths = {
            task_id: output_dir / shortuuid.uuid(str(task_id)) / task_id.path.name
            for task_id in git_task_ids
        }

        download_task_ids = {
            task_id: path
            for task_id, path in target_paths.items()
            if not path.exists()
            or overwrite
            or task_id.git_commit_id is None
            or (path.exists() and not any(path.iterdir()))
        }

        git_urls = {
            task_id.git_url
            for task_id in download_task_ids
            if task_id.git_url is not None
        }

        for git_url in git_urls:
            self._download_tasks_from_git_url(
                git_url=git_url,
                task_download_configs=[
                    TaskDownloadConfig(
                        git_commit_id=task_id.git_commit_id,
                        source_path=task_id.path,
                        target_path=path,
                    )
                    for task_id, path in target_paths.items()
                ],
            )

        result_paths = []
        for task_id in task_ids:
            if isinstance(task_id, GitTaskId):
                result_paths.append(target_paths[task_id])
            elif isinstance(task_id, LocalTaskId):
                result_paths.append(task_id.path)
            else:
                raise TypeError(f"Unsupported task id type: {type(task_id)}")
        return result_paths
