import os
import json
import boto3
from pathlib import Path
import requests
from cryptography.fernet import Fernet

# path
CLIENT_DIR = Path.home() / ".observability"
CONFIG_FILE = CLIENT_DIR / ".github"
CREDENTIALS_FILE = CLIENT_DIR / ".credentials.enc"
KEY_FILE = CLIENT_DIR / ".key.key"


def ensure_client_dir():
    CLIENT_DIR.mkdir(exist_ok=True, mode=0o700)


def get_github_token():
    token = os.environ.get("GITHUB_TOKEN")
    if token:
        return token
    if CONFIG_FILE.exists():
        with open(CONFIG_FILE, "r") as f:
            return f.read().strip()
    else:
        token = input("Please input GitHub token: ")
        with open(CONFIG_FILE, "w") as f:
            f.write(token)
        CONFIG_FILE.chmod(0o600)
        return token


def generate_key():
    if not KEY_FILE.exists():
        key = Fernet.generate_key()
        with open(KEY_FILE, "wb") as f:
            f.write(key)
        KEY_FILE.chmod(0o600)
    with open(KEY_FILE, "rb") as f:
        return f.read()


def encrypt_credentials(key, credentials):
    return Fernet(key).encrypt(json.dumps(credentials).encode())


def decrypt_credentials(key, encrypted):
    return json.loads(Fernet(key).decrypt(encrypted).decode())


def authenticate():
    response = requests.post(
        "https://www.ds.us-west-2.aws.observability.tidbcloud.com/client/auth",
        # "http://127.0.0.1:8080/auth",
        json={"github_token": get_github_token()}
    )
    if response.status_code == 200:
        return response.json()
    raise Exception(f"failed to authenticate: {response.text}")


def get_credentials():
    ensure_client_dir()

    aws_creds = authenticate()

    return aws_creds


if __name__ == "__main__":
    # os.remove(CONFIG_FILE)
    # os.remove(CREDENTIALS_FILE)
    # os.remove(KEY_FILE)
    creds = get_credentials()

    s3 = boto3.client(
        's3',
        aws_access_key_id=creds['access_key_id'],
        aws_secret_access_key=creds['secret_access_key'],
        aws_session_token=creds.get('session_token')
    )
    s3.download_file('o11y-prod-shared-us-west-2',
                     'diagnosis/data/10000108228606249578/slowlog-stats/1h/2025010708.db', '2025010708.db')