from flask import Flask, request, jsonify
import boto3
from botocore.exceptions import ClientError
import requests
import os
import logging
from auth0.authentication import GetToken

app = Flask(__name__)

app.logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
formatter = logging.Formatter(
    '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
app.logger.addHandler(handler)

# fixme: change to config
GITHUB_ORG = "tidbcloud"
AWS_ROLE_ARN = "arn:aws:iam::557537366020:role/autonomous-o11y-s3-readonly"


def check_github_org(token):
    headers = {
        "Authorization": f"token {token}",
        "Accept": "application/vnd.github.v3+json"
    }
    response = requests.get(
        "https://api.github.com/user/orgs", headers=headers)
    return response.ok and any(org['login'] == GITHUB_ORG for org in response.json())


def get_aws_credentials():
    sts = boto3.client('sts')
    try:
        role = sts.assume_role(
            RoleArn=AWS_ROLE_ARN,
            RoleSessionName="ClientAccessSession",
            DurationSeconds=3600
        )
        return role['Credentials']
    except ClientError as e:
        app.logger.error(f"Get AWS credentials failed: {e}")
        return None


def get_auth0_token():
    domain = 'tidb-soc2.us.auth0.com'
    get_token = GetToken(domain, os.getenv('AUTH0_CLIENT_ID'),
                         client_secret=os.getenv('AUTH0_CLIENT_SECRET'))
    token = get_token.client_credentials(
        'https://{}/api/v2/'.format(domain))
    auth0_token = token['access_token']
    return auth0_token
    # url = 'https://tidb-soc2.us.auth0.com/oauth/token'
    # headers = {
    #     'content-type': 'application/x-www-form-urlencoded'
    # }
    #
    # client_id = os.getenv('AUTH0_CLIENT_ID')
    # client_secret = os.getenv('AUTH0_CLIENT_SECRET')
    #
    # app.logger.info(f"Attempting to get Auth0 token. Client ID: {client_id[:4]}...{client_id[-4:]} (truncated), Client Secret: {'*' * len(client_secret) if client_secret else 'None'}") # 避免打印完整secret
    #
    # if not client_id or not client_secret:
    #     app.logger.error("Error: AUTH0_CLIENT_ID or AUTH0_CLIENT_SECRET environment variable not set.")
    #     return None
    #
    # data = {
    #     'grant_type': 'client_credentials',
    #     'client_id': client_id,
    #     'client_secret': client_secret,
    #     'audience': 'https://tidb-soc2.us.auth0.com/api/v2/'
    # }
    #
    # try:
    #     response = requests.post(url, headers=headers, data=data)
    #     app.logger.info(f"Auth0 token request response object: {response}") # 打印 response 对象
    #
    #     response.raise_for_status()
    #
    #     response_json = response.json()
    #     access_token = response_json.get('access_token')
    #
    #     if access_token:
    #         app.logger.info("Auth0 token retrieved successfully")
    #         return access_token
    #     else:
    #         app.logger.warning(f"Auth0 token not found in the HTTP response. Full response: {response.text}")
    #         return None
    #
    # except requests.exceptions.HTTPError as errh:
    #     app.logger.error(f"HTTP Error during Auth0 token retrieval: {errh}. Response content: {response.text}")
    #     return None
    # except requests.exceptions.ConnectionError as errc:
    #     app.logger.error(f"Connection Error during Auth0 token retrieval: {errc}")
    #     return None
    # except requests.exceptions.Timeout as errt:
    #     app.logger.error(f"Timeout Error during Auth0 token retrieval: {errt}")
    #     return None
    # except requests.exceptions.RequestException as err:
    #     app.logger.error(f"An unknown error occurred during Auth0 token request: {err}")
    #     return None
    # except json.JSONDecodeError:
    #     app.logger.error(f"Could not parse JSON response during Auth0 token retrieval. Raw response content: {response.text}")
    #     return None


@app.route('/auth', methods=['POST'])
def auth():
    data = request.get_json()
    if not data or 'github_token' not in data:
        app.logger.warning("Auth request: 'github_token' missing.")
        return jsonify({"error": "need GitHub token"}), 400

    if not check_github_org(data['github_token']):
        app.logger.warning("Auth request: GitHub org validation failed.")
        return jsonify({"error": "Failed to validate github org"}), 403

    app.logger.info("Attempting to get Auth0 token from function.")
    auth0_token = get_auth0_token()

    if auth0_token is None:
        app.logger.error("Failed to get auth0 token in /auth endpoint.")
        return jsonify({"error": "Failed to get auth0 token"}), 500

    creds = get_aws_credentials()
    if creds:
        app.logger.info("AWS credentials obtained successfully.")
        return jsonify({
            "access_key_id": creds['AccessKeyId'],
            "secret_access_key": creds['SecretAccessKey'],
            "session_token": creds['SessionToken'],
            "expiration": creds['Expiration'].isoformat(),
            "auth0_token": auth0_token,
        })
    app.logger.error("Failed to init AWS credentials in /auth endpoint.")
    return jsonify({"error": "Failed to init AWS credentials"}), 500


if __name__ == "__main__":
    app.run(host='0.0.0.0', port=8080)