import sys
import requests
import time
import zipfile
from io import BytesIO
from github.github_apis import create_git_headers

git_url = "https://api.github.com/repos"
headers = create_git_headers()

def wait_for_workflow_completion(repo_name, run_id, poll_interval=15):
    while True:
        response = requests.get(f'{git_url}/{repo_name}/actions/runs/{run_id}', headers=headers)
        if response.status_code == 200:
            data = response.json()
            status = data['status']
            conclusion = data.get('conclusion')
            if status == "completed":
                print(f"Workflow completed with status: {status}, and conclusion: {conclusion}")
                return conclusion == "success"
        else:
            print(f"Error polling workflow status: {response.status_code}")
            print(response.text)
            return False
        time.sleep(poll_interval)

def download_and_extract_logs(repo_name, run_id):
    response = requests.get(f'{git_url}/{repo_name}/actions/runs/{run_id}/logs', headers=headers)
    if response.status_code == 200:
        with zipfile.ZipFile(BytesIO(response.content)) as zip_file:
            zip_file.extractall("workflow_logs")
            print("Extracted logs to 'workflow_logs'.")
    else:
        print(f"Error downloading logs: {response.status_code}")
        print(response.text)

if __name__ == "__main__":
    repo_name = sys.argv[1].replace("repo_name=", "", 1)
    run_id = sys.argv[2].replace("run_id=", "", 1)

    print(f"Polling workflow run ID: {run_id} from repository: {repo_name}")
    success = wait_for_workflow_completion(repo_name, run_id)
    if not success:
        print("Triggered workflow status is failed.")
        exit(1)
    print("Triggered workflow status is success.")

    # Download and save workflow run logs to workflow_logs
    download_and_extract_logs(repo_name, run_id)