#
# Copyright (C) 2024 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.
#
# Implementation taken from external/perfetto/tools/record_android_trace.
#

import webbrowser
import socketserver
import http.server
import os
import subprocess
import types
from .base import ValidationError
from .handle_input import HandleInput
from .utils import path_exists, wait_for_process_or_ctrl_c, wait_for_output

TORQ_TEMP_DIR = "/tmp/.torq"
TRACE_PROCESSOR_BINARY = "/trace_processor"
TORQ_TEMP_TRACE_PROCESSOR = TORQ_TEMP_DIR + TRACE_PROCESSOR_BINARY
ANDROID_PERFETTO_TOOLS_DIR = "/external/perfetto/tools"
ANDROID_TRACE_PROCESSOR = ANDROID_PERFETTO_TOOLS_DIR + TRACE_PROCESSOR_BINARY
LARGE_FILE_SIZE = 1024 * 1024 * 1024 * 4  # 4 GiB
WAIT_FOR_TRACE_PROCESSOR_MS = 3000
MAX_SERVER_WAIT_TIME_SECS = 15
WEB_UI_ADDRESS = "https://ui.perfetto.dev"


class HttpHandler(http.server.SimpleHTTPRequestHandler):

  def end_headers(self):
    self.send_header("Access-Control-Allow-Origin", self.server.allow_origin)
    self.send_header("Cache-Control", "no-cache")
    super().end_headers()

  def do_GET(self):
    if self.path != "/" + self.server.expected_fname:
      self.send_error(404, "File not found")
      return
    self.server.fname_get_completed = True
    super().do_GET()

  def do_POST(self):
    self.send_error(404, "File not found")

  def log_message(self, format, *args):
    pass


def download_trace_processor(path):
  if (("ANDROID_BUILD_TOP" in os.environ and
       path_exists(os.environ["ANDROID_BUILD_TOP"] + ANDROID_TRACE_PROCESSOR))):
    return os.environ["ANDROID_BUILD_TOP"] + ANDROID_TRACE_PROCESSOR
  if path_exists(TORQ_TEMP_TRACE_PROCESSOR):
    return TORQ_TEMP_TRACE_PROCESSOR

  def download_accepted_callback():
    subprocess.run(
        ("mkdir -p %s && wget -P %s "
         "https://get.perfetto.dev/trace_processor && chmod +x "
         "%s/trace_processor" % (TORQ_TEMP_DIR, TORQ_TEMP_DIR, TORQ_TEMP_DIR)),
        shell=True)

    if not path_exists(TORQ_TEMP_TRACE_PROCESSOR):
      print("Could not download perfetto scripts. Continuing.")
      return None

    return TORQ_TEMP_TRACE_PROCESSOR

  def rejected_callback():
    print("Will continue without downloading perfetto scripts.")
    return None

  return (HandleInput(
      "You do not have $ANDROID_BUILD_TOP configured "
      "with the $ANDROID_BUILD_TOP%s directory.\nYour "
      "perfetto trace is larger than 512MB, so "
      "attempting to load the trace in the perfetto UI "
      "without the perfetto scripts might not work.\n"
      "torq can download the perfetto scripts to '%s'. "
      "Are you ok with this download? [Y/N]: " %
      (ANDROID_PERFETTO_TOOLS_DIR, TORQ_TEMP_DIR),
      "Please accept or reject the download.", {
          "y": download_accepted_callback,
          "n": rejected_callback
      }).handle_input())


def open_trace(path, origin, use_trace_processor):
  PORT = 9001
  path = os.path.abspath(path)
  trace_processor_path = None
  if os.path.getsize(path) >= LARGE_FILE_SIZE or use_trace_processor:
    trace_processor_path = download_trace_processor(path)
  if isinstance(trace_processor_path, ValidationError):
    return trace_processor_path
  if trace_processor_path is not None:
    process = subprocess.Popen(
        "%s --httpd %s" % (trace_processor_path, path),
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT)
    print("\033[93m##### Loading trace. #####")
    if wait_for_output("Trace loaded", process, WAIT_FOR_TRACE_PROCESSOR_MS):
      process.kill()
      return ValidationError("Trace took too long to load.",
                             "Please try again.")
    webbrowser.open_new_tab(origin)
    print("##### Follow the directions in the Perfetto UI. Do not "
          "exit out of torq until you are done viewing the trace. Press "
          "CTRL+C to exit torq and close the trace_processor. #####\033[0m")
    wait_for_process_or_ctrl_c(process)
  else:  # Open trace directly in UI

    def handle_timeout(self):
      self.timed_out = True
      print("File server timed out while waiting for a request.")

    os.chdir(os.path.dirname(path))
    fname = os.path.basename(path)
    socketserver.TCPServer.allow_reuse_address = True
    with (socketserver.TCPServer(("127.0.0.1", PORT), HttpHandler) as httpd):
      address = (f"{origin}/#!/?url=http://127.0.0.1:"
                 f"{PORT}/{fname}&referrer=open_trace_in_ui")
      webbrowser.open_new_tab(address)
      httpd.expected_fname = fname
      httpd.handle_timeout = types.MethodType(handle_timeout, httpd)
      httpd.fname_get_completed = None
      httpd.allow_origin = origin
      httpd.timeout = MAX_SERVER_WAIT_TIME_SECS
      httpd.timed_out = False
      while not httpd.timed_out and httpd.fname_get_completed is None:
        httpd.handle_request()
  return None