"""MLNGI Tracking API - MLflow 스타일의 고수준 추적 인터페이스

사용법:
    import mlngi

    mlngi.set_tracking_uri("http://localhost:8000/api/v1")
    mlngi.set_experiment("my-experiment")

    with mlngi.start_run(name="training-v1"):
        mlngi.log_param("lr", 0.001)
        mlngi.log_param("epochs", 100)

        for epoch in range(100):
            loss = train_one_epoch()
            mlngi.log_metric("loss", loss, step=epoch)
            mlngi.log_metric("accuracy", eval_accuracy(), step=epoch)

        mlngi.set_tag("model_type", "resnet50")
"""

import os
import time
from contextlib import contextmanager
from typing import Any

from mlngi.client import MLNGIClient

# ── 글로벌 상태 ──
_tracking_uri: str | None = None
_client: MLNGIClient | None = None
_experiment_id: str | None = None
_experiment_name: str | None = None
_project_id: str | None = None
_active_run: dict | None = None


def _get_client() -> MLNGIClient:
    global _client, _tracking_uri
    if _client is None:
        uri = _tracking_uri or os.environ.get("MLNGI_TRACKING_URI", "http://localhost:8000/api/v1")
        api_key = os.environ.get("MLNGI_API_KEY")
        token = os.environ.get("MLNGI_TOKEN")
        _client = MLNGIClient(uri, api_key=api_key, token=token)
    return _client


def set_tracking_uri(uri: str) -> None:
    """추적 서버 URI를 설정합니다.

    Args:
        uri: MLNGI 서버 API URL (예: "http://localhost:8000/api/v1")
    """
    global _tracking_uri, _client
    _tracking_uri = uri
    _client = None  # 새 URI로 클라이언트 재생성


def set_experiment(experiment_name: str, project_id: str | None = None) -> str:
    """활성 실험을 설정합니다. 없으면 자동 생성합니다.

    Args:
        experiment_name: 실험 이름
        project_id: 프로젝트 ID (생략 시 기본 프로젝트 사용)

    Returns:
        실험 ID
    """
    global _experiment_id, _experiment_name, _project_id
    client = _get_client()

    # 프로젝트 결정
    if project_id:
        _project_id = project_id
    elif _project_id is None:
        # 기본 프로젝트 찾기 또는 생성
        projects = client.list_projects()
        default = next((p for p in projects if p["name"] == "Default"), None)
        if default:
            _project_id = default["id"]
        else:
            result = client.create_project("Default", "SDK에서 자동 생성된 기본 프로젝트")
            _project_id = result["id"]

    # 실험 찾기 또는 생성
    experiments = client.list_experiments(_project_id)
    existing = next((e for e in experiments if e["name"] == experiment_name), None)
    if existing:
        _experiment_id = existing["id"]
    else:
        result = client.create_experiment(_project_id, experiment_name)
        _experiment_id = result["id"]

    _experiment_name = experiment_name
    return _experiment_id


@contextmanager
def start_run(name: str = "", tags: dict | None = None):
    """새 Run을 시작합니다. with 문으로 사용하면 자동으로 종료됩니다.

    Args:
        name: Run 이름 (선택)
        tags: 초기 태그 (선택)

    Yields:
        Run 정보 딕셔너리

    사용법:
        with mlngi.start_run(name="training-v1") as run:
            mlngi.log_metric("loss", 0.5)
    """
    global _active_run
    if _experiment_id is None:
        raise RuntimeError("실험이 설정되지 않았습니다. mlngi.set_experiment()을 먼저 호출하세요.")
    if _active_run is not None:
        raise RuntimeError("이미 실행 중인 Run이 있습니다. 먼저 end_run()을 호출하세요.")

    client = _get_client()
    run = client.create_run(_experiment_id, name=name, tags=tags)
    _active_run = run

    try:
        yield run
    except Exception:
        # 실패 시 Run 상태를 FAILED로 변경
        try:
            client.patch(f"/runs/{run['id']}", {"status": "FAILED"})
        except Exception:
            pass
        _active_run = None
        raise
    else:
        # 정상 완료 시 Run 종료
        try:
            client.stop_run(run["id"])
        except Exception:
            pass
    finally:
        _active_run = None


def end_run() -> None:
    """현재 활성 Run을 종료합니다."""
    global _active_run
    if _active_run is None:
        return
    try:
        _get_client().stop_run(_active_run["id"])
    except Exception:
        pass
    _active_run = None


