"""
AHAD QUANT — Export Modèle Unifié
Combine rl_agent.zip + model_ensemble.pkl + rl_scaler.pkl
en un seul fichier ahad_quant_unified.zip.

Usage :
    python export_unified.py
    python export_unified.py --output mon_modele.zip

Structure du zip produit :
    ahad_quant_unified.zip
    ├── ppo/            ← contenu SB3 intact (policy.pth, optimizer, ...)
    ├── ensemble.pkl    ← ensemble supervisé
    ├── rl_scaler.pkl   ← scaler z-score features
    └── metadata.json   ← version, date, config
"""

import os
import sys
import json
import zipfile
import pickle
import argparse
import io
from datetime import datetime

sys.stdout.reconfigure(encoding="utf-8", errors="replace")

import config


def export_unified(
    ppo_path:      str = None,
    ensemble_path: str = None,
    scaler_path:   str = None,
    output_path:   str = "ahad_quant_unified.zip",
):
    ppo_path      = ppo_path      or getattr(config, "RL_MODEL_PATH",       "rl_agent")
    ensemble_path = ensemble_path or getattr(config, "ENSEMBLE_MODEL_PATH", "model_ensemble.pkl")
    scaler_path   = scaler_path   or getattr(config, "RL_SCALER_PATH",      "rl_scaler.pkl")

    # Ajouter .zip si absent
    ppo_zip = ppo_path if ppo_path.endswith(".zip") else ppo_path + ".zip"

    # ── Vérifications ────────────────────────────────────────────────────────
    missing = []
    for label, path in [("PPO", ppo_zip), ("Ensemble", ensemble_path), ("Scaler", scaler_path)]:
        if not os.path.exists(path):
            missing.append(f"  ❌ {label} : {path}")
    if missing:
        print("\n".join(missing))
        print("\nExécuter dans l'ordre :")
        print("  1. python train.py")
        print("  2. python rl_train.py")
        sys.exit(1)

    print("=" * 55)
    print("  AHAD QUANT — Export Modèle Unifié")
    print("=" * 55)
    print(f"  PPO       : {ppo_zip}")
    print(f"  Ensemble  : {ensemble_path}")
    print(f"  Scaler    : {scaler_path}")
    print(f"  Sortie    : {output_path}")
    print()

    # ── Metadata ─────────────────────────────────────────────────────────────
    metadata = {
        "version":        "ahad_quant-v6-unified",
        "exported_at":    datetime.now().isoformat(),
        "ppo_source":     ppo_zip,
        "ensemble_source": ensemble_path,
        "scaler_source":  scaler_path,
        "coins":          getattr(config, "COINS", []),
        "leverage":       getattr(config, "LEVERAGE", 30),
        "obs_dim":        69,
    }

    # ── Construction du zip unifié ────────────────────────────────────────────
    with zipfile.ZipFile(output_path, "w", compression=zipfile.ZIP_DEFLATED) as out_zip:

        # 1. Contenu PPO → sous-dossier ppo/
        with zipfile.ZipFile(ppo_zip, "r") as ppo_zip_f:
            for item in ppo_zip_f.namelist():
                data = ppo_zip_f.read(item)
                out_zip.writestr(f"ppo/{item}", data)
        print(f"  ✅ PPO intégré (ppo/)")

        # 2. Ensemble pkl
        with open(ensemble_path, "rb") as f:
            out_zip.writestr("ensemble.pkl", f.read())
        print(f"  ✅ Ensemble intégré (ensemble.pkl)")

        # 3. Scaler pkl
        with open(scaler_path, "rb") as f:
            out_zip.writestr("rl_scaler.pkl", f.read())
        print(f"  ✅ Scaler intégré (rl_scaler.pkl)")

        # 4. Metadata
        out_zip.writestr("metadata.json", json.dumps(metadata, indent=2))
        print(f"  ✅ Metadata intégrée (metadata.json)")

    size_mb = os.path.getsize(output_path) / (1024 * 1024)
    print()
    print(f"  💾 {output_path}  ({size_mb:.1f} MB)")
    print("=" * 55)
    print("  Chargement :")
    print(f"    from rl_agent import RLAgent")
    print(f"    agent = RLAgent(unified_path='{output_path}')")
    print("=" * 55)
    return output_path


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ppo",      type=str, default=None)
    parser.add_argument("--ensemble", type=str, default=None)
    parser.add_argument("--scaler",   type=str, default=None)
    parser.add_argument("--output",   type=str, default="ahad_quant_unified.zip")
    args = parser.parse_args()

    export_unified(args.ppo, args.ensemble, args.scaler, args.output)