from __future__ import annotations

from typing import Optional

from .._http import HttpClient
from .._environment import Environment, resolve_platform_url
from ..models.design import (
    DesignValidation,
    DesignJobSubmission,
    DesignJob,
    DesignJobStatus,
    DesignContent,
)


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

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

    def validate(
        self,
        project: str,
        protein_name: str,
        config_yaml_content: str,
        ligand_name: Optional[str] = None,
        environment: Environment = None,
    ) -> DesignValidation:
        payload: dict = {
            "job_name": project,
            "protein_name": protein_name,
            "config_yaml_content": config_yaml_content,
        }
        if ligand_name:
            payload["ligand_name"] = ligand_name
        data = self._http.post(
            "/design/validate", json=payload, base_url=self._url(environment),
        )
        return DesignValidation.from_dict(data)

    def submit(
        self,
        project: str,
        job_name: str,
        design_model_name: str,
        design_config_name: str,
        protein_name: str,
        config_yaml_content: str,
        ligand_name: Optional[str] = None,
        num_designs: Optional[int] = None,
        inverse_fold_num_sequences: Optional[int] = None,
        budget: Optional[int] = None,
        runner_config: Optional[dict] = None,
        environment: Environment = None,
    ) -> DesignJobSubmission:
        payload: dict = {
            "job_name": job_name,
            "design_model_name": design_model_name,
            "design_config_name": design_config_name,
            "protein_name": protein_name,
            "config_yaml_content": config_yaml_content,
        }
        if ligand_name is not None:
            payload["ligand_name"] = ligand_name
        if num_designs is not None:
            payload["num_designs"] = num_designs
        if inverse_fold_num_sequences is not None:
            payload["inverse_fold_num_sequences"] = inverse_fold_num_sequences
        if budget is not None:
            payload["budget"] = budget
        if runner_config is not None:
            payload["runner_config"] = runner_config
        data = self._http.post(
            "/design/submit", json=payload, base_url=self._url(environment),
        )
        return DesignJobSubmission.from_dict(data)

    def validate_and_submit(
        self,
        project: str,
        job_name: str,
        design_model_name: str,
        design_config_name: str,
        protein_name: str,
        config_yaml_content: str,
        ligand_name: Optional[str] = None,
        num_designs: Optional[int] = None,
        inverse_fold_num_sequences: Optional[int] = None,
        budget: Optional[int] = None,
        runner_config: Optional[dict] = None,
        environment: Environment = None,
    ) -> DesignJobSubmission:
        validation = self.validate(
            project=project,
            protein_name=protein_name,
            config_yaml_content=config_yaml_content,
            ligand_name=ligand_name,
            environment=environment,
        )
        if not validation.valid:
            return DesignJobSubmission(
                success=False,
                message=f"Validation failed: {validation.message}",
            )
        return self.submit(
            project=project,
            job_name=job_name,
            design_model_name=design_model_name,
            design_config_name=design_config_name,
            protein_name=protein_name,
            config_yaml_content=config_yaml_content,
            ligand_name=ligand_name,
            num_designs=num_designs,
            inverse_fold_num_sequences=inverse_fold_num_sequences,
            budget=budget,
            runner_config=runner_config,
            environment=environment,
        )

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

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

    def get_content(
        self,
        project: str,
        job_name: str,
        design_model_name: str,
        design_config_name: str,
        environment: Environment = None,
    ) -> DesignContent:
        data = self._http.get(
            f"/design/content/{job_name}/{design_model_name}/{design_config_name}",
            base_url=self._url(environment),
        )
        return DesignContent.from_dict(data)
