"""
Armory Run - Session management for ML experiments

This module provides the core Run class that handles:
- Session creation and lifecycle
- Metric logging with real-time streaming
- Event tracking (train_start, epoch_end, etc.)
- Artifact management
- Summary metrics
- Run resume functionality
"""

import os
import time
import atexit
import threading
import json
import hashlib
from typing import Any, Dict, List, Optional, Union, Callable
from dataclasses import dataclass, field
from pathlib import Path
import requests
from datetime import datetime

from armory.system import SystemMonitor
from armory.logger import get_logger

logger = get_logger("run")


class Summary(dict):
    """
    Summary metrics dictionary that auto-syncs to server.
    Similar to wandb.run.summary
    """
    def __init__(self, run: "Run"):
        super().__init__()
        self._run = run

    def __setitem__(self, key: str, value: Any):
        super().__setitem__(key, value)
        self._sync()

    def update(self, *args, **kwargs):
        super().update(*args, **kwargs)
        self._sync()

    def _sync(self):
        """Sync summary to server"""
        if self._run and self._run._run_id:
            try:
                self._run._api_request(
                    "POST",
                    f"/api/runs/{self._run._run_id}/summary",
                    json={"summary": dict(self)}
                )
            except Exception:
                pass  # Silent fail for summary sync


class Config(dict):
    """
    Configuration dictionary with attribute access.
    Similar to wandb.config
    """
    def __init__(self, data: Dict[str, Any] = None):
        super().__init__(data or {})

    def __getattr__(self, key: str) -> Any:
        try:
            return self[key]
        except KeyError:
            raise AttributeError(f"Config has no attribute '{key}'")

    def __setattr__(self, key: str, value: Any):
        self[key] = value

    def update(self, *args, **kwargs):
        super().update(*args, **kwargs)


class MetricDefinition:
    """Definition for a custom metric"""
    def __init__(
        self,
        name: str,
        step_metric: Optional[str] = None,
        summary: str = "last",  # last, min, max, mean, best, none
        goal: Optional[str] = None,  # minimize, maximize
    ):
        self.name = name
        self.step_metric = step_metric
        self.summary = summary
        self.goal = goal


@dataclass
class RunConfig:
    """Configuration for an Armory run"""
    project: str = "default"
    name: Optional[str] = None
    config: Dict[str, Any] = field(default_factory=dict)
    tags: list = field(default_factory=list)
    notes: str = ""

    # Resume settings
    resume: Union[bool, str] = False  # True, False, "allow", "must", "never", or run_id
    id: Optional[str] = None  # Specific run ID to resume

    # System metrics settings
    log_system_metrics: bool = True  # Automatically log GPU/CPU/RAM
    system_metrics_interval: float = 10.0  # Seconds between system metric logs

    # Armory server settings
    api_url: str = field(default_factory=lambda: os.getenv("ARMORY_API_URL", "http://localhost:8080"))
    api_key: Optional[str] = field(default_factory=lambda: os.getenv("ARMORY_API_KEY"))


