import json
import logging
import time
from os import PathLike
from pathlib import Path
from typing import Literal, Optional

import requests

from ui import ConstData
from .orm import (
    ResourceCompleteRspSchema,
    ResourceCreateRspSchema,
    ResultRspSchema,
    ResultStateEnum,
    TaskCreateRspSchema,
)

__version__ = "0.0.3"

API_BASE_URL = "https://member.bilibili.com/x/bcut/rubick-interface"

# 申请上传
API_REQ_UPLOAD = API_BASE_URL + "/resource/create"

# 提交上传
API_COMMIT_UPLOAD = API_BASE_URL + "/resource/create/complete"

# 创建任务
API_CREATE_TASK = API_BASE_URL + "/task"

# 查询结果
API_QUERY_RESULT = API_BASE_URL + "/task/result"

SUPPORT_SOUND_FORMAT = Literal["flac", "aac", "m4a", "mp3", "wav"]

INFILE_FMT = ["flac", "aac", "m4a", "mp3", "wav"]
OUTFILE_FMT = ["srt", "json", "lrc", "txt"]
"""必剪 语音识别接口"""

headers = {
    "User-Agent": "Bilibili/1.0.0 (https://www.bilibili.com)",
    "Content-Type": "application/json",
}


class APIError(Exception):
    "接口调用错误"

    def __init__(self, code, msg) -> None:
        self.code = code
        self.msg = msg
        super().__init__()

    def __str__(self) -> str:
        return f"{self.code}:{self.msg}"


class BcutASR:
    "必剪 语音识别接口"
    session: requests.Session
    sound_name: str
    sound_bin: bytes
    sound_fmt: SUPPORT_SOUND_FORMAT
    __in_boss_key: str
    __resource_id: str
    __upload_id: str
    __upload_urls: list[str]
    __per_size: int
    __clips: int
    __etags: list[str]
    __download_url: str
    task_id: str

    def __init__(self, file: Optional[str] = None) -> None:
        self.session = requests.Session()
        self.task_id = None
        self.__etags = []
        if file:
            self.set_data(file)

    def set_data(
            self,
            file: Optional[str] = None,
            raw_data: Optional[bytes] = None,
            data_fmt: Optional[SUPPORT_SOUND_FORMAT] = None,
    ) -> None:
        "设置欲识别的数据"
        if file:
            if not isinstance(file, (str, PathLike)):
                raise TypeError("unknow file ptr")
            # 文件类
            file = Path(file)
            self.sound_bin = open(file, "rb").read()
            suffix = data_fmt or file.suffix[1:]
            self.sound_name = file.name
        elif raw_data:
            # bytes类
            self.sound_bin = raw_data
            suffix = data_fmt
            self.sound_name = f"{int(time.time())}.{suffix}"
        else:
            raise ValueError("none set data")
        if suffix not in SUPPORT_SOUND_FORMAT.__args__:
            raise TypeError("format is not support")
        self.sound_fmt = suffix
        logging.info(f"加载文件成功: {self.sound_name}___{self.sound_fmt}")

    def upload(self) -> None:
        "申请上传"
        if not self.sound_bin or not self.sound_fmt:
            raise ValueError("none set data")
        resp = self.session.post(
            API_REQ_UPLOAD,
            data=json.dumps({
                "type": 2,
                "name": self.sound_name,
                "size": len(self.sound_bin),
                "ResourceFileType": self.sound_fmt,
                "model_id": "8",
            }),
            headers=headers
        )
        resp.raise_for_status()
        resp = resp.json()
        code = resp["code"]
        if code:
            raise APIError(code, resp["message"])
        resp_data = ResourceCreateRspSchema.parse_obj(resp["data"])
        self.__in_boss_key = resp_data.in_boss_key
        self.__resource_id = resp_data.resource_id
        self.__upload_id = resp_data.upload_id
        self.__upload_urls = resp_data.upload_urls
        self.__per_size = resp_data.per_size
        self.__clips = len(resp_data.upload_urls)
        logging.info(
            f"申请上传成功, 总计大小{resp_data.size // 1024}KB, {self.__clips}分片, 分片大小{resp_data.per_size // 1024}KB: {self.__in_boss_key}"
        )
        if ConstData.cut_stop:
            return None
        self.__upload_part()
        if ConstData.cut_stop:
            return None
        self.__commit_upload()

    def __upload_part(self) -> None:
        "上传音频数据"
        for clip in range(self.__clips):
            if ConstData.cut_stop:
                return None
            start_range = clip * self.__per_size
            end_range = (clip + 1) * self.__per_size
            logging.info(f"开始上传分片{clip}: {start_range}-{end_range}")
            for attempt in range(3):
                try:
                    resp = self.session.put(
                        self.__upload_urls[clip],
                        data=self.sound_bin[start_range:end_range],
                    )
                    resp.raise_for_status()
                    etag = resp.headers.get("Etag")
                    self.__etags.append(etag)
                    logging.info(f"分片{clip}上传成功: {etag}")
                    break
                except requests.exceptions.RequestException as e:
                    logging.warning(f"上传分片 {clip} 失败，第 {attempt + 1} 次重试中... 错误: {e}")
                    time.sleep(2)

    def __commit_upload(self) -> None:
        "提交上传数据"
        resp = self.session.post(
            API_COMMIT_UPLOAD,
            data=json.dumps(
                {
                    "InBossKey": self.__in_boss_key,
                    "ResourceId": self.__resource_id,
                    "Etags": ",".join(self.__etags) if self.__etags else "",
                    "UploadId": self.__upload_id,
                    "model_id": "8",
                }
            ),
            headers=headers
        )
        resp.raise_for_status()
        resp = resp.json()
        code = resp["code"]
        if code:
            raise APIError(code, resp["message"])
        resp_data = ResourceCompleteRspSchema.model_validate(resp["data"])
        self.__download_url = resp_data.download_url
        logging.info(f"提交成功")

    def create_task(self) -> str:
        "开始创建转换任务"
        resp = self.session.post(
            API_CREATE_TASK, json={"resource": self.__download_url, "model_id": "8"}, headers=headers
        )
        resp.raise_for_status()
        resp = resp.json()
        code = resp["code"]
        if code:
            raise APIError(code, resp["message"])
        resp_data = TaskCreateRspSchema.model_validate(resp["data"])
        self.task_id = resp_data.task_id
        logging.info(f"任务已创建: {self.task_id}")
        return self.task_id

    def result(self, task_id: Optional[str] = None) -> ResultRspSchema:
        "查询转换结果"
        resp = self.session.get(
            API_QUERY_RESULT, params={"model_id": 8, "task_id": task_id or self.task_id},
            headers=headers,
        )
        resp.raise_for_status()
        resp = resp.json()
        code = resp["code"]
        if code:
            raise APIError(code, resp["message"])
        return ResultRspSchema.model_validate(resp["data"])
