import requests
import argparse
import os
import sys
import time
import numpy as np

# Define the server URL
SERVER_URL = "http://192.168.9.51:8888"


def save_results(results: dict, output_dir: str):
    """Saves the numpy arrays received from the server."""
    if not results or "segment" not in results or "pred" not in results:
        print("Error: The server response did not contain the expected 'segment' or 'pred' arrays.", file=sys.stderr)
        return

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"Created output directory: {output_dir}")

    try:
        # Convert lists back to numpy arrays
        segment_array = np.array(results['segment'])
        pred_array = np.array(results['pred'])

        # Save the .npy files
        segment_path = os.path.join(output_dir, 'segment.npy')
        pred_path = os.path.join(output_dir, 'pred.npy')

        np.save(segment_path, segment_array)
        print(f"Successfully saved: {segment_path}")

        np.save(pred_path, pred_array)
        print(f"Successfully saved: {pred_path}")

    except Exception as e:
        print(f"An error occurred while saving the result arrays: {e}", file=sys.stderr)


def poll_status(task_id: str, output_dir: str):
    """Polls the server for the task status and retrieves the results."""
    status_url = f"{SERVER_URL}/status/{task_id}"

    print(f"\nPolling for status updates...")

    while True:
        try:
            response = requests.get(status_url, timeout=10)
            response.raise_for_status()

            data = response.json()
            status = data.get("status")
            progress = data.get("progress")

            # Display progress on a single line
            sys.stdout.write(f"\rStatus: {status} | Progress: {progress}      ")
            sys.stdout.flush()

            if status == "completed":
                print("\n\nTask complete. Retrieving results...")
                save_results(data.get("results"), output_dir)
                break
            elif status == "error":
                print(f"\n\nTask failed on the server with an error: {progress}", file=sys.stderr)
                break

            time.sleep(3)

        except requests.exceptions.RequestException as e:
            print(f"\nAn error occurred while polling for status: {e}", file=sys.stderr)
            break
        except KeyboardInterrupt:
            print("\nPolling stopped by user.")
            break


def upload_and_start_task(zip_path: str):
    """Uploads the scene zip file to the server to start the processing task."""
    start_url = f"{SERVER_URL}/run_test/"

    if not os.path.exists(zip_path) or not zip_path.endswith('.zip'):
        print(f"Error: The specified file does not exist or is not a .zip file: '{zip_path}'", file=sys.stderr)
        return None

    print(f"Uploading {os.path.basename(zip_path)} to the server to start processing...")

    try:
        with open(zip_path, 'rb') as f:
            files = {'file': (os.path.basename(zip_path), f, 'application/zip')}
            response = requests.post(start_url, files=files, timeout=60)  # Increased timeout for upload

        response.raise_for_status()

        task_id = response.json().get("task_id")
        if task_id:
            print(f"File uploaded successfully. Task started with ID: {task_id}")
            return task_id
        else:
            print("Error: The server did not return a valid task ID.", file=sys.stderr)
            return None

    except requests.exceptions.ConnectionError:
        print(f"\nConnection Error: Could not connect to the server at {SERVER_URL}.", file=sys.stderr)
        print("Please ensure 'server.py' is running and accessible.", file=sys.stderr)
        return None
    except requests.exceptions.RequestException as e:
        print(f"\nAn error occurred during file upload: {e}", file=sys.stderr)
        return None


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="输入单个场景文件夹用于向服务器传输数据",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        '-i',
        '--input_folder',
        type=str,
        required=True,
        help='指定输入场景的路径。'
    )
    args = parser.parse_args()
    scene_folder = args.input_folder

    DEBUG_ZIP_FILE = os.path.join(scene_folder, "scene.zip")
    DEBUG_OUTPUT_DIR = os.path.join(scene_folder, "output")
    if not os.path.exists(DEBUG_OUTPUT_DIR):
        os.makedirs(DEBUG_OUTPUT_DIR)

    zip_file_path = DEBUG_ZIP_FILE
    output_directory = DEBUG_OUTPUT_DIR
    print(f"DEBUG模式启用: 使用 ZIP文件: {zip_file_path}, 输出目录: {output_directory}")
    task_id = upload_and_start_task(zip_file_path)
    if task_id:
        poll_status(task_id, output_directory)