#!/usr/bin/env python3
"""Download files from HF Space via SSH"""

from ssh_helper import execute_ssh_command
import os

# Create results directory
os.makedirs("results", exist_ok=True)

# Download summary CSV
print("[*] Downloading summary CSV...")
result = execute_ssh_command("cat /tmp/chronos2_forecast_summary.csv")
if result['success']:
    with open("results/chronos2_forecast_summary.csv", 'w') as f:
        f.write(result['stdout'])
    print(f"[+] Saved: results/chronos2_forecast_summary.csv")
else:
    print(f"[!] Failed: {result['stderr']}")

# Download full inference log
print("\n[*] Downloading full inference log...")
result = execute_ssh_command("cat /tmp/full_inference.log")
if result['success']:
    with open("results/full_inference.log", 'w') as f:
        f.write(result['stdout'])
    print(f"[+] Saved: results/full_inference.log")
else:
    print(f"[!] Failed: {result['stderr']}")

# For parquet file, use base64 encoding
print("\n[*] Downloading forecast parquet file (base64 encoded)...")
result = execute_ssh_command("base64 -w 0 /tmp/chronos2_forecasts_14day.parquet")
if result['success']:
    import base64
    parquet_data = base64.b64decode(result['stdout'])
    with open("results/chronos2_forecasts_14day.parquet", 'wb') as f:
        f.write(parquet_data)
    file_size = len(parquet_data) / 1024
    print(f"[+] Saved: results/chronos2_forecasts_14day.parquet ({file_size:.2f} KB)")
else:
    print(f"[!] Failed: {result['stderr']}")

print("\n[+] All files downloaded to results/")