class Run:
    """
    Armory Run - tracks a single training session.

    Similar to wandb.Run, this class handles metric logging, event tracking,
    and artifact management for ML experiments.

    Attributes:
        summary: Final metrics summary (like wandb.run.summary)
        config: Run configuration with attribute access
        name: Run name
        id: Run/session ID
        project: Project name
        dir: Local directory for artifacts
    """

    _current: Optional["Run"] = None

    def __init__(self, run_config: RunConfig):
        self._run_config = run_config
        self._run_id: Optional[str] = None
        self._project_id: Optional[str] = None
        self._step: int = 0
        self._start_time: float = time.time()
        self._finished: bool = False
        self._resumed: bool = False
        self._lock = threading.Lock()
        self._system_monitor: Optional[SystemMonitor] = None

        # Public attributes (W&B compatible)
        self.summary = Summary(self)
        self.config = Config(run_config.config)
        self.name = run_config.name
        self.project = run_config.project
        self.tags = list(run_config.tags)
        self.notes = run_config.notes

        # Metric definitions
        self._metric_definitions: Dict[str, MetricDefinition] = {}

        # Watch tracking
        self._watched_models: List[Any] = []
        self._gradient_log_freq: int = 0

        # Alerts
        self._alert_callbacks: List[Callable] = []

        # Local artifact directory
        self.dir = os.path.join(os.getcwd(), "armory_runs", run_config.name or "default")
        os.makedirs(self.dir, exist_ok=True)

        # Register cleanup on exit
        atexit.register(self._cleanup)
    
    @property
    def run_id(self) -> Optional[str]:
        return self._run_id

    @property
    def id(self) -> Optional[str]:
        """Run ID"""
        return self._run_id

    @property
    def session_id(self) -> Optional[str]:
        """Alias for run_id (backwards compatibility)"""
        return self._run_id

    @property
    def step(self) -> int:
        return self._step

    @property
    def resumed(self) -> bool:
        """Whether this run was resumed from a previous run"""
        return self._resumed

    def _api_request(self, method: str, endpoint: str, **kwargs) -> requests.Response:
        """Make an API request to the Armory server"""
        url = f"{self._run_config.api_url}{endpoint}"
        headers = kwargs.pop("headers", {})
        headers["Content-Type"] = "application/json"

        if self._run_config.api_key:
            headers["Authorization"] = f"Bearer {self._run_config.api_key}"

        return requests.request(method, url, headers=headers, **kwargs)
    
    def _get_or_create_project(self) -> str:
        """Get existing project or create a new one"""
        project_resp = self._api_request("GET", "/api/projects")
        projects = project_resp.json() if project_resp.ok else []

        for p in projects:
            if p.get("name") == self._run_config.project:
                self._project_id = p.get("id")
                return self._project_id

        # Create project
        resp = self._api_request("POST", "/api/projects", json={
            "name": self._run_config.project,
            "description": "Created by Armory SDK"
        })
        if resp.ok:
            self._project_id = resp.json().get("id")
            return self._project_id
        else:
            raise RuntimeError(f"Failed to create project: {resp.text}")

    def _try_resume_run(self) -> Optional[str]:
        """Try to resume an existing run"""
        resume = self._run_config.resume

        # If resume is a specific run ID
        if isinstance(resume, str) and resume not in ("allow", "must", "never", "True", "False"):
            try:
                resp = self._api_request("GET", f"/api/runs/{resume}")
                if resp.ok:
                    run_data = resp.json()
                    self._step = run_data.get("last_step", 0)
                    self._resumed = True
                    logger.info(f"Resuming run {resume} from step {self._step}")
                    return resume
            except Exception as e:
                if resume == "must" or self._run_config.resume == "must":
                    raise RuntimeError(f"Failed to resume run {resume}: {e}")
            return None

        # If resume is True or "allow" or "must", try to find a matching run
        if resume in (True, "allow", "must", "True"):
            run_id = self._run_config.id
            if run_id:
                try:
                    resp = self._api_request("GET", f"/api/runs/{run_id}")
                    if resp.ok:
                        run_data = resp.json()
                        self._step = run_data.get("last_step", 0)
                        self._resumed = True
                        logger.info(f"Resuming run {run_id} from step {self._step}")
                        return run_id
                except Exception:
                    pass

            if resume == "must":
                raise RuntimeError("Resume mode is 'must' but no run found to resume")

        return None

    def _create_run(self) -> str:
        """Create a new run on the Armory server"""
        # Get or create project
        project_id = self._get_or_create_project()

        # Try to resume if configured
        if self._run_config.resume and self._run_config.resume != "never":
            resumed_id = self._try_resume_run()
            if resumed_id:
                return resumed_id

        # Create new run
        run_name = self._run_config.name or f"run-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
        self.name = run_name

        resp = self._api_request("POST", f"/api/projects/{project_id}/runs", json={
            "name": run_name,
            "model_type": "embedding",  # Default, can be overridden
            "run_type": "training",
            "goal": self._run_config.notes or f"Training run: {run_name}",
            "config": dict(self.config),
            "tags": self.tags,
        })

        if not resp.ok:
            raise RuntimeError(f"Failed to create run: {resp.text}")

        run = resp.json()
        return run["id"]
    
    def start(self) -> "Run":
        """Start the run and create it on the server"""
        if self._run_id:
            return self

        self._run_id = self._create_run()

        # Log train_start event (only if not resumed)
        if not self._resumed:
            self._log_event("train_start", {
                "config": dict(self.config),
                "tags": self.tags,
            })

        # Start system metrics monitoring if enabled
        if self._run_config.log_system_metrics:
            self._system_monitor = SystemMonitor(
                callback=lambda m: self.log_system(m),
                interval=self._run_config.system_metrics_interval,
            )
            self._system_monitor.start()

        Run._current = self
        return self
    
    def _log_event(self, event: str, data: Optional[Dict[str, Any]] = None):
        """Log an event to the server"""
        if not self._run_id:
            return
        
        payload = {
            "step": self._step,
            "event": event,
            "metrics": data.get("metrics", {}) if data else {},
            "message": data.get("message", "") if data else "",
        }
        
        try:
            self._api_request("POST", f"/api/runs/{self._run_id}/log", json=payload)
        except Exception as e:
            logger.warning(f"Failed to log event: {e}")
    
    def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True):
        """
        Log metrics to the current step.

        Args:
            metrics: Dictionary of metric names to values
            step: Optional step number (auto-increments if not provided)
            commit: Whether to commit this step (increment step counter)
        """
        if not self._run_id:
            self.start()

        with self._lock:
            if step is not None:
                self._step = step

            # Extract numeric metrics
            numeric_metrics = {k: float(v) if isinstance(v, (int, float)) else v
                              for k, v in metrics.items() if isinstance(v, (int, float))}

            # Update summary with latest values (for final metrics on complete)
            for k, v in numeric_metrics.items():
                dict.__setitem__(self.summary, k, v)

            payload = {
                "step": self._step,
                "metrics": numeric_metrics,
            }

            try:
                resp = self._api_request("POST", f"/api/runs/{self._run_id}/log", json=payload)
                if not resp.ok:
                    logger.warning(f"Failed to log metrics: {resp.status_code} - {resp.text}")
            except Exception as e:
                logger.warning(f"Failed to log metrics: {e}")

            if commit:
                self._step += 1
    
    def log_event(self, event: str, message: str = "", metrics: Optional[Dict[str, float]] = None):
        """
        Log a training event.

        Supported events:
        - train_start: Training started
        - epoch_start: Epoch started
        - epoch_end: Epoch ended
        - eval_start: Evaluation started
        - eval_end: Evaluation ended
        - checkpoint: Checkpoint saved
        - train_end: Training completed
        - error: Error occurred
        """
        if not self._run_id:
            self.start()

        payload = {
            "step": self._step,
            "event": event,
            "message": message,
            "metrics": metrics or {},
        }

        try:
            self._api_request("POST", f"/api/runs/{self._run_id}/log", json=payload)
        except Exception as e:
            logger.warning(f"Failed to log event: {e}")

    def log_system(self, system_metrics=None):
        """
        Log system metrics (GPU/CPU/Memory usage).

        Args:
            system_metrics: Optional SystemMetrics object. If not provided,
                           will collect current system metrics automatically.
        """
        if not self._run_id:
            self.start()

        if system_metrics is None:
            try:
                from armory.system import collect_system_metrics
                system_metrics = collect_system_metrics()
            except Exception as e:
                logger.warning(f"Failed to collect system metrics: {e}")
                return

        payload = {
            "step": self._step,
            "system": system_metrics.to_dict() if hasattr(system_metrics, 'to_dict') else system_metrics,
        }

        try:
            self._api_request("POST", f"/api/runs/{self._run_id}/log", json=payload)
        except Exception as e:
            logger.warning(f"Failed to log system metrics: {e}")

    def define_metric(
        self,
        name: str,
        step_metric: Optional[str] = None,
        summary: str = "last",
        goal: Optional[str] = None,
    ):
        """
        Define how a metric should be tracked and summarized.

        Args:
            name: Metric name (supports wildcards like "train/*")
            step_metric: Name of metric to use as x-axis (e.g., "epoch")
            summary: How to summarize: "last", "min", "max", "mean", "best", "none"
            goal: Optimization goal: "minimize" or "maximize"

        Example:
            run.define_metric("val/loss", summary="min", goal="minimize")
            run.define_metric("val/*", step_metric="epoch")
        """
        self._metric_definitions[name] = MetricDefinition(
            name=name,
            step_metric=step_metric,
            summary=summary,
            goal=goal,
        )

    def watch(
        self,
        models: Any,
        log: str = "gradients",
        log_freq: int = 1000,
        log_graph: bool = False,
    ):
        """
        Watch a PyTorch model for gradient and parameter logging.

        Args:
            models: Model or list of models to watch
            log: What to log - "gradients", "parameters", "all", or None
            log_freq: How often to log (in steps)
            log_graph: Whether to log the model graph

        Example:
            run.watch(model, log="all", log_freq=100)
        """
        if not isinstance(models, (list, tuple)):
            models = [models]

        self._watched_models.extend(models)
        self._gradient_log_freq = log_freq

        # Try to register gradient hooks if PyTorch is available
        try:
            import torch

            def make_hook(name):
                def hook(grad):
                    if self._step % log_freq == 0:
                        grad_data = grad.detach().cpu()
                        self.log({
                            f"gradients/{name}/mean": float(grad_data.mean()),
                            f"gradients/{name}/std": float(grad_data.std()),
                            f"gradients/{name}/max": float(grad_data.abs().max()),
                        }, commit=False)
                    return grad
                return hook

            for model in models:
                if hasattr(model, 'named_parameters'):
                    for name, param in model.named_parameters():
                        if param.requires_grad and log in ("gradients", "all"):
                            param.register_hook(make_hook(name))

            if log_graph:
                # Log model architecture
                for i, model in enumerate(models):
                    self.log_event("model_graph", f"Model {i}: {type(model).__name__}")

        except ImportError:
            logger.warning("PyTorch not available, watch() disabled")

    def alert(
        self,
        title: str,
        text: str,
        level: str = "INFO",
        wait_duration: float = 0,
    ):
        """
        Send an alert notification.

        Args:
            title: Alert title
            text: Alert message
            level: Severity level - "INFO", "WARN", "ERROR"
            wait_duration: Minimum time between alerts of same title (seconds)

        Example:
            run.alert("Training Complete", "Final accuracy: 95%", level="INFO")
        """
        if not self._run_id:
            self.start()

        payload = {
            "step": self._step,
            "event": "alert",
            "message": f"[{level}] {title}: {text}",
            "metrics": {"alert_level": level},
        }

        try:
            self._api_request("POST", f"/api/runs/{self._run_id}/log", json=payload)
        except Exception as e:
            logger.warning(f"Failed to send alert: {e}")

        # Call registered alert callbacks
        for callback in self._alert_callbacks:
            try:
                callback(title, text, level)
            except Exception:
                pass

    def save(self, path: str, policy: str = "live"):
        """
        Save a file to the run's artifact directory.

        Args:
            path: Path to file to save (glob patterns supported)
            policy: When to upload - "live", "now", "end"

        Example:
            run.save("model.pt")
            run.save("checkpoints/*.pt")
        """
        from pathlib import Path as P
        import shutil
        import glob as g

        paths = g.glob(path) if "*" in path else [path]
        for p in paths:
            src = P(p).resolve()
            if src.exists():
                dst = (P(self.dir) / src.name).resolve()

                # Only copy if source and destination are different
                if src != dst:
                    shutil.copy2(src, dst)
                    logger.debug(f"Saved {src.name} to {self.dir}")

                # Upload to server if policy is "now" or "live"
                if policy in ("now", "live"):
                    self._upload_file(dst)

    def _upload_file(self, path: Path):
        """Upload a file to the Armory server"""
        if not self._run_id:
            return

        try:
            with open(path, "rb") as f:
                files = {"file": (path.name, f)}
                resp = requests.post(
                    f"{self._run_config.api_url}/api/runs/{self._run_id}/artifacts",
                    files=files,
                    headers={"Authorization": f"Bearer {self._run_config.api_key}"} if self._run_config.api_key else {}
                )
                if resp.ok:
                    logger.debug(f"Uploaded {path.name}")
        except Exception as e:
            logger.warning(f"Failed to upload {path.name}: {e}")

    def log_artifact(
        self,
        artifact: Any,
        name: Optional[str] = None,
        type: str = "dataset",
    ):
        """
        Log an artifact (dataset, model, etc.).

        Args:
            artifact: Artifact object or path
            name: Artifact name
            type: Artifact type - "dataset", "model", "checkpoint", etc.

        Example:
            run.log_artifact("data/train.csv", name="training-data", type="dataset")
            run.log_artifact(model, name="final-model", type="model")
        """
        if isinstance(artifact, str):
            # It's a file path
            path = Path(artifact)
            if path.exists():
                self.save(str(path))
        else:
            # Try to save as pickle or torch
            artifact_name = name or f"artifact_{self._step}"
            try:
                import torch
                artifact_path = Path(self.dir) / f"{artifact_name}.pt"
                torch.save(artifact, artifact_path)
                self._upload_file(artifact_path)
            except ImportError:
                import pickle
                artifact_path = Path(self.dir) / f"{artifact_name}.pkl"
                with open(artifact_path, "wb") as f:
                    pickle.dump(artifact, f)
                self._upload_file(artifact_path)

    def finish(self, exit_code: int = 0, quiet: bool = False):
        """
        Finish the run and mark the session as completed.

        Args:
            exit_code: Exit code (0 for success, non-zero for failure)
            quiet: Suppress output messages
        """
        if self._finished:
            logger.debug("Run already finished, skipping")
            return

        self._finished = True

        # Stop system metrics monitoring
        if self._system_monitor:
            self._system_monitor.stop()
            self._system_monitor = None

        if not self._run_id:
            logger.debug("No run_id, cannot finish")
            return

        logger.info(f"Finishing run {self._run_id}...")

        # Log train_end event
        self._log_event("train_end", {
            "message": f"Training completed with exit code {exit_code}",
        })

        # Mark run as completed or failed
        try:
            run_resp = self._api_request("GET", f"/api/runs/{self._run_id}")
            if run_resp.ok:
                run_data = run_resp.json()
                project_id = run_data.get("project_id", "default")

                if exit_code == 0:
                    final_metrics = dict(self.summary)
                    logger.info(f"Sending final metrics: {len(final_metrics)} metrics")
                    logger.debug(f"Final metrics: {final_metrics}")
                    resp = self._api_request("POST", f"/api/projects/{project_id}/runs/{self._run_id}/complete",
                                     json={"metrics": final_metrics})
                    if resp.ok:
                        logger.info("Run marked as completed")
                    else:
                        logger.error(f"Failed to mark run as completed: {resp.status_code} - {resp.text}")
                else:
                    resp = self._api_request("POST", f"/api/projects/{project_id}/runs/{self._run_id}/fail",
                                     json={"error": f"Exit code: {exit_code}"})
                    if resp.ok:
                        logger.info("Run marked as failed")
                    else:
                        logger.error(f"Failed to mark run as failed: {resp.status_code} - {resp.text}")
            else:
                logger.error(f"Failed to get run data: {run_resp.status_code} - {run_resp.text}")
        except Exception as e:
            logger.error(f"Failed to finish run: {e}")

        duration = time.time() - self._start_time
        logger.info(f"Run finished in {duration:.2f}s")
        
        Run._current = None
    
    def _cleanup(self):
        """Cleanup handler called on exit"""
        if not self._finished:
            self.finish(quiet=True)
    
    def __enter__(self) -> "Run":
        return self.start()
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        exit_code = 1 if exc_type else 0
        self.finish(exit_code=exit_code)


