from __future__ import annotations

import importlib
import platform
import socket
import subprocess
import sys
from collections.abc import Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable, Self

from log_utils import (
    LOG_EVENT_TYPE,
    RUNTIME_ERROR_EVENT_TYPE,
    RUNTIME_INTERRUPTION_EVENT_TYPE,
    exception_log_payload,
)
from models import EventRecord, RunContext, RunStatus
from storage import LocalSessionStore, utc_now


def _detect_frameworks() -> dict[str, str]:
    versions: dict[str, str] = {}
    for module_name in ("torch", "tensorflow", "jax"):
        try:
            module = importlib.import_module(module_name)
        except Exception:
            continue
        version = getattr(module, "__version__", None)
        if version:
            versions[module_name] = str(version)
    return versions


def _discover_git_commit(storage_dir: Path) -> str | None:
    try:
        result = subprocess.run(
            ["git", "-C", str(storage_dir), "rev-parse", "HEAD"],
            capture_output=True,
            check=True,
            text=True,
            timeout=5,
        )
    except Exception:
        return None
    commit = result.stdout.strip()
    return commit or None


def _capture_environment(storage_dir: Path) -> dict[str, Any]:
    return {
        "captured_at": utc_now(),
        "python": {
            "version": platform.python_version(),
            "implementation": platform.python_implementation(),
            "executable": sys.executable,
        },
        "platform": {
            "system": platform.system(),
            "release": platform.release(),
            "version": platform.version(),
            "machine": platform.machine(),
        },
        "host": {
            "hostname": socket.gethostname(),
        },
        "process": {
            "cwd": str(storage_dir),
            "argv": sys.argv,
        },
        "frameworks": _detect_frameworks(),
    }


def _normalize_tags(tags: Iterable[str] | None) -> list[str]:
    if tags is None:
        return []
    return [tag for tag in tags]


