# Copyright (C) 2025 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import flask
import pathlib
import json
import requests
import io
import zipfile

from utils import ArgParseImpl

from flask import Flask, request, make_response, send_from_directory

DIR = pathlib.Path(__file__).parent.absolute()
HTML_DIR = os.path.join(DIR, "viewer_html")

def _download_and_extract_artifacts(run_id, headers, output_dir):
  OWNER_REPO = 'google/filament'
  artifacts_url = f"https://api.github.com/repos/{OWNER_REPO}/actions/runs/{run_id}/artifacts"
  downloaded_any_artifact = False
  try:
    response = requests.get(artifacts_url, headers=headers)
    response.raise_for_status()
    artifacts_data = response.json()
    artifacts = artifacts_data.get("artifacts", [])

    if not artifacts:
      print(f"  No artifacts found for workflow run ID {run_id}.")
      return downloaded_any_artifact

    for artifact in artifacts:
      artifact_id = artifact["id"]
      artifact_name = artifact["name"]
      archive_download_url = artifact["archive_download_url"]

      print(f"  Found artifact: '{artifact_name}' (ID: {artifact_id})")

      # Perform the download request
      print(f"  Downloading '{artifact_name}'...")
      # Use a copy of headers and specific Accept for ZIP download
      download_headers = headers.copy()
      download_headers["Accept"] = "application/vnd.github.v3+zip"
      download_response = requests.get(archive_download_url, headers=download_headers, stream=True)
      download_response.raise_for_status() # Check for errors in download

      # --- Step 5: Extract the contents ---
      # Use BytesIO to handle the zip file content in memory without saving to a temporary file
      with io.BytesIO(download_response.content) as zip_buffer:
        try:
          with zipfile.ZipFile(zip_buffer, 'r') as zip_ref:
            # Create a unique subdirectory for each artifact to avoid file name conflicts
            extract_path = os.path.join(output_dir, f"{artifact_name}_{artifact_id}")
            os.makedirs(extract_path, exist_ok=True)
            zip_ref.extractall(extract_path)
            print(f"  Successfully extracted '{artifact_name}' to '{extract_path}/'")
            downloaded_any_artifact = True
        except zipfile.BadZipFile:
          print(f"  Error: Downloaded file for '{artifact_name}' is not a valid zip file. Skipping extraction.")
        except Exception as e:
          print(f"  An error occurred during extraction of '{artifact_name}': {e}")
    return downloaded_any_artifact
  except requests.exceptions.HTTPError as e:
    print(f"  An HTTP error occurred while fetching artifacts for run {run_id}: {e}")
    if e.response.status_code == 403:
      print("  This often means you need a GitHub Personal Access Token with 'repo' scope (even for public repos for artifact downloads).")
  except requests.exceptions.RequestException as e:
    print(f"  A network error occurred while fetching artifacts for run {run_id}: {e}")

def _download_github_artifacts_by_run(run_id, github_token, output_dir= ".") -> None:
  headers = {"Accept": "application/vnd.github.v3+json"}
  if github_token:
    headers["Authorization"] = f"token {github_token}"
  os.makedirs(output_dir, exist_ok=True)
  return _download_and_extract_artifacts(run_id, headers, output_dir)

# Generated by gemini
def _download_github_artifacts(pr_number, github_token, output_dir= ".") -> None:
  # Prepare HTTP headers for GitHub API requests
  headers = {"Accept": "application/vnd.github.v3+json"}
  if github_token:
    headers["Authorization"] = f"token {github_token}"

  OWNER_REPO = 'google/filament'

  # --- Step 1: Get PR details to find the head commit SHA ---
  print(f"Fetching details for PR #{pr_number} in {OWNER_REPO}...")
  pr_url = f"https://api.github.com/repos/{OWNER_REPO}/pulls/{pr_number}"
  try:
    response = requests.get(pr_url, headers=headers)
    response.raise_for_status() # Raise an exception for HTTP errors (4xx or 5xx)
    pr_data = response.json()
    commit_sha = pr_data["head"]["sha"]
    print(f"PR #{pr_number} is associated with commit SHA: {commit_sha}")
  except requests.exceptions.HTTPError as e:
    if e.response.status_code == 404:
      print(f"Error: PR #{pr_number} not found in {OWNER_REPO}. Please check the PR number, owner, and repository name.")
    elif e.response.status_code == 403:
      print(f"Error: Access forbidden to PR #{pr_number}. You might be hitting API rate limits or need a valid GitHub Token.")
    else:
      print(f"An HTTP error occurred while fetching PR details: {e}")
    return # Exit function on error
  except requests.exceptions.RequestException as e:
    print(f"A network error occurred while fetching PR details: {e}")
    return # Exit function on error

  # --- Step 2: Find workflow runs associated with the commit SHA ---
  print(f"Searching for workflow runs for commit SHA: {commit_sha}...")
  workflow_runs_url = f"https://api.github.com/repos/{OWNER_REPO}/actions/runs"
  # Filter by head_sha and event='pull_request' for precision
  params = {"head_sha": commit_sha, "event": "pull_request"}
  try:
    response = requests.get(workflow_runs_url, headers=headers, params=params)
    response.raise_for_status()
    runs_data = response.json()
    workflow_runs = runs_data.get("workflow_runs", [])

    if not workflow_runs:
      print(f"No workflow runs found directly associated with PR #{pr_number} (commit SHA: {commit_sha}).")
      print("This might happen if the workflow was triggered by a push after the PR was opened,")
      print("or if the PR head branch was updated without triggering a new workflow run with this exact SHA.")
      print("Consider checking GitHub Actions runs manually for this PR's branch on GitHub.")
      return None

    # Do not filter for runs that completed successfully
    successful_runs = sorted(workflow_runs, key=lambda run: run['id'], reverse=True)
    if not successful_runs:
      print(f"No *successful and completed* workflow runs found for PR #{pr_number} with commit SHA {commit_sha}. Exiting.")
      return None

  except requests.exceptions.HTTPError as e:
    print(f"An HTTP error occurred while searching for workflow runs: {e}")
    return None
  except requests.exceptions.RequestException as e:
    print(f"A network error occurred while searching for workflow runs: {e}")
    return None

  # Create the main output directory if it doesn't exist
  os.makedirs(output_dir, exist_ok=True)
  print(f"Ensuring output directory exists: {os.path.abspath(output_dir)}")

  downloaded_any_artifact = False # Flag to track if any artifact was downloaded

  # --- Step 3 & 4: List and Download Artifacts for each successful run ---
  for run in successful_runs:
    run_id = run["id"]
    run_name = run["name"]
    print(f"\nProcessing workflow run '{run_name}' (ID: {run_id})...")
    downloaded_any_artifact = _download_and_extract_artifacts(run_id, headers, output_dir)
    if downloaded_any_artifact:
        break

  if not downloaded_any_artifact:
    print("\nNo artifacts were downloaded for the specified PR.")
  else:
    print("\nAll available artifacts have been processed.")
  return 'Done'

