from typing import Annotated, Union
from fastapi import FastAPI, Header, status
from fastapi.responses import Response
from pydantic import BaseModel
from enum import Enum
import os
import time
import jwt
import json
import base64
import click
import ffmpeg
import shutil
import logging
import uvicorn
import uuid
from config import settings
os.environ["TF_CPP_MIN_LOG_LEVEL"] = settings["TF_CPP_MIN_LOG_LEVEL"]
from spleeter.separator import Separator


class Payload(BaseModel):
    filename: str
    data: str  # base64 encoded
    actions: list[str]


class TokenResponse(Enum):
    TOKEN_OK = 0
    TOKEN_SCHEME_ERROR = 1
    TOKEN_EXPIRED = 2
    TOKEN_INVALID = 3
    TOKEN_DECODE_ERROR = 4
    TOKEN_USED = 5
    
class ProcessStatus():
    uuid = ""
    name = ""
    status = 0
    total = 0
    tempfiles = []

    def __init__(self, uuid, name, total):
        self.uuid = uuid
        self.name = name
        self.status = 0
        self.total = total
    
    def update(self):
        self.status += 1
        
    def update_temp_files(self, tempfile):
        self.tempfiles.append(tempfile)

    def is_done(self):
        return self.status == self.total

    def get_status(self):
        return self.status, self.total
    
    def __eq__(self, other):
        return self.uuid == other.uuid or self.name == other.name
    
    def __dict__(self):
        return {"uuid": self.uuid, "name": self.name, "status": self.status, "total": self.total, "tempfiles": self.tempfiles}
    

queue_status = []

tempfiles = []

app = FastAPI(docs_url=None, redoc_url=None)

headers = {
    "typ": "JWT",
    "alg": "HS256"
}

usedJTI = set()  # Avoid replay attacks
separator = Separator("spleeter:2stems")
separator_lock = False
LOGFORMAT = "[%(asctime)s] [%(levelname)s] <%(threadName)s> %(message)s"
logging.basicConfig(format=LOGFORMAT, level=settings.LOG_LEVEL)


def verify_token(auth):
    scheme = auth.split()[0]
    if scheme != "Bearer":
        return TokenResponse.TOKEN_SCHEME_ERROR
    token = auth.split()[1]
    try:
        result = jwt.decode(token,
                            settings["JWT_SECRET"],
                            verify=True,
                            algorithms=["HS256"],
                            audience="galaxymusichall.convert")
        if result["jti"] in usedJTI:
            return TokenResponse.TOKEN_USED
        usedJTI.add(result["jti"])
        return TokenResponse.TOKEN_OK
    except jwt.ExpiredSignatureError:
        return TokenResponse.TOKEN_EXPIRED
    except jwt.InvalidTokenError:
        return TokenResponse.TOKEN_INVALID
    except jwt.DecodeError:
        return TokenResponse.TOKEN_DECODE_ERROR
    

@app.post("/clear")
def clear(authorization: Annotated[Union[str, None], Header()],
          response: Response):
    if verify := verify_token(authorization):
        if verify == TokenResponse.TOKEN_OK:
            pass
        elif verify == TokenResponse.TOKEN_SCHEME_ERROR:
            response.status_code = status.HTTP_400_BAD_REQUEST
            return {"error": "Invalid authorization header, 'bearer' required"}
        elif verify == TokenResponse.TOKEN_EXPIRED:
            response.status_code = status.HTTP_403_FORBIDDEN
            return {"error": "Token has expired"}
        elif verify == TokenResponse.TOKEN_INVALID:
            response.status_code = status.HTTP_403_FORBIDDEN
            return {"error": "Invalid token"}
        elif verify == TokenResponse.TOKEN_DECODE_ERROR:
            response.status_code = status.HTTP_400_BAD_REQUEST
            return {"error": "Token decoding failed"}
        elif verify == TokenResponse.TOKEN_USED:
            response.status_code = status.HTTP_403_FORBIDDEN
            return {"error": "Token has been used"}
        else:
            response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
            return {"error": "Uncaught error"}
    for i in tempfiles:
        try:
            os.remove(i)
            logging.info("删除文件: "+i)
        except FileNotFoundError:
            logging.warning("文件已被删除: "+i)
            pass
        except Exception as e:
            logging.error("删除文件失败: "+str(e))
            response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
            return {"error": "Failed to clear temporary files"}
    tempfiles.clear()
    return {"message": "Temporary files cleared"}


@app.get("/status")
def status(authorization: Annotated[Union[str, None], Header()]):
    if verify := verify_token(authorization):
        if verify == TokenResponse.TOKEN_OK:
            pass
        elif verify == TokenResponse.TOKEN_SCHEME_ERROR:
            response.status_code = status.HTTP_400_BAD_REQUEST
            return {"error": "Invalid authorization header, 'bearer' required"}
        elif verify == TokenResponse.TOKEN_EXPIRED:
            response.status_code = status.HTTP_403_FORBIDDEN
            return {"error": "Token has expired"}
        elif verify == TokenResponse.TOKEN_INVALID:
            response.status_code = status.HTTP_403_FORBIDDEN
            return {"error": "Invalid token"}
        elif verify == TokenResponse.TOKEN_DECODE_ERROR:
            response.status_code = status.HTTP_400_BAD_REQUEST
            return {"error": "Token decoding failed"}
        elif verify == TokenResponse.TOKEN_USED:
            response.status_code = status.HTTP_403_FORBIDDEN
            return {"error": "Token has been used"}
        else:
            response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
            return {"error": "Uncaught error"}
    return json.dumps([i.__dict__() for i in queue_status])

