import requests
import time
from typing import Dict, Any, Optional
import os

class MLTracker:
    """Python client for ML Tracker - Free W&B Alternative"""

    def __init__(self, api_key: str, base_url: str = "https://your-space-url.hf.space"):
        self.api_key = api_key
        self.base_url = base_url.rstrip('/')
        self.current_experiment = None
        self.step = 0

    def init(self, experiment_name: str, config: Optional[Dict[str, Any]] = None):
        self.current_experiment = experiment_name
        self.step = 0
        if config:
            self.log_config(experiment_name, config)

    def log_config(self, experiment_name: str, config: Dict[str, Any]):
        try:
            response = requests.post(
                f"{self.base_url}/api/log",
                json={
                    "api_key": self.api_key,
                    "experiment": experiment_name,
                    "step": 0,
                    "metrics": {},
                    "config": config,
                    "timestamp": time.time()
                }
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Warning: Failed to log config: {e}")
            return None

    def log(self, metrics: Dict[str, Any], step: Optional[int] = None):
        if not self.current_experiment:
            raise ValueError("No experiment initialized. Call init() first.")
        if step is None:
            self.step += 1
            step = self.step
        else:
            self.step = max(self.step, step)
        try:
            response = requests.post(
                f"{self.base_url}/api/log",
                json={
                    "api_key": self.api_key,
                    "experiment": self.current_experiment,
                    "step": step,
                    "metrics": metrics,
                    "timestamp": time.time()
                }
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Warning: Failed to log metrics: {e}")
            return None

    def delete_experiment(self, experiment_name: str):
        try:
            response = requests.delete(
                f"{self.base_url}/api/experiment/{experiment_name}",
                params={"api_key": self.api_key}
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Error deleting experiment: {e}")
            return None

    def get_experiments(self):
        try:
            response = requests.get(
                f"{self.base_url}/api/experiments",
                params={"api_key": self.api_key}
            )
            response.raise_for_status()
            return response.json().get("experiments", [])
        except Exception as e:
            print(f"Error fetching experiments: {e}")
            return []

    def get_experiment(self, experiment_name: str):
        try:
            response = requests.get(
                f"{self.base_url}/api/experiment/{experiment_name}",
                params={"api_key": self.api_key}
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Error fetching experiment: {e}")
            return None

    def finish(self):
        self.current_experiment = None
        self.step = 0


# Global tracker and convenience functions
_global_tracker = None

def init(experiment_name: str, config: Optional[Dict[str, Any]] = None, 
         api_key: Optional[str] = None, base_url: Optional[str] = None):
    global _global_tracker
    if api_key is None:
        api_key = os.environ.get("ML_TRACKER_API_KEY")
        if not api_key:
            raise ValueError("API key not provided and ML_TRACKER_API_KEY not set")
    if base_url is None:
        base_url = os.environ.get("ML_TRACKER_BASE_URL", "https://your-space-url.hf.space")
    _global_tracker = MLTracker(api_key, base_url)
    _global_tracker.init(experiment_name, config)

def log(metrics: Dict[str, Any], step: Optional[int] = None):
    if _global_tracker is None:
        raise ValueError("No experiment initialized. Call init() first.")
    return _global_tracker.log(metrics, step)

def finish():
    global _global_tracker
    if _global_tracker:
        _global_tracker.finish()
        _global_tracker = None


if __name__ == "__main__":
    # Example usage
    tracker = MLTracker(
        api_key="your-api-key-here",
        base_url="https://your-space-url.hf.space"
    )
    tracker.init("my_experiment", config={
        "model": "ResNet50",
        "dataset": "CIFAR-10",
        "learning_rate": 0.001,
        "batch_size": 32
    })
    for epoch in range(10):
        loss = 2.0 * (0.9 ** epoch) + 0.1
        accuracy = 1.0 - (0.9 ** epoch)
        tracker.log({
            "loss": loss,
            "accuracy": accuracy,
            "epoch": epoch
        })
    tracker.finish()

    init("another_experiment", config={"model": "BERT"})
    for step in range(5):
        log({"loss": 1.0 / (step + 1), "step": step})
    finish()

    experiments = tracker.get_experiments()
    print(f"Found {len(experiments)} experiments")

    if experiments:
        exp_data = tracker.get_experiment(experiments[0]["experiment"])
        if exp_data:
            print(f"Experiment has {len(exp_data.get('metrics', []))} metric entries")