@dataclass(slots=True)
class TrainingSession:
    """Stable run context shared by SDK modules during a training session."""

    _store: LocalSessionStore
    context: RunContext
    _flushables: list[Any] = field(default_factory=list, init=False, repr=False)

    @classmethod
    def start_run(
        cls,
        *,
        project: str,
        config: dict[str, Any] | None = None,
        tags: Iterable[str] | None = None,
        parent_run_id: str | None = None,
        storage_dir: str | Path | None = None,
        strict_mode: bool = False,
    ) -> Self:
        resolved_storage_dir = Path(storage_dir or Path.cwd()).resolve()
        store = LocalSessionStore(resolved_storage_dir)
        context = store.create_run(
            project_name=project,
            config=config,
            tags=_normalize_tags(tags),
            parent_run_id=parent_run_id,
            environment=_capture_environment(resolved_storage_dir),
            git_commit=_discover_git_commit(resolved_storage_dir),
            strict_mode=strict_mode,
        )
        return cls(_store=store, context=context)

    @classmethod
    def resume_run(
        cls,
        *,
        run_id: str,
        storage_dir: str | Path | None = None,
    ) -> Self:
        resolved_storage_dir = Path(storage_dir or Path.cwd()).resolve()
        store = LocalSessionStore(resolved_storage_dir)
        context = store.resume_run(run_id)
        return cls(_store=store, context=context)

    @property
    def project_id(self) -> str:
        return self.context.project_id

    @property
    def project_name(self) -> str:
        return self.context.project_name

    @property
    def run_id(self) -> str:
        return self.context.run_id

    @property
    def run_path(self) -> str:
        return self.context.run_path

    @property
    def status(self) -> RunStatus:
        return self.context.status

    def reload(self) -> RunContext:
        self.context = self._store.load_run(self.run_id)
        return self.context

    def add_tag(self, tag: str) -> RunContext:
        self.context = self._store.add_tag(self.run_id, tag)
        return self.context

    def add_log(
        self,
        message: str,
        *,
        identifier: str | None = "Info",
        step: int | None = None,
        epoch: int | None = None,
        status: RunStatus | None = None,
        payload: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
        event_type: str = LOG_EVENT_TYPE,
    ) -> EventRecord:
        self._flush_registered_writers()
        return self._store.add_log(
            self.run_id,
            message,
            identifier=identifier,
            step=step,
            epoch=epoch,
            status=status,
            payload=dict(payload) if payload is not None else None,
            timestamp=timestamp,
            event_type=event_type,
        )

    def end_run(self, status: RunStatus = "completed") -> RunContext:
        self._flush_registered_writers()
        self.context = self._store.end_run(self.run_id, status)
        return self.context

    def to_dict(self) -> dict[str, Any]:
        return self.context.to_dict()

    def __enter__(self) -> Self:
        return self

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        if self.status != "running":
            return
        if exc is None:
            try:
                self.end_run("completed")
            except Exception:
                pass
            return

        status: RunStatus = "interrupted" if self._is_interruption(exc_type) else "failed"
        self._record_runtime_exit_log(exc_type, exc, exc_tb, status=status)
        try:
            self.end_run(status)
        except Exception:
            pass

    def _register_flushable(self, flushable: Any) -> None:
        if flushable not in self._flushables:
            self._flushables.append(flushable)

    def _unregister_flushable(self, flushable: Any) -> None:
        self._flushables = [item for item in self._flushables if item is not flushable]

    def _flush_registered_writers(self) -> None:
        active_flushables: list[Any] = []
        for flushable in self._flushables:
            if getattr(flushable, "_closed", False):
                continue
            flush = getattr(flushable, "flush", None)
            if callable(flush):
                flush()
                active_flushables.append(flushable)
        self._flushables = active_flushables

    def _record_runtime_exit_log(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        exc_tb: Any,
        *,
        status: RunStatus,
    ) -> None:
        identifier = "Warning" if status == "interrupted" else "Error"
        event_type = RUNTIME_INTERRUPTION_EVENT_TYPE if status == "interrupted" else RUNTIME_ERROR_EVENT_TYPE
        exception_name = exc_type.__name__ if exc_type is not None else type(exc).__name__
        details = str(exc).strip() if exc is not None else ""
        if status == "interrupted":
            message = f"Run interrupted by {exception_name}"
        else:
            message = f"Run failed with {exception_name}"
        if details:
            message = f"{message}: {details}"

        try:
            self.add_log(
                message,
                identifier=identifier,
                status=status,
                payload=exception_log_payload(
                    exc_type,
                    exc,
                    exc_tb,
                    source="TrainingSession.__exit__",
                ),
                event_type=event_type,
            )
        except Exception:
            pass

    @staticmethod
    def _is_interruption(exc_type: type[BaseException] | None) -> bool:
        if exc_type is None:
            return False
        return issubclass(exc_type, (KeyboardInterrupt, SystemExit))


def start_run(
    *,
    project: str,
    config: dict[str, Any] | None = None,
    tags: Iterable[str] | None = None,
    parent_run_id: str | None = None,
    storage_dir: str | Path | None = None,
    strict_mode: bool = False,
) -> TrainingSession:
    """Start a new run and persist project and run metadata locally."""

    return TrainingSession.start_run(
        project=project,
        config=config,
        tags=tags,
        parent_run_id=parent_run_id,
        storage_dir=storage_dir,
        strict_mode=strict_mode,
    )


def resume_run(*, run_id: str, storage_dir: str | Path | None = None) -> TrainingSession:
    """Resume an existing locally persisted run."""

    return TrainingSession.resume_run(run_id=run_id, storage_dir=storage_dir)


def end_run(session: TrainingSession, status: RunStatus = "completed") -> RunContext:
    """End an active run and persist its final lifecycle state."""

    return session.end_run(status=status)


def add_tag(session: TrainingSession, tag: str) -> RunContext:
    """Attach a custom tag to a run."""

    return session.add_tag(tag)


def add_log(
    session: TrainingSession,
    message: str,
    *,
    identifier: str | None = "Info",
    step: int | None = None,
    epoch: int | None = None,
    status: RunStatus | None = None,
    payload: Mapping[str, Any] | None = None,
    timestamp: str | None = None,
    event_type: str = LOG_EVENT_TYPE,
) -> EventRecord:
    """Attach a timestamped run log to event history."""

    return session.add_log(
        message,
        identifier=identifier,
        step=step,
        epoch=epoch,
        status=status,
        payload=payload,
        timestamp=timestamp,
        event_type=event_type,
    )