@app.post("/process")
def process(data: Payload,
            authorization: Annotated[Union[str, None], Header()],
            response: Response):
    global separator_lock
    #print(data.filename, data.actions, authorization)
    logging.info(f"处理文件{data.filename}，操作{data.actions}")
    processUUID = str(uuid.uuid4())
    process_status = ProcessStatus(uuid=processUUID, name=data.filename, total=len(data.actions))
    queue_status.append(process_status)
    if verify := verify_token(authorization):
        if verify == TokenResponse.TOKEN_OK:
            pass
        elif verify == TokenResponse.TOKEN_SCHEME_ERROR:
            response.status_code = status.HTTP_400_BAD_REQUEST
            return {"error": "Invalid authorization header, 'bearer' required"}
        elif verify == TokenResponse.TOKEN_EXPIRED:
            response.status_code = status.HTTP_403_FORBIDDEN
            return {"error": "Token has expired"}
        elif verify == TokenResponse.TOKEN_INVALID:
            response.status_code = status.HTTP_403_FORBIDDEN
            return {"error": "Invalid token"}
        elif verify == TokenResponse.TOKEN_DECODE_ERROR:
            response.status_code = status.HTTP_400_BAD_REQUEST
            return {"error": "Token decoding failed"}
        elif verify == TokenResponse.TOKEN_USED:
            response.status_code = status.HTTP_403_FORBIDDEN
            return {"error": "Token has been used"}
        else:
            response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
            return {"error": "Uncaught error"}
    tempfilename = "/tmp/"+data.filename
    try:
        with open(tempfilename, "wb") as wrfile:
            wrfile.write(base64.b64decode(data.data))
    except FileExistsError:
        response.status_code = status.HTTP_409_CONFLICT
        return {"error": "File already exists"}
    except Exception as e:
        response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
        logging.error("Unexpected exception:",e)
        return {"error": "Uncaught error"}
    current_name = tempfilename
    process_status.update_temp_files(tempfilename)
    basename = tempfilename[:tempfilename.rfind(".")]
    try:
        for i in data.actions:
            logging.info(f"正在处理{process_status.get_status()}/{process_status.total}：{i}")
            if i == "separate":
                # TODO:性能优化，并发
                while separator_lock:
                    logging.info("等待上一Spleeter任务完成...")
                    time.sleep(2)
                separator_lock = True
                current_name = separate(current_name, basename)
                process_status.update_temp_files(basename+"/vocals.wav")
                process_status.update_temp_files(current_name)
                separator_lock = False
            elif i == "audio":
                current_name = audio(current_name, basename, target="aac")
                process_status.update_temp_files(current_name)
            elif "convert" in i:
                if "->" not in i:
                    response.status_code = status.HTTP_400_BAD_REQUEST
                    return {"error": "操作格式错误", "hint": "使用'->'指定目标格式"}
                current_name = convert(current_name, basename, target=i.split("->")[1])
                process_status.update_temp_files(current_name)
            else:
                logging.warning("无效操作"+i+"，跳过中...")
                pass
            process_status.update()
            if current_name is None:
                response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
                return {"error": "处理失败", "status": process_status.get_status(), "required": process_status.total}
        data = {"filename": current_name[current_name.rfind('/') + 1:],
                "data": base64.b64encode(open(current_name, "rb").read()).decode()}
        if process_status.is_done():
            tempfiles.extend(process_status.tempfiles)
            queue_status.remove(process_status)
        else:
            response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
            return {"error": "未预期的错误：流水线在完成前就退出", "status": process_status.get_status(), "required": process_status.total}
        return json.dumps(data)
    except Exception as e:
        logging.error("未预期的错误：" + str(e))
        response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
        return {"error": "Processing failed"}


def separate(filename, basename):
    separator.separate_to_file(filename, "/tmp")
    shutil.move(basename + "/accompaniment.wav", basename + ".wav")
    return basename + ".wav"


def audio(filename, basename, target="aac"):
    audio = ffmpeg.input(filename).audio
    out = ffmpeg.output(audio, basename+"."+target)
    out.overwrite_output().run()
    return basename+"."+target


def convert(filename, basename, target="wav"):
    audio = ffmpeg.input(filename)
    out = ffmpeg.output(audio, basename+"."+target)
    out.overwrite_output().run()
    return basename+"."+target


@click.group()
def cli():
    pass


@cli.command()
def secret():
    """显示JWT密钥"""
    click.echo(settings["JWT_SECRET"])


@cli.command()
@click.argument("token")
def verify(token):
    """验证JWT token"""
    try:
        result = jwt.decode(token,
                            settings["JWT_SECRET"],
                            verify=True,
                            algorithms=["HS256"],
                            audience="galaxymusichall.convert")
        click.echo(result)
    except jwt.DecodeError as e:
        click.echo("JWT token解码失败"+str(e))
    except jwt.ExpiredSignatureError:
        click.echo("token已过期。")
    except jwt.InvalidTokenError:
        click.echo("token无效。")
    except Exception as e:
        click.echo("未知错误：")
        click.echo(e)


@cli.command()
def run():
    """启动转换服务器"""
    uvicorn.run(app, host=settings.HOST, port=settings.PORT)


if __name__ == "__main__":
    cli()