"""Session management for Armory SDK.

This module provides MLflow-like experiment tracking with automatic
retrospective generation and skill creation.
"""

import atexit
import threading
import time
import traceback
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional, Union

from armory.client import ArmoryClient
from armory.config import ArmoryConfig, get_config
from armory.models import Session as SessionModel, Skill
from armory.logger import get_logger

logger = get_logger("session")


@dataclass
class MetricEntry:
    """A single metric entry with optional step."""

    name: str
    value: float
    step: Optional[int] = None
    timestamp: float = field(default_factory=time.time)


@dataclass
class ArtifactEntry:
    """An artifact entry."""

    name: str
    path: str
    artifact_type: str = "file"


class Session:
    """Active experiment session.

    Tracks metrics, parameters, and artifacts for an experiment.
    Automatically calls retrospective API on session end.

    Usage:
        with armory.start_session(project="my-project", goal="train model") as session:
            armory.log_param("lr", 0.001)
            for epoch in range(10):
                loss = train()
                armory.log_metric("loss", loss, step=epoch)
    """

    def __init__(
        self,
        project_id: str,
        goal: str,
        model_type: str = "llm",
        session_type: str = "training",
        base_model: str = "",
        dataset_path: str = "",
        experiment_id: Optional[str] = None,
        config: Optional[ArmoryConfig] = None,
        client: Optional[ArmoryClient] = None,
    ):
        self.config = config or get_config()
        self.client = client or ArmoryClient(self.config)
        self._owns_client = client is None

        # Session info
        self.project_id = project_id
        self.goal = goal
        self.model_type = model_type
        self.session_type = session_type
        self.base_model = base_model
        self.dataset_path = dataset_path
        self.experiment_id = experiment_id

        # Tracking data
        self.params: Dict[str, Any] = {}
        self.metrics: Dict[str, List[MetricEntry]] = {}
        self.tags: Dict[str, str] = {}
        self.artifacts: List[ArtifactEntry] = []
        self.errors: List[str] = []

        # Session state
        self._session_model: Optional[SessionModel] = None
        self._started = False
        self._ended = False
        self._start_time: Optional[float] = None
        self._end_time: Optional[float] = None

        # System metrics collection
        self._system_metrics_thread: Optional[threading.Thread] = None
        self._stop_system_metrics = threading.Event()

    @property
    def session_id(self) -> Optional[str]:
        """Get the session ID."""
        return self._session_model.id if self._session_model else None

    @property
    def status(self) -> str:
        """Get current session status."""
        if self._session_model:
            return self._session_model.status
        return "not_started"

    @property
    def duration_seconds(self) -> float:
        """Get session duration in seconds."""
        if self._start_time is None:
            return 0.0
        end = self._end_time or time.time()
        return end - self._start_time

    def start(self) -> "Session":
        """Start the session."""
        if self._started:
            return self

        # Create session on server
        self._session_model = self.client.create_session(
            project_id=self.project_id,
            goal=self.goal,
            model_type=self.model_type,
            session_type=self.session_type,
            base_model=self.base_model,
            dataset_path=self.dataset_path,
        )

        # Start session on server
        self._session_model = self.client.start_session(
            project_id=self.project_id,
            session_id=self._session_model.id,
        )

        self._started = True
        self._start_time = time.time()

        # Start system metrics collection
        if self.config.auto_log_system_metrics:
            self._start_system_metrics_collection()

        if self.config.verbose:
            logger.info(f"Session started: {self._session_model.id}")

        return self

    def end(
        self,
        status: str = "completed",
        error_message: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        """End the session and trigger retrospective.

        Args:
            status: Session status ("completed", "failed")
            error_message: Error message if failed

        Returns:
            Retrospective response if auto_retrospective is enabled
        """
        if self._ended or not self._started:
            return None

        self._ended = True
        self._end_time = time.time()

        # Stop system metrics collection
        self._stop_system_metrics.set()
        if self._system_metrics_thread:
            self._system_metrics_thread.join(timeout=2.0)

        # Get final metrics
        final_metrics = self.get_final_metrics()

        try:
            if status == "completed":
                # Complete session on server
                self._session_model = self.client.complete_session(
                    project_id=self.project_id,
                    session_id=self._session_model.id,
                    metrics=final_metrics,
                )
            else:
                # Fail session on server
                self._session_model = self.client.fail_session(
                    project_id=self.project_id,
                    session_id=self._session_model.id,
                    error_message=error_message or "Unknown error",
                )
        except Exception as e:
            if self.config.verbose:
                logger.warning(f"Failed to update session status: {e}")

        # Auto-retrospective
        retrospective_response = None
        if self.config.auto_retrospective:
            retrospective_response = self._run_retrospective(status, final_metrics)

        if self.config.verbose:
            logger.info(f"Session ended: {self._session_model.id} ({status})")

        # Clean up client if we own it
        if self._owns_client:
            self.client.close()

        return retrospective_response

    def _run_retrospective(
        self,
        status: str,
        metrics: Dict[str, float],
    ) -> Optional[Dict[str, Any]]:
        """Run retrospective analysis."""
        try:
            # Determine outcome
            if status == "completed":
                # Check if it's actually successful based on metrics
                success_metric = self.config.success_metric_name
                if success_metric in metrics:
                    if metrics[success_metric] >= self.config.min_success_metric:
                        outcome = "success"
                    else:
                        outcome = "partial"
                else:
                    outcome = "success"  # Assume success if no metric to check
            else:
                outcome = "failure"

            # Build error log
            error_log = None
            if self.errors:
                error_log = "\n".join(self.errors)

            # Call retrospective API
            response = self.client.retrospective(
                session_id=self._session_model.id,
                outcome=outcome,
                metrics=metrics,
                error_log=error_log,
                hyperparameters=self.params,
                observations=None,  # Could add observations feature
                auto_create_skill=self.config.auto_skill_creation and outcome == "success",
            )

            if self.config.verbose:
                logger.info(f"Retrospective: {response.summary}")
                if response.skill_created:
                    logger.info(f"Skill created: {response.skill_id}")

            return response.model_dump()

        except Exception as e:
            if self.config.verbose:
                logger.error(f"Retrospective failed: {e}")
            return None

    def get_final_metrics(self) -> Dict[str, float]:
        """Get final metric values (last value for each metric)."""
        final = {}
        for name, entries in self.metrics.items():
            if entries:
                final[name] = entries[-1].value
        return final

    def log_param(self, key: str, value: Any) -> None:
        """Log a parameter."""
        self.params[key] = value

    def log_params(self, params: Dict[str, Any]) -> None:
        """Log multiple parameters."""
        self.params.update(params)

    def log_metric(
        self,
        key: str,
        value: float,
        step: Optional[int] = None,
    ) -> None:
        """Log a metric value."""
        if key not in self.metrics:
            self.metrics[key] = []

        self.metrics[key].append(MetricEntry(name=key, value=value, step=step))

    def log_metrics(
        self,
        metrics: Dict[str, float],
        step: Optional[int] = None,
    ) -> None:
        """Log multiple metrics."""
        for key, value in metrics.items():
            self.log_metric(key, value, step)

    def log_artifact(
        self,
        name: str,
        path: str,
        artifact_type: str = "file",
    ) -> None:
        """Log an artifact."""
        self.artifacts.append(ArtifactEntry(name=name, path=path, artifact_type=artifact_type))

    def log_error(self, error: Union[str, Exception]) -> None:
        """Log an error."""
        if isinstance(error, Exception):
            self.errors.append(f"{type(error).__name__}: {str(error)}\n{traceback.format_exc()}")
        else:
            self.errors.append(str(error))

    def set_tag(self, key: str, value: str) -> None:
        """Set a tag."""
        self.tags[key] = value

    def set_tags(self, tags: Dict[str, str]) -> None:
        """Set multiple tags."""
        self.tags.update(tags)

    def _start_system_metrics_collection(self) -> None:
        """Start collecting system metrics in background."""

        def collect_metrics():
            try:
                import psutil
            except ImportError:
                return

            while not self._stop_system_metrics.is_set():
                try:
                    # CPU
                    self.log_metric("system/cpu_percent", psutil.cpu_percent())

                    # Memory
                    mem = psutil.virtual_memory()
                    self.log_metric("system/memory_percent", mem.percent)
                    self.log_metric("system/memory_used_gb", mem.used / (1024**3))

                    # GPU (if available)
                    try:
                        import pynvml

                        pynvml.nvmlInit()
                        device_count = pynvml.nvmlDeviceGetCount()
                        for i in range(device_count):
                            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                            util = pynvml.nvmlDeviceGetUtilizationRates(handle)
                            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
                            self.log_metric(f"system/gpu{i}_util", util.gpu)
                            self.log_metric(
                                f"system/gpu{i}_memory_used_gb",
                                mem_info.used / (1024**3),
                            )
                    except Exception:
                        pass

                except Exception:
                    pass

                # Wait 30 seconds between collections
                self._stop_system_metrics.wait(30.0)

        self._system_metrics_thread = threading.Thread(target=collect_metrics, daemon=True)
        self._system_metrics_thread.start()

    def __enter__(self) -> "Session":
        """Context manager entry."""
        return self.start()

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Context manager exit."""
        if exc_type is not None:
            # Exception occurred
            self.log_error(exc_val)
            self.end(status="failed", error_message=str(exc_val))
        else:
            self.end(status="completed")


# =========================================================================
# Global session management
# =========================================================================

_active_session: Optional[Session] = None
_session_lock = threading.Lock()


def get_active_session() -> Optional[Session]:
    """Get the currently active session."""
    return _active_session


def _set_active_session(session: Optional[Session]) -> None:
    """Set the active session."""
    global _active_session
    with _session_lock:
        _active_session = session


@contextmanager
def start_session(
    project: str,
    goal: str,
    model_type: str = "llm",
    session_type: str = "training",
    base_model: str = "",
    dataset_path: str = "",
    experiment_id: Optional[str] = None,
    nested: bool = False,
):
    """Start a new Armory session.

    This is the main entry point for experiment tracking.

    Args:
        project: Project ID or name
        goal: What you're trying to achieve
        model_type: Type of model (llm, vision, audio, etc.)
        session_type: Type of session (training, inference, evaluation)
        base_model: Base model being used
        dataset_path: Path to dataset
        experiment_id: Optional experiment ID to group sessions
        nested: Allow nested sessions

    Yields:
        Active session object

    Example:
        with armory.start_session(project="my-project", goal="finetune llama") as session:
            armory.log_param("lr", 0.001)
            armory.log_metric("loss", 0.5)
    """
    global _active_session

    if _active_session is not None and not nested:
        raise RuntimeError(
            "A session is already active. Use nested=True to allow nested sessions, "
            "or end the current session first."
        )

    session = Session(
        project_id=project,
        goal=goal,
        model_type=model_type,
        session_type=session_type,
        base_model=base_model,
        dataset_path=dataset_path,
        experiment_id=experiment_id,
    )

    previous_session = _active_session
    _set_active_session(session)

    try:
        with session:
            yield session
    finally:
        _set_active_session(previous_session)


# Clean up any active session on exit
def _cleanup_on_exit():
    global _active_session
    if _active_session is not None:
        try:
            _active_session.end(status="failed", error_message="Process terminated unexpectedly")
        except Exception:
            pass


atexit.register(_cleanup_on_exit)
