from __future__ import annotations

from typing import Optional

from .._http import HttpClient
from .._environment import Environment, resolve_platform_url
from ..models.docking import (
    PocketDetectionResult,
    DockingJobSubmission,
    DockingJob,
    DockingJobStatus,
    DockingResult,
    DockingSummary,
)


class DockingNamespace:
    def __init__(self, http: HttpClient):
        self._http = http

    def _url(self, environment: Environment) -> str | None:
        return resolve_platform_url(environment)

    # ── pocket detection (one-shot) ─────────────────────────────

    def find_pockets(
        self,
        project: str,
        protein_file_name: str,
        environment: Environment = None,
    ) -> PocketDetectionResult:
        data = self._http.post("/docking/pockets/find", json={
            "job_name": project,
            "protein_file_name": protein_file_name,
        }, base_url=self._url(environment))
        return PocketDetectionResult.from_dict(data)

    # ── docking job lifecycle ───────────────────────────────────

    def submit(
        self,
        project: str,
        job_name: str,
        protein_file_name: str,
        ligand_file_names: list[str],
        pocket_coords: list[float],
        box_sizes: list[float],
        num_poses: int = 10,
        exhaustiveness: int = 128,
        energy_range: float = 3.0,
        scoring_function: str = "vina",
        environment: Environment = None,
    ) -> DockingJobSubmission:
        data = self._http.post("/docking/submit", json={
            "job_name": job_name,
            "protein_file_name": protein_file_name,
            "ligand_file_names": ligand_file_names,
            "pocket_coords": pocket_coords,
            "box_sizes": box_sizes,
            "num_poses": num_poses,
            "exhaustiveness": exhaustiveness,
            "energy_range": energy_range,
            "scoring_function": scoring_function,
        }, base_url=self._url(environment))
        return DockingJobSubmission.from_dict(data)

    def list_jobs(
        self,
        project: str,
        job_name: Optional[str] = None,
        environment: Environment = None,
    ) -> list[DockingJob]:
        params: dict = {}
        if job_name:
            params["job_name"] = job_name
        data = self._http.get(
            "/docking/jobs", params=params or None, base_url=self._url(environment),
        )
        return [DockingJob.from_dict(j) for j in data.get("jobs", [])]

    def get_status(
        self,
        project: str,
        job_name: str,
        protein_file_name: str,
        environment: Environment = None,
    ) -> DockingJobStatus:
        data = self._http.get(
            f"/docking/status/{job_name}/{protein_file_name}",
            base_url=self._url(environment),
        )
        return DockingJobStatus.from_dict(data)

    def get_result(
        self,
        project: str,
        job_name: str,
        protein_file_name: str,
        ligand_name: str,
        pose_id: Optional[int] = None,
        environment: Environment = None,
    ) -> DockingResult:
        params = {"pose_id": pose_id} if pose_id is not None else None
        data = self._http.get(
            f"/docking/result/{job_name}/{protein_file_name}/{ligand_name}",
            params=params,
            base_url=self._url(environment),
        )
        return DockingResult.from_dict(data)

    def get_summary(
        self,
        project: str,
        job_name: str,
        protein_file_name: str,
        environment: Environment = None,
    ) -> DockingSummary:
        data = self._http.get(
            f"/docking/summary/{job_name}/{protein_file_name}",
            base_url=self._url(environment),
        )
        return DockingSummary.from_dict(data)
