# -*- coding:utf-8 -*-

import os
import time
from http.cookies import SimpleCookie
from threading import Thread

import requests

from utils import COMMON_HEADERS


class SunoCookie:
    def __init__(self):
        self.cookie = SimpleCookie()
        self.session_id = None
        self.token = None
        self.is_disabled = False

    def load_cookie(self, cookie_str):
        self.cookie.load(cookie_str)

    def get_cookie(self):
        return ";".join([f"{i}={self.cookie.get(i).value}" for i in self.cookie.keys()])

    def set_session_id(self, session_id):
        self.session_id = session_id

    def get_session_id(self):
        return self.session_id

    def get_token(self):
        return self.token

    def set_token(self, token: str):
        self.token = token

    def disable(self):
        self.is_disabled = True

    def enable(self):
        self.is_disabled = False

def load_env_cookies():
    suno_auths = {}
    for key, value in os.environ.items():
        if key.startswith("SESSION_ID"):
            try:
                i = int(key[10:])  # 提取 SESSION_ID 后的数字
                cookie = os.getenv(f"COOKIE{i}")
                if cookie:
                    suno_auth = SunoCookie()
                    suno_auth.set_session_id(value)
                    suno_auth.load_cookie(cookie)
                    suno_auths[i] = suno_auth
            except ValueError:
                continue
    return suno_auths

def update_token(suno_cookie: SunoCookie):
    headers = {"cookie": suno_cookie.get_cookie()}
    headers.update(COMMON_HEADERS)
    session_id = suno_cookie.get_session_id()

    resp = requests.post(
        url=f"https://clerk.suno.com/v1/client/sessions/{session_id}/tokens?_clerk_js_version=5.22.3",
        headers=headers,
    )

    resp_headers = dict(resp.headers)
    set_cookie = resp_headers.get("Set-Cookie")
    suno_cookie.load_cookie(set_cookie)
    token = resp.json().get("jwt")
    suno_cookie.set_token(token)
    # print(set_cookie)
    # print(f"*** token -> {token} ***")


def keep_alive(suno_cookie: SunoCookie):
    while True:
        try:
            update_token(suno_cookie)
        except Exception as e:
            print(e)
        finally:
            time.sleep(5)

def start_keep_alive(suno_auths):
    for suno_auth in suno_auths.values():
        t = Thread(target=keep_alive, args=(suno_auth,))
        t.start()

suno_auths = load_env_cookies()
start_keep_alive(suno_auths)