"""
Experiments API
"""

from typing import Optional, Dict, Any, List
from .base import BaseAPI


class ExperimentsAPI(BaseAPI):
    """Experiment management endpoints"""
    
    def list(
        self,
        page: int = 1,
        page_size: int = 20,
    ) -> Dict[str, Any]:
        """List experiments"""
        params = {"page": page, "page_size": min(page_size, 100)}
        return self.client._request("GET", "experiments/", params=params)
    
    def get(self, experiment_id: int) -> Dict[str, Any]:
        """Get experiment by ID"""
        return self.client._request("GET", f"experiments/{experiment_id}")
    
    def create(
        self,
        name: str,
        description: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Create a new experiment"""
        data = {"name": name}
        if description:
            data["description"] = description
        return self.client._request("POST", "experiments/", json=data)
    
    def update(
        self,
        experiment_id: int,
        name: Optional[str] = None,
        description: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Update experiment"""
        data = {}
        if name:
            data["name"] = name
        if description is not None:
            data["description"] = description
        return self.client._request(
            "PUT",
            f"experiments/{experiment_id}",
            json=data
        )
    
    def delete(self, experiment_id: int) -> None:
        """Delete experiment"""
        self.client._request("DELETE", f"experiments/{experiment_id}")
    
    # Experiment Runs
    def list_runs(
        self,
        experiment_id: int,
        page: int = 1,
        page_size: int = 20,
    ) -> Dict[str, Any]:
        """List runs của một experiment. GET experiments/{experiment_id}/runs"""
        params = {"page": page, "page_size": min(page_size, 100)}
        return self.client._request(
            "GET",
            f"experiments/{experiment_id}/runs",
            params=params
        )

    def list_runs_global(
        self,
        experiment_id: Optional[int] = None,
        model_id: Optional[int] = None,
        page: int = 1,
        page_size: int = 20,
    ) -> list:
        """
        List tất cả experiment runs (có lọc). GET /experiments/runs.
        Returns list of runs (không paginated wrapper).
        """
        params = {"page": page, "page_size": min(page_size, 100)}
        if experiment_id is not None:
            params["experiment_id"] = experiment_id
        if model_id is not None:
            params["model_id"] = model_id
        return self.client._request("GET", "experiments/runs", params=params)
    
    def get_run(
        self,
        experiment_id: int,
        run_id: int,
    ) -> Dict[str, Any]:
        """Get experiment run by ID. GET /experiments/runs/{run_id} (backend không dùng experiment_id trong path)."""
        return self.client._request("GET", f"experiments/runs/{run_id}")
    
    def create_run(
        self,
        experiment_id: int,
        model_id: int,
        dataset_split_id: int,
        config: Dict[str, Any],
        git_commit: str = "manual-training",
    ) -> Dict[str, Any]:
        """Create a new experiment run (triggers Celery background job). Có thể fail nếu Celery/Redis/MLFlow lỗi."""
        data = {
            "model_id": model_id,
            "dataset_split_id": dataset_split_id,
            "config": config,
            "git_commit": git_commit,
        }
        return self.client._request(
            "POST",
            f"experiments/{experiment_id}/runs",
            json=data
        )

    def create_run_sync(
        self,
        experiment_id: int,
        model_id: int,
        dataset_split_id: int,
        config: Dict[str, Any],
        git_commit: str = "manual-training",
    ) -> Dict[str, Any]:
        """
        Tạo experiment run đồng bộ (POST /experiments/runs).
        Dùng khi job create_run bị failed (Celery/Redis/MLFlow) nhưng vẫn cần run_id để lưu metrics.
        Trả về run với experiment_runs_id (hoặc id).
        """
        data = {
            "experiment_id": experiment_id,
            "model_id": model_id,
            "dataset_split_id": dataset_split_id,
            "config": config,
            "git_commit": git_commit,
        }
        return self.client._request("POST", "experiments/runs", json=data)

    def update_run(
        self,
        experiment_id: int,
        run_id: int,
        name: Optional[str] = None,
        description: Optional[str] = None,
        hyperparameters: Optional[Dict] = None,
    ) -> Dict[str, Any]:
        """Update experiment run"""
        data = {}
        if name:
            data["name"] = name
        if description is not None:
            data["description"] = description
        if hyperparameters is not None:
            data["hyperparameters"] = hyperparameters
        return self.client._request(
            "PUT",
            f"experiments/{experiment_id}/runs/{run_id}",
            json=data
        )
    
    def delete_run(
        self,
        experiment_id: int,
        run_id: int,
    ) -> None:
        """Delete experiment run"""
        self.client._request(
            "DELETE",
            f"experiments/{experiment_id}/runs/{run_id}"
        )