def init(
    project: str = "default",
    name: Optional[str] = None,
    config: Optional[Dict[str, Any]] = None,
    tags: Optional[list] = None,
    notes: str = "",
    resume: Union[bool, str] = False,
    id: Optional[str] = None,
    log_system_metrics: bool = True,
    system_metrics_interval: float = 10.0,
    **kwargs
) -> Run:
    """
    Initialize a new Armory run.

    This is the main entry point for the Armory SDK, similar to wandb.init().

    Args:
        project: Project name (will be created if doesn't exist)
        name: Optional run name
        config: Hyperparameters and configuration dict
        tags: List of tags for the run
        notes: Description or notes about the run
        resume: Resume mode - True, False, "allow", "must", "never", or specific run_id
        id: Specific run ID (for resume or to set a custom ID)
        log_system_metrics: Whether to automatically log GPU/CPU/RAM (default True)
        system_metrics_interval: Seconds between system metric logs (default 10.0)

    Returns:
        Run object for logging metrics and events

    Example:
        # Basic usage
        run = armory.init(
            project="my-embeddings",
            config={"learning_rate": 1e-4, "batch_size": 32}
        )

        for epoch in range(10):
            run.log({"loss": 0.5, "epoch": epoch})

        run.finish()

        # Resume a run
        run = armory.init(
            project="my-embeddings",
            resume=True,
            id="abc123"
        )
    """
    run_config = RunConfig(
        project=project,
        name=name,
        config=config or {},
        tags=tags or [],
        notes=notes,
        resume=resume,
        id=id,
        log_system_metrics=log_system_metrics,
        system_metrics_interval=system_metrics_interval,
    )

    run = Run(run_config)
    return run.start()


def finish(exit_code: int = 0, quiet: bool = False):
    """Finish the current run if one is active"""
    if Run._current:
        Run._current.finish(exit_code=exit_code, quiet=quiet)
