"""API client wrapping all Nbench REST endpoints."""

import requests


class NbenchClient:
    """HTTP client for the Nbench API."""

    def __init__(self, api_url: str = "http://localhost:8080", api_key: str | None = None):
        self.api_url = api_url.rstrip("/")
        self.session = requests.Session()
        if api_key:
            self.session.headers["Authorization"] = f"Bearer {api_key}"

    # ------------------------------------------------------------------
    # Project endpoints
    # ------------------------------------------------------------------

    def get_project(self, project_id: str) -> dict:
        """GET /v1/projects/:id"""
        resp = self.session.get(f"{self.api_url}/v1/projects/{project_id}")
        resp.raise_for_status()
        return resp.json()

    def get_experiments(self, project_id: str, limit: int = 50) -> list[dict]:
        """GET /v1/projects/:id/experiments"""
        resp = self.session.get(
            f"{self.api_url}/v1/projects/{project_id}/experiments",
            params={"limit": limit},
        )
        resp.raise_for_status()
        return resp.json()

    def get_categories(self, project_id: str) -> list[dict]:
        """GET /v1/projects/:id/categories"""
        resp = self.session.get(
            f"{self.api_url}/v1/projects/{project_id}/categories"
        )
        resp.raise_for_status()
        return resp.json()

    def get_queue(self, project_id: str) -> list[dict]:
        """GET /v1/projects/:id/queue (requires auth)"""
        resp = self.session.get(
            f"{self.api_url}/v1/projects/{project_id}/queue"
        )
        resp.raise_for_status()
        return resp.json()

    def get_fee(self, project_id: str) -> dict:
        """GET /v1/projects/:id/fee"""
        resp = self.session.get(
            f"{self.api_url}/v1/projects/{project_id}/fee"
        )
        resp.raise_for_status()
        return resp.json()

    def get_code(self, project_id: str) -> str:
        """GET /v1/projects/:id/code"""
        resp = self.session.get(
            f"{self.api_url}/v1/projects/{project_id}/code"
        )
        resp.raise_for_status()
        return resp.text

    def deactivate_project(self, project_id: str) -> None:
        """POST /v1/projects/:id/deactivate"""
        resp = self.session.post(f"{self.api_url}/v1/projects/{project_id}/deactivate")
        resp.raise_for_status()

    def create_project(
        self,
        name: str,
        description: str,
        baseline_val_bpb: float,
        hardware: str = "unspecified",
        repo_url: str | None = None,
        branch_name: str | None = None,
    ) -> dict:
        """POST /v1/projects"""
        data = {
            "name": name,
            "description": description,
            "baseline_val_bpb": baseline_val_bpb,
            "hardware": hardware,
        }
        if repo_url:
            data["repo_url"] = repo_url
        if branch_name:
            data["branch_name"] = branch_name
        resp = self.session.post(f"{self.api_url}/v1/projects", json=data)
        resp.raise_for_status()
        return resp.json()

    def get_account(self) -> dict:
        """GET /v1/account"""
        resp = self.session.get(f"{self.api_url}/v1/account")
        resp.raise_for_status()
        return resp.json()

    # ------------------------------------------------------------------
    # Suggestion submission
    # ------------------------------------------------------------------

    def submit_suggestion(
        self,
        project_id: str,
        prompt_text: str,
        submitter_type: str = "agent",
    ) -> dict:
        """POST /v1/projects/:id/suggestions"""
        resp = self.session.post(
            f"{self.api_url}/v1/projects/{project_id}/suggestions",
            json={"prompt_text": prompt_text, "submitter_type": submitter_type},
        )
        resp.raise_for_status()
        return resp.json()

    def update_suggestion_status(
        self, project_id: str, suggestion_id: str, status: str
    ) -> None:
        """PATCH /v1/projects/:id/suggestions/:sid"""
        resp = self.session.patch(
            f"{self.api_url}/v1/projects/{project_id}/suggestions/{suggestion_id}",
            json={"status": status},
        )
        resp.raise_for_status()

    # ------------------------------------------------------------------
    # Branch endpoints
    # ------------------------------------------------------------------

    def register_branch(self, project_id, branch_name, description="", claimed_metric=None):
        """POST /v1/projects/:id/branches"""
        data = {"branch_name": branch_name, "description": description}
        if claimed_metric is not None:
            data["claimed_metric"] = claimed_metric
        resp = self.session.post(f"{self.api_url}/v1/projects/{project_id}/branches", json=data)
        resp.raise_for_status()
        return resp.json()

    def get_branches(self, project_id):
        """GET /v1/projects/:id/branches"""
        resp = self.session.get(f"{self.api_url}/v1/projects/{project_id}/branches")
        resp.raise_for_status()
        return resp.json()

    def update_branch(self, project_id, branch_id, status=None, verified_metric=None, error_log=None):
        """PATCH /v1/projects/:id/branches/:bid"""
        data = {}
        if status: data["status"] = status
        if verified_metric is not None: data["verified_metric"] = verified_metric
        if error_log: data["error_log"] = error_log
        resp = self.session.patch(f"{self.api_url}/v1/projects/{project_id}/branches/{branch_id}", json=data)
        resp.raise_for_status()
        return resp.json()
