from __future__ import annotations

from typing import Optional

from .._http import HttpClient
from ..models.dynamo import (
    ActionResponse,
    ExperimentDetail,
    ExperimentResponse,
    ExperimentSummary,
    ResourceResponse,
    SimulationResults,
    SimulationStatus,
    SimulationSummary,
    StartSimulationResponse,
    TopologyResponse,
)


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

    # ── Experiments ───────────────────────────────────────────────

    def create_experiment(
        self,
        job_name: str,
        experiment_name: str,
        description: Optional[str] = None,
    ) -> ExperimentResponse:
        data = self._http.post("/md/experiments", json={
            "job_name": job_name,
            "experiment_name": experiment_name,
            "description": description,
        })
        return ExperimentResponse.from_dict(data)

    def list_experiments(
        self, job_name: Optional[str] = None,
    ) -> list[ExperimentSummary]:
        params = {"job_name": job_name} if job_name else None
        data = self._http.get("/md/experiments", params=params)
        return [ExperimentSummary.from_dict(e) for e in data.get("experiments") or []]

    def get_experiment(self, experiment_id: str) -> ExperimentDetail:
        data = self._http.get(f"/md/experiments/{experiment_id}")
        return ExperimentDetail.from_dict(data)

    def update_experiment(
        self,
        experiment_id: str,
        experiment_name: Optional[str] = None,
        description: Optional[str] = None,
    ) -> ActionResponse:
        payload: dict = {}
        if experiment_name is not None:
            payload["experiment_name"] = experiment_name
        if description is not None:
            payload["description"] = description
        data = self._http.patch(f"/md/experiments/{experiment_id}", json=payload)
        return ActionResponse.from_dict(data)

    # ── Simulation lifecycle ─────────────────────────────────────

    def start_simulation(
        self,
        job_name: str,
        experiment_name: str,
        protein_file_name: str,
        *,
        components: Optional[list[dict]] = None,
        sim_type: str = "md",
        forcefields: Optional[list[str]] = None,
        water_model: str = "tip3p",
        padding: float = 1.5,
        ionic_strength: float = 0.15,
        protein_pH: float = 7.2,
        temperature: float = 300.0,
        pressure: float = 1.0,
        timestep: float = 4.0,
        total_time: float = 10.0,
        total_time_unit: str = "nanosecond",
        step_size: int = 1000,
        platform: str = "CUDA",
        precision: str = "mixed",
        log_frequency: int = 250,
        chkpt_frequency: int = 2500,
        remove_heterogens: bool = True,
        hydrogen_mass: float = 2.0,
        has_membrane: bool = False,
        membrane_lipids: str = "POPC",
        membrane_preoriented: bool = False,
        membrane_system: str = "PMf",
        membrane_topology: str = "out",
        membrane_bioassembly: bool = False,
        membrane_ihetero: int = 1,
        resume_from_md_id: Optional[int] = None,
    ) -> StartSimulationResponse:
        payload: dict = {
            "job_name": job_name,
            "experiment_name": experiment_name,
            "protein_file_name": protein_file_name,
            "components": components or [],
            "sim_type": sim_type,
            "forcefields": forcefields or [
                "amber/protein.ff14SB.xml",
                "amber/tip3p_standard.xml",
                "amber/tip3p_HFE_multivalent.xml",
            ],
            "water_model": water_model,
            "padding": padding,
            "ionic_strength": ionic_strength,
            "protein_pH": protein_pH,
            "temperature": temperature,
            "pressure": pressure,
            "timestep": timestep,
            "total_time": total_time,
            "total_time_unit": total_time_unit,
            "step_size": step_size,
            "platform": platform,
            "precision": precision,
            "log_frequency": log_frequency,
            "chkpt_frequency": chkpt_frequency,
            "remove_heterogens": remove_heterogens,
            "hydrogen_mass": hydrogen_mass,
            "has_membrane": has_membrane,
            "membrane_lipids": membrane_lipids,
            "membrane_preoriented": membrane_preoriented,
            "membrane_system": membrane_system,
            "membrane_topology": membrane_topology,
            "membrane_bioassembly": membrane_bioassembly,
            "membrane_ihetero": membrane_ihetero,
        }
        if resume_from_md_id is not None:
            payload["resume_from_md_id"] = resume_from_md_id
        data = self._http.post("/md/start", json=payload)
        return StartSimulationResponse.from_dict(data)

    def restart_simulation(self, md_id: int) -> ActionResponse:
        data = self._http.post(f"/md/restart/{md_id}")
        return ActionResponse.from_dict(data)

    def continue_simulation(
        self,
        md_id: int,
        total_time: Optional[float] = None,
        total_time_unit: str = "nanosecond",
    ) -> ActionResponse:
        payload: dict = {"total_time_unit": total_time_unit}
        if total_time is not None:
            payload["total_time"] = total_time
        data = self._http.post(f"/md/continue/{md_id}", json=payload)
        return ActionResponse.from_dict(data)

    def reanalyze_simulation(self, md_id: int) -> ActionResponse:
        data = self._http.post(f"/md/reanalyze/{md_id}")
        return ActionResponse.from_dict(data)

    def cancel_simulation(self, md_id: int) -> ActionResponse:
        data = self._http.post(f"/md/cancel/{md_id}")
        return ActionResponse.from_dict(data)

    # ── Status / Results ─────────────────────────────────────────

    def get_status(self, experiment_id: str) -> SimulationStatus:
        data = self._http.get(f"/md/status/{experiment_id}")
        return SimulationStatus.from_dict(data)

    def list_simulations(self, experiment_id: str) -> list[SimulationSummary]:
        data = self._http.get(f"/md/simulations/{experiment_id}")
        return [SimulationSummary.from_dict(s) for s in data.get("simulations") or []]

    def get_results(self, experiment_id: str) -> SimulationResults:
        data = self._http.get(f"/md/results/{experiment_id}")
        return SimulationResults.from_dict(data)

    def get_topology(self, experiment_id: str) -> TopologyResponse:
        data = self._http.get(f"/md/topology/{experiment_id}")
        return TopologyResponse.from_dict(data)

    def get_resource(self, experiment_id: str, filename: str) -> ResourceResponse:
        data = self._http.get(
            f"/md/resource/{experiment_id}", params={"filename": filename},
        )
        return ResourceResponse.from_dict(data)

    def get_trajectory_stream_url(
        self,
        experiment_id: str,
        skip: int = 1,
        batch_size: int = 15,
        start_frame: int = 0,
        end_frame: Optional[int] = None,
    ) -> str:
        base = self._http._client.base_url
        url = f"{base}/md/trajectory/stream/{experiment_id}?skip={skip}&batch_size={batch_size}&start_frame={start_frame}"
        if end_frame is not None:
            url += f"&end_frame={end_frame}"
        return url
