"""
MCA SDK Example: Autolog Demo with sklearn

Demonstrates zero-code instrumentation using autolog() with Registry API integration.
"""

import os
import sys
import time
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification

# Add both sdk-examples (for shared.py) and mca-prototype root (for mca_sdk) to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

from shared import (
    print_gcp_auth_status,
    print_collector_status,
    COLLECTOR_HTTP_ENDPOINT,
    REGISTRY_URL,
    INTERNAL_MODEL_ID,
)

from mca_sdk import MCAClient, autolog
from mca_sdk.registry import RegistryClient


def main():
    """Demonstrate sklearn autolog instrumentation with GCP integration."""
    print("MCA SDK - Autolog Demo (Scikit-learn)\n")

    # Step 1-2: Verify prerequisites
    if not print_gcp_auth_status() or not print_collector_status():
        sys.exit(1)

    # Step 3: Registry integration
    print("\n=== Registry Integration ===")
    model_id = os.environ.get("MODEL_ID", INTERNAL_MODEL_ID)

    registry_client = RegistryClient(url=REGISTRY_URL, use_gcp_auth=True, timeout=10.0)

    # Fetch specific model config
    model_config = registry_client.fetch_model_config(model_id)
    print(f"✓ Using: {model_config.service_name} ({model_config.team_name})")
    if model_config.thresholds:
        print(f"✓ Thresholds: {model_config.thresholds}")
    print()

    # Step 4: Initialize MCAClient
    print("=== SDK Initialization ===")
    client = MCAClient(
        service_name=model_config.service_name,
        model_id=model_id,
        team_name=model_config.team_name,
        collector_endpoint=COLLECTOR_HTTP_ENDPOINT,
        allow_insecure_collector=True,
        registry_url=REGISTRY_URL,
        use_gcp_auth=True,
    )
    print("✓ MCAClient initialized\n")

    try:
        # Step 5: Enable autolog (one line!)
        print("=== Autolog Activation ===")
        autolog()
        print("✓ Autolog enabled (sklearn auto-detected and patched)\n")

        # Step 6: Use sklearn normally - automatically instrumented!
        print("=== Model Training ===")
        X, y = make_classification(n_samples=1000, n_features=20, random_state=42)
        model = RandomForestClassifier(n_estimators=10, random_state=42)
        model.fit(X, y)
        print("✓ Model trained\n")

        # Step 7: Make predictions - automatically instrumented!
        print("=== Automatic Prediction Tracking ===")
        for i in range(5):
            X_test = np.random.randn(10, 20)
            predictions = model.predict(X_test)
            print(f"  ✓ Prediction {i + 1}: {len(predictions)} results (auto-tracked)")
            time.sleep(0.1)

        print("\n✓ All predictions automatically tracked!")
        print("  - Metrics: latency, model_class=RandomForestClassifier, framework=sklearn")
        print("  - Traces: sklearn.RandomForestClassifier.predict spans\n")

        # Step 8: Flush
        print("=== Flush Telemetry ===")
        client.shutdown(timeout_millis=5000)
        registry_client.close()
        print("✓ Telemetry sent to collector\n")

        # Step 9: Verification
        print("=== Verify in GCP ===")
        print("Cloud Logging: https://console.cloud.google.com/logs/query")
        print('  Filter: logName="projects/bhsf-ai-team-projects/logs/emms-model-telemetry"')
        print(f'          AND labels."service.name"="{model_config.service_name}"')
        print("\nCloud Trace: https://console.cloud.google.com/traces/list")
        print(f'  Search: service.name="{model_config.service_name}"')
        print("  Look for: sklearn.RandomForestClassifier.predict spans")

    except Exception as e:
        print(f"\n✗ Error: {e}")
        client.shutdown()
        registry_client.close()
        raise


if __name__ == "__main__":
    main()