def _create_app(config):
  app = Flask(__name__)

  client_config = config.copy()
  diff_dir = client_config['diff_dir']

  base_dir = os.path.join(diff_dir, client_config['base_dir'])
  comparison_dir = os.path.join(diff_dir, client_config['comparison_dir'])

  del client_config['base_dir']
  del client_config['comparison_dir']
  del client_config['diff_dir']

  @app.route('/r/', methods=['GET'])
  def get_r():
    return json.dumps(client_config)

  @app.route('/g/<path:filepath>', methods=['GET'])
  def get_g(filepath):
    return send_from_directory(base_dir, filepath)

  @app.route('/d/<path:filepath>', methods=['GET'])
  def get_d(filepath):
    return send_from_directory(diff_dir, filepath)

  @app.route('/c/<path:filepath>', methods=['GET'])
  def get_c(filepath):
    return send_from_directory(comparison_dir, filepath)

  @app.route('/<path:filepath>')
  def get_static_file(filepath):
    return send_from_directory(HTML_DIR, filepath)

  @app.route('/')
  def get_index():
    return send_from_directory(HTML_DIR, 'index.html')

  app.url_map.strict_slashes = False
  return app

if __name__ == '__main__':
  PORT = 8901
  parser = ArgParseImpl()
  parser.add_argument('--diff', type=str, help='Diff result directory')
  parser.add_argument('--pr_number', type=str, help='Pull request artifacts to examine')
  parser.add_argument('--run_number', type=str, help='Run number to examine')
  parser.add_argument('--github_token', type=str, help='Necessary for pull PR artifacts')
  args, _ = parser.parse_known_args(sys.argv[1:])

  if not args.diff and not args.pr_number and not args.run_number:
    print('Need to specify either a diff result directory, a Pull Request number, or a run number')
    exit(1)

  if args.diff and (args.pr_number or args.run_number):
    print('Cannot specify both a diff result directory and a Pull Request/run number')
    exit(1)

  if args.pr_number and args.run_number:
    print('Cannot specify both a PR number and a run number')
    exit(1)

  fdir = args.diff
  if args.pr_number:
    if not args.github_token:
      print('Must provide --github_token to be able to download artifacts')
      exit(1)
    output_dir = f'/tmp/filament-pr{args.pr_number}-rdiff-result'
    res = _download_github_artifacts(args.pr_number, args.github_token, output_dir)
    if not res:
      print('Failed to retrieve PR artifacts')
      exit(1)

    # TODO: Clean up the following so that we're not so specific on the paths diffs/presubmit
    directory_name = list(os.listdir(output_dir))[0]
    fdir = os.path.join(os.path.join(output_dir, directory_name), 'diffs/presubmit')

  if args.run_number:
    if not args.github_token:
      print('Must provide --github_token to be able to download artifacts')
      exit(1)
    output_dir = f'/tmp/filament-run{args.run_number}-rdiff-result'
    res = _download_github_artifacts_by_run(args.run_number, args.github_token, output_dir)
    if not res:
      print('Failed to retrieve run artifacts')
      exit(1)

    # TODO: Clean up the following so that we're not so specific on the paths diffs/presubmit
    directory_name = list(os.listdir(output_dir))[0]
    fdir = os.path.join(os.path.join(output_dir, directory_name), 'diffs/presubmit')

  with open(os.path.join(fdir, 'compare_results.json'), 'r') as f:
    config = json.loads(f.read())
  config['diff_dir'] = os.path.abspath(fdir)

  app = _create_app(config)
  from waitress import serve

  print(f'Point your browser to http://localhost:{PORT} to see the diff results')
  serve(app, host="127.0.0.1", port=PORT)