"""
══════════════════════════════════════════════════════════════════════════════
PolicyBridge — Master Agent Training Runner
══════════════════════════════════════════════════════════════════════════════
Trains all 4 ML agents in sequence from the Bronze layer data.
Run this ONCE after loading bronze_historical_data.sql.

Usage:
    cd agents/
    pip install pymysql pandas scikit-learn xgboost sqlalchemy
    python train_all_agents.py

Output models saved to agents/models/:
    agent1_kyc_classifier.pkl     — KYC binary classifier
    agent2_property_risk.pkl      — Property peril regressors + risk classifier
    agent3_underwriting.pkl       — UW binary classifier
    agent4_pricing.pkl            — Premium XGBoost + GLM ensemble

Agent 5 (Issuance) is deterministic — no training required.

Environment:
    Local         → connects to localhost:3306 / bronze  (root / root@123)
    HuggingFace   → reads MYSQL_ADDON_* or MYSQL_* Secrets → Clever Cloud
══════════════════════════════════════════════════════════════════════════════
"""

import sys
import os
import time
import datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

# ── Dynamic environment detection (mirrors agent DB config) ───────────────────
def _is_huggingface() -> bool:
    return (
        os.environ.get("SPACE_ID")            is not None
        or os.environ.get("HUGGINGFACE_SPACE") is not None
        or os.environ.get("MYSQL_ADDON_HOST")  is not None
        or os.environ.get("MYSQL_HOST")        is not None
    )

def _env(addon_key: str, generic_key: str, default: str = "") -> str:
    return os.environ.get(addon_key) or os.environ.get(generic_key) or default

def _db_display() -> str:
    """Returns a human-readable DB connection string for the startup banner."""
    if _is_huggingface():
        host = _env("MYSQL_ADDON_HOST", "MYSQL_HOST", "clever-cloud")
        port = _env("MYSQL_ADDON_PORT", "MYSQL_PORT", "3306")
        db   = _env("MYSQL_ADDON_DB",   "MYSQL_DATABASE", "btvbbpqhvnttzvptguj3")
        return f"Clever Cloud  {host}:{port}/{db}  (bronze_* prefix)"
    return "localhost:3306/bronze  (separate schemas)"

# ─────────────────────────────────────────────────────────────────────────────

def print_banner(title):
    print(f"\n{'═'*65}")
    print(f"  {title}")
    print(f"{'═'*65}")

def train_all():
    env_label = "HuggingFace → Clever Cloud" if _is_huggingface() else "Local → MySQL"

    print_banner("PolicyBridge — Agent Training Pipeline")
    print(f"  Started:     {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"  Environment: {env_label}")
    print(f"  DB Target:   {_db_display()}")
    print(f"  Records:     500 historical submissions (2024)")

    results = {}
    total_start = time.time()

    # ── Agent 1: KYC ─────────────────────────────────────────────────────────
    try:
        from agent1_ssn_kyc import train_kyc_model
        t0  = time.time()
        art = train_kyc_model()
        results["Agent1_KYC"] = {"status": "OK", "elapsed": round(time.time()-t0, 1)}
    except Exception as e:
        results["Agent1_KYC"] = {"status": f"FAILED: {e}", "elapsed": 0}

    # ── Agent 2: Property Risk ────────────────────────────────────────────────
    try:
        from agent2_property_risk import train_property_model
        t0  = time.time()
        art = train_property_model()
        results["Agent2_Property"] = {"status": "OK", "elapsed": round(time.time()-t0, 1)}
    except Exception as e:
        results["Agent2_Property"] = {"status": f"FAILED: {e}", "elapsed": 0}

    # ── Agent 3: Underwriting ────────────────────────────────────────────────
    try:
        from agent3_underwriting import train_uw_model
        t0  = time.time()
        art = train_uw_model()
        results["Agent3_UW"] = {"status": "OK", "elapsed": round(time.time()-t0, 1)}
    except Exception as e:
        results["Agent3_UW"] = {"status": f"FAILED: {e}", "elapsed": 0}

    # ── Agent 4: Pricing ─────────────────────────────────────────────────────
    try:
        from agent4_pricing import train_pricing_model
        t0  = time.time()
        art = train_pricing_model()
        results["Agent4_Pricing"] = {"status": "OK", "elapsed": round(time.time()-t0, 1)}
    except Exception as e:
        results["Agent4_Pricing"] = {"status": f"FAILED: {e}", "elapsed": 0}

    # ── Summary ───────────────────────────────────────────────────────────────
    total = round(time.time() - total_start, 1)
    print_banner("Training Summary")

    all_ok = True
    for agent, res in results.items():
        status = res["status"]
        icon   = "✓" if status == "OK" else "✗"
        print(f"  {icon}  {agent:<25}  {status:<10}  {res['elapsed']}s")
        if status != "OK":
            all_ok = False

    print(f"\n  Total elapsed: {total}s")
    print(f"\n  Models saved to: agents/models/")
    for pkl in sorted(Path("models").glob("*.pkl")):
        size = pkl.stat().st_size // 1024
        print(f"    {pkl.name:<40} {size} KB")

    if all_ok:
        print(f"\n  All 4 agents trained successfully.")
        print(f"  Run the pipeline:")
        print(f"    python agent5_issuance_orchestrator.py")
        print(f"    python agent5_issuance_orchestrator.py --submission SUB-2024-00001")
        print(f"    python agent5_issuance_orchestrator.py --batch --limit 100")
    else:
        print(f"\n  Some agents failed. Check errors above.")
        if _is_huggingface():
            print(f"  HuggingFace: verify MYSQL_ADDON_* Secrets are set in Space Settings.")
        else:
            print(f"  Local: ensure MySQL is running and bronze_historical_data.sql is loaded.")

    return results

if __name__ == "__main__":
    train_all()