from __future__ import annotations

from typing import Optional

from .._http import HttpClient
from .._environment import Environment, resolve_platform_url
from ..models.sbdd import (
    SBDDJobSubmission,
    SBDDJob,
    SBDDJobStatus,
    SBDDResults,
    SBDDMoleculeDetail,
    SBDDSummary,
)


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

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

    def submit(
        self,
        project: str,
        job_name: str,
        protein_file_name: str,
        pocket_coords: list[float],
        num_mols: int = 100,
        radius: int = 13,
        environment: Environment = None,
    ) -> SBDDJobSubmission:
        data = self._http.post("/sbdd/submit", json={
            "job_name": job_name,
            "protein_file_name": protein_file_name,
            "pocket_coords": pocket_coords,
            "num_mols": num_mols,
            "radius": radius,
        }, base_url=self._url(environment))
        return SBDDJobSubmission.from_dict(data)

    def list_jobs(
        self,
        project: str,
        job_name: Optional[str] = None,
        environment: Environment = None,
    ) -> list[SBDDJob]:
        params: dict = {}
        if job_name:
            params["job_name"] = job_name
        data = self._http.get(
            "/sbdd/jobs", params=params or None, base_url=self._url(environment),
        )
        return [SBDDJob.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,
    ) -> SBDDJobStatus:
        data = self._http.get(
            f"/sbdd/status/{job_name}/{protein_file_name}",
            base_url=self._url(environment),
        )
        return SBDDJobStatus.from_dict(data)

    def get_results(
        self,
        project: str,
        job_name: str,
        protein_file_name: str,
        denovo_job_id: Optional[str] = None,
        limit: int = 50,
        offset: int = 0,
        sort_by: str = "vina_score",
        sort_order: str = "asc",
        environment: Environment = None,
    ) -> SBDDResults:
        params: dict = {
            "limit": limit,
            "offset": offset,
            "sort_by": sort_by,
            "sort_order": sort_order,
        }
        if denovo_job_id:
            params["denovo_job_id"] = denovo_job_id
        data = self._http.get(
            f"/sbdd/results/{job_name}/{protein_file_name}",
            params=params,
            base_url=self._url(environment),
        )
        return SBDDResults.from_dict(data)

    def get_molecule(
        self,
        project: str,
        job_name: str,
        mol_id: str,
        environment: Environment = None,
    ) -> SBDDMoleculeDetail:
        data = self._http.get(
            f"/sbdd/molecule/{job_name}/{mol_id}",
            base_url=self._url(environment),
        )
        return SBDDMoleculeDetail.from_dict(data)

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