from typing import Dict

from ._http import HttpClient


class RestClient:
    def __init__(self, http_client: HttpClient, user_info: dict):
        self.http_client: HttpClient = http_client
        self.user_info = user_info

    def get_user_info(self) -> Dict:
        """获取用户信息"""
        return self.user_info

    def list_weights(self, page=1, page_size=20) -> Dict:
        """
        列出模型权重列表

        Parameters
        ----------
        page: 当前页，默认为 1
        page_size: 每页的大小，默认为 20

        Returns
        -------

        """
        return self.http_client.list_weights(self.user_info["user_id"], page, page_size)

    def get_archive_url(self, checkpoint_id: str) -> Dict:
        """
        获取模型临时下载链接

        Parameters
        ----------
        checkpoint_id: 模型对应的 id，通过 list_weights 接口获取

        Returns
        -------

        """
        return self.http_client.get_archive_url(
            self.user_info["user_id"], checkpoint_id
        )

    def get_training_run(self, training_run_id: str) -> Dict:
        """
        获取指定训练运行的详细信息

        Parameters
        ----------
        training_run_id: 训练运行 id

        Returns
        -------

        """
        return self.http_client.get_training_run(training_run_id)

    def list_training_runs(self, limit: int = 20, offset: int = 0) -> Dict:
        """
        分页列出当前用户的训练运行

        Parameters
        ----------
        limit: 返回条数，默认为 20
        offset: 偏移量，默认为 0

        Returns
        -------

        """
        return self.http_client.list_training_runs(limit, offset)

    def list_checkpoints(self, training_run_id: str) -> Dict:
        """
        列出指定训练运行下的检查点列表

        Parameters
        ----------
        training_run_id: 训练运行 id

        Returns
        -------

        """
        return self.http_client.list_checkpoints(training_run_id)

    def get_checkpoint_archive_url(
        self, training_run_id: str, checkpoint_id: str
    ) -> Dict:
        """
        获取指定检查点的临时下载链接

        Parameters
        ----------
        training_run_id: 训练运行 id
        checkpoint_id: 检查点 id

        Returns
        -------

        """
        return self.http_client.get_checkpoint_archive_url(
            training_run_id, checkpoint_id
        )

    def delete_checkpoint(self, training_run_id: str, checkpoint_id: str) -> None:
        """
        删除指定训练运行下的检查点

        Parameters
        ----------
        training_run_id: 训练运行 id
        checkpoint_id: 检查点 id

        Returns
        -------

        """
        self.http_client.delete_checkpoint(training_run_id, checkpoint_id)

    def list_user_checkpoints(self, limit: int = 100, offset: int = 0) -> Dict:
        """
        分页列出当前用户的所有检查点

        Parameters
        ----------
        limit: 返回条数，默认为 100
        offset: 偏移量，默认为 0

        Returns
        -------

        """
        return self.http_client.list_user_checkpoints(limit, offset)

    def get_session(self, session_id: str) -> Dict:
        """
        获取指定会话的信息

        Parameters
        ----------
        session_id: 会话 id

        Returns
        -------

        """
        return self.http_client.get_session(session_id)

    def list_sessions(self, limit: int = 20, offset: int = 0) -> Dict:
        """
        分页列出当前用户的所有会话

        Parameters
        ----------
        limit: 返回条数，默认为 20
        offset: 偏移量，默认为 0

        Returns
        -------

        """
        return self.http_client.list_sessions(limit, offset)

    def get_sampler(self, sampler_id: str) -> Dict:
        """
        获取指定采样器的信息

        Parameters
        ----------
        sampler_id: 采样器 id

        Returns
        -------

        """
        return self.http_client.get_sampler(sampler_id)

    def get_checkpoint_info(self, path: str) -> Dict:
        """
        获取指定路径的检查点信息

        Parameters
        ----------
        path: 检查点路径

        Returns
        -------

        """
        try:
            response = self.http_client.http_client.get(
                "/checkpoints/info",
                params={
                    "path": path,
                },
            )
            response.raise_for_status()
        except Exception as e:
            raise ValueError(f"Failed to get checkpoint info for path {path}: {e}")

        return response.json()