def active_run() -> dict | None:
    """현재 활성 Run 정보를 반환합니다."""
    return _active_run


def get_run(run_id: str) -> dict:
    """Run 정보를 조회합니다."""
    return _get_client().get_run(run_id)


def _require_active_run() -> str:
    if _active_run is None:
        raise RuntimeError("활성 Run이 없습니다. mlngi.start_run()을 먼저 호출하세요.")
    return _active_run["id"]


# ── 메트릭 ──

def log_metric(key: str, value: float, step: int | None = None) -> None:
    """현재 Run에 메트릭을 기록합니다.

    Args:
        key: 메트릭 이름 (예: "accuracy", "loss")
        value: 메트릭 값
        step: 학습 스텝 (선택, 자동 증가)
    """
    run_id = _require_active_run()
    _get_client().log_metric(run_id, key, value, step=step or 0)


def log_metrics(metrics: dict[str, float], step: int | None = None) -> None:
    """현재 Run에 여러 메트릭을 한 번에 기록합니다.

    Args:
        metrics: {"accuracy": 0.95, "loss": 0.05} 형태
        step: 학습 스텝
    """
    run_id = _require_active_run()
    batch = [{"key": k, "value": v, "step": step or 0} for k, v in metrics.items()]
    _get_client().log_metrics_batch(run_id, batch)


# ── 파라미터 ──

def log_param(key: str, value: Any) -> None:
    """현재 Run에 파라미터를 기록합니다.

    Args:
        key: 파라미터 이름 (예: "learning_rate", "batch_size")
        value: 파라미터 값 (문자열로 변환됨)
    """
    run_id = _require_active_run()
    _get_client().log_param(run_id, key, str(value))


def log_params(params: dict[str, Any]) -> None:
    """현재 Run에 여러 파라미터를 한 번에 기록합니다.

    Args:
        params: {"lr": 0.001, "epochs": 100} 형태
    """
    run_id = _require_active_run()
    batch = [{"key": k, "value": str(v)} for k, v in params.items()]
    _get_client().log_params_batch(run_id, batch)


# ── 태그 ──

def set_tag(key: str, value: str) -> None:
    """현재 Run에 태그를 설정합니다."""
    run_id = _require_active_run()
    _get_client().set_tags(run_id, {key: value})


def set_tags(tags: dict[str, str]) -> None:
    """현재 Run에 여러 태그를 설정합니다."""
    run_id = _require_active_run()
    _get_client().set_tags(run_id, tags)


# ── 아티팩트 ──

def log_artifact(local_path: str, artifact_path: str | None = None) -> None:
    """현재 Run에 로컬 파일을 아티팩트로 업로드합니다.

    Args:
        local_path: 업로드할 로컬 파일 경로 (예: "model.pt", "./output/results.csv")
        artifact_path: 서버에 저장할 경로 (선택, 기본: 파일명 그대로)

    사용법:
        # 모델 파일 업로드
        torch.save(model, "model.pt")
        mlngi.log_artifact("model.pt")

        # 하위 폴더에 저장
        mlngi.log_artifact("config.yaml", "configs/config.yaml")
    """
    import os
    run_id = _require_active_run()
    if not os.path.exists(local_path):
        raise FileNotFoundError(f"파일을 찾을 수 없습니다: {local_path}")
    _get_client().upload_artifact(run_id, local_path, artifact_path)


def log_artifacts(local_dir: str, artifact_path: str | None = None) -> None:
    """현재 Run에 디렉토리 내 모든 파일을 아티팩트로 업로드합니다.

    Args:
        local_dir: 업로드할 로컬 디렉토리 경로
        artifact_path: 서버에 저장할 상위 경로 (선택)

    사용법:
        mlngi.log_artifacts("./output")  # output/ 내 모든 파일 업로드
    """
    import os
    run_id = _require_active_run()
    if not os.path.isdir(local_dir):
        raise NotADirectoryError(f"디렉토리를 찾을 수 없습니다: {local_dir}")
    client = _get_client()
    for root, dirs, files in os.walk(local_dir):
        for f in files:
            file_path = os.path.join(root, f)
            relative = os.path.relpath(file_path, local_dir)
            target = os.path.join(artifact_path, relative) if artifact_path else relative
            client.upload_artifact(run_id, file_path, target.replace("\\", "/"))
