# src/core_sdk/reforge_core/api.py
import os, hashlib
import time
import zipfile

from pathlib import Path
from datetime import datetime, timezone

ROBOT_MODELS_PATH = os.path.join("src", "robot", "models", "current")
ALL_MODELS_PATH = os.path.join("src", "robot", "models")

CHUNK_SIZE = 1024 * 1024  # 1 byte
DEFAULT_API_TOKEN = os.getenv("REFORGE_API_TOKEN", "dev-token-change-me")
LOCAL_HOST = os.getenv("LOCAL_HOST", "localhost")
LOCAL_PORT = int(os.getenv("LOCAL_PORT", "8000"))
PROD_HOST = os.getenv("PROD_HOST", "api.reforgerobotics.com")
REQUEST_TIMEOUT_SEC = int(os.getenv("REQUEST_TIMEOUT_SEC", "30"))

UPLOAD_TIMEOUT_SEC = REQUEST_TIMEOUT_SEC * 30
JOB_POLL_INTERVAL_SEC = float(os.getenv("JOB_POLL_INTERVAL_SEC", "10.0"))
JOB_POLL_TIMEOUT_SEC = float(os.getenv("JOB_POLL_TIMEOUT_SEC", "1200.0"))


def _env_bool(name: str, default: bool = False) -> bool:
    """Read a boolean value from an environment variable.

    Args:
        name: Environment variable name to read.
        default: Fallback value when the variable is unset.

    Returns:
        Parsed boolean value from the environment or `default` if unset.

    Side Effects:
        Reads process environment variables.

    Raises:
        None.

    Preconditions:
        `name` is a valid environment variable key.
    """
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "t", "yes", "y", "on"}


LOCAL_DEV = _env_bool("LOCAL_DEV", default=False)


class ReforgeAPIManager:
    """Coordinate calibration data upload and model retrieval.

    Args:
        None.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        None.
    """

    def __init__(
        self,
        reforge_api_token: str,
        robot_id: str,
    ) -> None:
        """Initialize the API manager with credentials and robot identity.

        Args:
            reforge_api_token: API token for cloud authentication.
            robot_id: Identifier for the robot.

        Side Effects:
            Stores credentials in the instance.

        Raises:
            None.

        Preconditions:
            `reforge_api_token` and `robot_id` are non-empty strings.
        """
        self.user_api_token = reforge_api_token
        self.robot_id = robot_id
        self._api_base_url = os.getenv(
            "REFORGE_API_BASE_URL", "https://api.reforgerobotics.com"
        )

    def _get_requests_module(self):
        """Import and return the HTTP client dependency.

        Args:
            None.

        Returns:
            The imported `requests` module.

        Side Effects:
            Imports an optional dependency at runtime.

        Raises:
            RuntimeError: If `requests` is not installed.

        Preconditions:
            None.
        """
        try:
            import requests  # type: ignore[import-untyped]
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "The 'requests' package is required for cloud API calls. Install it with `pip install requests`."
            ) from exc
        return requests

    def _auth_headers(self) -> dict[str, str]:
        """Build auth headers for outbound cloud API requests.

        Args:
            None.

        Returns:
            `dict[str, str]` HTTP headers with bearer token authorization.

        Side Effects:
            None.

        Raises:
            ValueError: If no usable token is available.

        Preconditions:
            API token is provided via constructor or environment.
        """
        # token = self.user_api_token or DEFAULT_API_TOKEN # Uncomment this when we have individual API tokens.
        token = DEFAULT_API_TOKEN
        if not token:
            raise ValueError("Missing API token.")
        return {"Authorization": f"Bearer {token}"}

    def compute_file_fingerprint(self, path: str | Path) -> tuple[str, int]:
        """Compute a SHA-256 fingerprint and size for a file.

        Args:
            path: File path to hash.

        Returns:
            `tuple[str, int]` containing the hex SHA-256 digest and file size in bytes.

        Side Effects:
            Reads file contents from disk.

        Raises:
            FileNotFoundError: If `path` does not point to an existing file.
            OSError: If the file cannot be read.

        Preconditions:
            `path` points to a readable file.
        """
        file_path = Path(path)
        if not file_path.is_file():
            raise FileNotFoundError(f"File not found: {file_path}")
        with file_path.open("rb") as file_obj:
            digest = hashlib.file_digest(file_obj, "sha256").hexdigest()
        return (digest, file_path.stat().st_size)

    def get_api_base_url(self, local_dev: bool | None = None) -> str:
        """Resolve the API base URL for local or production environments.

        Args:
            local_dev: Override for local (`True`) or production (`False`) selection.
                If `None`, uses the `LOCAL_DEV` environment-derived default.

        Returns:
            Base API URL string.

        Side Effects:
            Prints the selected base URL to stdout.

        Raises:
            None.

        Preconditions:
            Host and port environment values are valid for URL construction.
        """
        use_local = LOCAL_DEV if local_dev is None else local_dev
        if use_local:
            url = f"http://{LOCAL_HOST}:{LOCAL_PORT}"
            print(f"Using local API base URL: {url}")
            return url
        url = f"https://{PROD_HOST}"
        print(f"Using production API base URL: {url}")
        return url

    def get_url(self, endpoint: str = "", local_dev: bool | None = None) -> str:
        """Build a full API URL from a base URL and endpoint path.

        Args:
            endpoint: Relative endpoint path with or without leading slash.
            local_dev: Override for local (`True`) or production (`False`) base URL.

        Returns:
            Fully qualified URL string for the requested endpoint.

        Side Effects:
            Calls `get_api_base_url()`, which may print the selected base URL.

        Raises:
            None.

        Preconditions:
            `endpoint` is a valid URL path fragment.
        """
        base = self.get_api_base_url(local_dev=local_dev).rstrip("/")
        if not endpoint:
            return base
        normalized = endpoint if endpoint.startswith("/") else f"/{endpoint}"
        return f"{base}{normalized}"

    def _manual_fallback_hint(self, local_data_path: str | None = None) -> str:
        """Build manual fallback instructions for cloud processing failures.

        Args:
            local_data_path: Optional local path to calibration artifact for support.

        Returns:
            `str` user-facing fallback instructions.

        Side Effects:
            None.

        Raises:
            None.

        Preconditions:
            None.
        """
        hint = (
            " If this issue persists, send the calibration ZIP to "
            "calibration@reforgerobotics.com for manual processing."
        )
        if local_data_path:
            hint += f" Local data path: {os.path.abspath(local_data_path)}"
        return hint

    def _compress_data(
        self, data_folder: str, temp_folder: str = "temp_calibration_data.zip"
    ) -> str:
        """Create a zip archive of the calibration data folder.

        Args:
            data_folder: Path to the folder to compress.
            temp_folder: Zip filename to create within the folder.

        Returns:
            `str` path to the created zip file.

        Side Effects:
            Writes a zip file to disk.

        Raises:
            OSError: If files cannot be read or written.

        Preconditions:
            `data_folder` exists and is readable.
        """
        if os.path.isfile(data_folder) and zipfile.is_zipfile(data_folder):
            return data_folder

        zip_filename = os.path.join(data_folder, temp_folder)
        if zipfile.is_zipfile(zip_filename):
            return zip_filename

        with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
            for root, _, files in os.walk(data_folder):
                for file in files:
                    if file != temp_folder:
                        file_path = os.path.join(root, file)
                        arcname = os.path.relpath(file_path, start=data_folder)
                        zipf.write(file_path, arcname)
        return zip_filename

    def _save_models_locally(
        self, destination_folder: str, local_folder: str = ROBOT_MODELS_PATH
    ) -> None:
        """Save model artifacts from a cloud folder to local storage.

        Args:
            destination_folder: Folder containing model artifacts to save.
            local_folder: Folder for extracted models.

        Returns:
            `None`.

        Side Effects:
            Creates directories and writes model files to disk.

        Raises:
            OSError: If files cannot be read or written.
            ValueError: If `destination_folder` is not a zip artifact.

        Preconditions:
            `destination_folder` exists and is readable.
        """
        if os.path.isfile(destination_folder) and not zipfile.is_zipfile(
            destination_folder
        ):
            raise ValueError(
                f"Expected a compressed zip artifact, but got non-zip file: {destination_folder}"
            )

        # Unzip and save to src/robot/models/current
        with zipfile.ZipFile(destination_folder, "r") as zipf:
            extract_path = local_folder
            os.makedirs(extract_path, exist_ok=True)
            zipf.extractall(extract_path)

    def resolve_sinesweep_data(
        self,
        upload_sha256: str,
    ) -> str | None:
        """Resolve existing sine-sweep data by uploaded-file SHA-256.

        Args:
            upload_sha256: SHA-256 fingerprint of the uploaded calibration zip.

        Returns:
            `str | None` existing `data_id` if the upload already exists; otherwise `None`.

        Side Effects:
            Sends an authenticated GET request to the cloud API.
            Prints cache hit/miss and compatibility messages.

        Raises:
            requests.HTTPError: If the API returns an unexpected non-success status.
            ValueError: If the response payload is missing a valid `data_id`.

        Preconditions:
            `upload_sha256` is a valid SHA-256 hex digest string.
        """
        requests = self._get_requests_module()
        url = self.get_url(
            f"/identification/sinesweepdata/by-upload-sha/{upload_sha256}",
            local_dev=False,
        )
        response = requests.get(
            url,
            headers=self._auth_headers(),
            timeout=REQUEST_TIMEOUT_SEC,
        )

        if response.status_code in {404, 405}:
            print(
                "Data upload unavailable on this API deployment; "
                "falling back to direct upload"
            )
            return None

        response.raise_for_status()
        payload = response.json()

        # Data has not been previously uploaded
        if payload is None:
            print("[resolve_sinesweep_data] cache miss, upload required")
            return None

        # Data has been previously uploaded
        data_id = payload.get("data_id")
        if not data_id or not isinstance(data_id, str):
            raise ValueError(f"Resolve response missing valid data_id: {payload}")

        print(f"Previously uploaded data, reusing data_id={data_id}")
        return data_id

    def upload_data_to_cloud(self, zip_file_path: str) -> tuple[str, float]:
        """Upload zipped calibration data to the cloud and return the data_id.
        Reuses existing data when possible, otherwise uploads new data and returns data_id.

        Args:
            zip_file_path: Path to a `.zip` file containing calibration data.

        Returns:
            `tuple[str, float]` containing:
            - data_id returned by the cloud upload endpoint
            - upload duration in seconds

        Raises:
            FileNotFoundError: If `zip_file_path` does not exist.
            ValueError: If `zip_file_path` is not a zip file or response is invalid.
            RuntimeError: If request fails or the server returns an error.
        """
        requests = self._get_requests_module()

        if not os.path.isfile(zip_file_path):
            raise FileNotFoundError(f"Zip file not found: {zip_file_path}")
        if not zip_file_path.lower().endswith(".zip"):
            raise ValueError(f"Expected a .zip file, got: {zip_file_path}")
        if not self.robot_id:
            raise ValueError("Missing robot_id.")

        # Check if the file already exists
        upload_sha256, _ = self.compute_file_fingerprint(zip_file_path)
        request_started_at = time.time()
        resolved = self.resolve_sinesweep_data(upload_sha256)
        if resolved is not None:
            upload_duration = time.time() - request_started_at
            return resolved, upload_duration

        url = self.get_url("/identification/sinesweepdata")
        uploaded_at = datetime.now(timezone.utc).isoformat()
        manual_fallback_hint = self._manual_fallback_hint(local_data_path=zip_file_path)
        headers = self._auth_headers()
        data = {
            "robot_id": self.robot_id,
            "upload_started_at": str(uploaded_at),
        }

        print(
            f"Uploading data to Reforge Robotics cloud. This may take a few minutes depending on your connection..."
        )
        request_started_at = time.time()

        try:
            with open(zip_file_path, "rb") as file_obj:
                files = {
                    "file": (
                        os.path.basename(zip_file_path),
                        file_obj,
                        "application/zip",
                    )
                }
                response = requests.post(
                    url,
                    headers=headers,
                    data=data,
                    files=files,
                    timeout=UPLOAD_TIMEOUT_SEC,
                )
        except requests.exceptions.ConnectionError as exc:
            raise RuntimeError(
                "Could not reach Reforge cloud API. Check internet connection and API base URL."
                f" Details: {exc}.{manual_fallback_hint}"
            ) from exc
        except requests.exceptions.Timeout as exc:
            raise RuntimeError(
                f"Upload request timed out after {UPLOAD_TIMEOUT_SEC} seconds.{manual_fallback_hint}"
            ) from exc
        except Exception as exc:
            raise RuntimeError(
                f"Upload request failed: {exc}.{manual_fallback_hint}"
            ) from exc

        response.raise_for_status()
        if response.status_code != 201:
            raise RuntimeError(
                f"Upload failed with status {response.status_code}: {response.text}.{manual_fallback_hint}"
            )

        try:
            payload = response.json()
        except ValueError as exc:
            raise RuntimeError(
                f"Upload response is not valid JSON: {response.text}"
            ) from exc

        data_id = payload.get("data_id")
        upload_duration = payload.get("upload_duration_sec")
        if not data_id or not isinstance(data_id, str):
            raise ValueError(f"Upload response missing valid `data_id`: {payload}")
        if not isinstance(upload_duration, (float, int)):
            upload_duration = time.time() - request_started_at
        return data_id, upload_duration

    def _create_identification_job(
        self, data_id: str, fine_tune: bool = False, parent_job_id: str | None = None
    ) -> str:
        """Trigger cloud identification and return the created job id.

        Args:
            data_id: Identifier returned by the upload endpoint.
            fine_tune: If `True`, run fine-tuning instead of identification.
            parent_job_id: Optional job_id of the original identification job when triggering a fine-tuning job.

        Returns:
            `str` job id for the asynchronous identification request.

        Raises:
            ValueError: If the response does not contain a valid `job_id`.
            RuntimeError: If request fails, returns non-202, or JSON decoding fails.
        """
        requests = self._get_requests_module()
        # TODO: Implement fine-tuning jobs in the API
        url = self.get_url(
            endpoint=("/identification/jobs"),
        )
        headers = self._auth_headers()
        payload = {
            "data_id": data_id,
            "job_type": "identify" if not fine_tune else "fine_tune",
            "parent_job_id": parent_job_id,
        }
        try:
            response = requests.post(
                url, json=payload, headers=headers, timeout=REQUEST_TIMEOUT_SEC
            )
        except requests.exceptions.ConnectionError as exc:
            raise RuntimeError(
                f"Could not reach Reforge cloud API while creating the job. "
                f"Check internet connection.{self._manual_fallback_hint()}"
            ) from exc
        except requests.exceptions.Timeout as exc:
            raise RuntimeError(
                f"Job creation timed out after {REQUEST_TIMEOUT_SEC} seconds.{self._manual_fallback_hint()}"
            ) from exc
        if response.status_code != 201:
            raise RuntimeError(
                f"Failed to create identification job (status {response.status_code}): {response.text}.{self._manual_fallback_hint()}"
            )

        try:
            body = response.json()
        except ValueError as exc:
            raise RuntimeError(
                f"Identification job response is not valid JSON: {response.text}"
            ) from exc

        job_id = body.get("job_id")
        if not job_id or not isinstance(job_id, str):
            raise ValueError(
                f"Identification job response missing valid `job_id`: {body}"
            )
        return job_id

    def _wait_for_job_completion(
        self,
        job_id: str,
        poll_interval_seconds: float = JOB_POLL_INTERVAL_SEC,
        timeout_seconds: float = JOB_POLL_TIMEOUT_SEC,
    ) -> dict:
        """Poll `/jobs/{job_id}` until the job reaches a terminal state.

        Args:
            job_id: Identification job id returned by `_create_identification_job`.
            poll_interval_seconds: Delay between job status checks.
            timeout_seconds: Maximum total wait time before giving up.

        Returns:
            `dict` final job payload when status is `completed`.

        Raises:
            RuntimeError: If status checks fail, JSON decoding fails, or status is `failed`.
            TimeoutError: If `timeout_seconds` is exceeded before completion.
        """
        requests = self._get_requests_module()

        deadline = datetime.now(timezone.utc).timestamp() + timeout_seconds
        url = self.get_url(f"/jobs/{job_id}")
        headers = self._auth_headers()
        while datetime.now(timezone.utc).timestamp() < deadline:
            try:
                response = requests.get(
                    url, headers=headers, timeout=REQUEST_TIMEOUT_SEC
                )
            except requests.exceptions.ConnectionError as exc:
                raise RuntimeError(
                    "Lost connection to Reforge cloud API while polling job status. "
                    f"Check internet connection.{self._manual_fallback_hint()}"
                ) from exc
            except requests.exceptions.Timeout as exc:
                raise RuntimeError(
                    f"Job status request timed out after {REQUEST_TIMEOUT_SEC} seconds.{self._manual_fallback_hint()}"
                ) from exc
            if response.status_code != 200:
                raise RuntimeError(
                    f"Failed to fetch job status (status {response.status_code}): {response.text}.{self._manual_fallback_hint()}"
                )
            response.raise_for_status()
            try:
                payload = response.json()
            except ValueError as exc:
                raise RuntimeError(
                    f"Job status response is not valid JSON: {response.text}"
                ) from exc

            status = payload.get("status")
            print(
                f"[poll_job] job_id={job_id} status={status} "
                f"progress={payload.get('progress')} result_msg={payload.get('result_msg')}"
            )
            if status in {"succeeded", "failed"}:
                return payload

            time.sleep(poll_interval_seconds)

        raise TimeoutError(
            f"Timed out waiting for identification job {job_id} after {timeout_seconds} seconds. Please email calibration@reforgerobotics.com to manually process your data."
        )

    def get_latest_model_by_robot_id(self) -> dict:
        """Fetch metadata for the latest model artifact for this robot.

        Args:
            None.

        Returns:
            `dict` JSON payload describing the latest model-producing job.

        Side Effects:
            Sends an authenticated GET request to the cloud API.

        Raises:
            requests.HTTPError: If the API returns a non-success status.
            ValueError: If the response payload is not a JSON object.

        Preconditions:
            `self.robot_id` is set to a valid robot identifier.
        """
        requests = self._get_requests_module()
        url = self.get_url(f"/identification/robots/{self.robot_id}/models/latest")
        headers = self._auth_headers()
        response = requests.get(
            url,
            headers=headers,
            timeout=REQUEST_TIMEOUT_SEC,
        )

        response.raise_for_status()
        payload = response.json()
        if not isinstance(payload, dict):
            raise ValueError(f"Expected object payload for latest model: {payload}")
        return payload

    def _download_file(self, job_id: str, destination_path: str) -> str:
        """Download a file from `/download/{filename}`.

        Args:
            job_id: Remote job_id returned by the completed job.
            destination_path: Local output path. Can be a directory or a zip filepath.

        Returns:
            `str` full local path to the downloaded file.

        Raises:
            RuntimeError: If the download request fails or returns a non-200 status.
        """
        requests = self._get_requests_module()

        url = self.get_url(f"/jobs/{job_id}/models")
        headers = self._auth_headers()
        default_filename = f"downloaded_models.zip"

        # Accept either a directory or an explicit .zip filepath.
        if os.path.isdir(destination_path):
            if destination_path.lower().endswith(".zip"):
                parent_dir = os.path.dirname(destination_path) or "."
                output_path = os.path.join(parent_dir, default_filename)
            else:
                output_path = os.path.join(destination_path, default_filename)
        elif destination_path.lower().endswith(".zip"):
            output_path = destination_path
        else:
            os.makedirs(destination_path, exist_ok=True)
            output_path = os.path.join(destination_path, default_filename)

        output_dir = os.path.dirname(output_path) or "."
        os.makedirs(output_dir, exist_ok=True)
        manual_fallback_hint = self._manual_fallback_hint(local_data_path=output_path)
        tic = time.time()
        try:
            with requests.get(
                url, headers=headers, stream=True, timeout=UPLOAD_TIMEOUT_SEC
            ) as response:
                response.raise_for_status()
                if response.status_code != 200:
                    raise RuntimeError(
                        f"Download failed for '{job_id}' (status {response.status_code}): {response.text}.{manual_fallback_hint}"
                    )

                with open(output_path, "wb") as downloaded_file:
                    for chunk in response.iter_content(chunk_size=CHUNK_SIZE):
                        if chunk:
                            downloaded_file.write(chunk)
        except requests.exceptions.ConnectionError as exc:
            raise RuntimeError(
                "Could not reach Reforge cloud API while downloading models. "
                "Check internet connection and API base URL."
                f"{manual_fallback_hint}"
            ) from exc
        except requests.exceptions.Timeout as exc:
            raise RuntimeError(
                f"Download request timed out after {UPLOAD_TIMEOUT_SEC} seconds.{manual_fallback_hint}"
            ) from exc

        duration = time.time() - tic
        print(f"Downloaded files in {duration: .2f} seconds.")
        return output_path

    def has_local_models(self, models_path: str = ROBOT_MODELS_PATH) -> bool:
        """Check whether local model files are available for fine-tuning.

        Args:
            models_path: Folder expected to contain `.pt` model artifacts.

        Returns:
            `bool` indicating whether at least one `.pt` file is present.

        Side Effects:
            Reads directory metadata and file names from local storage.

        Raises:
            OSError: If `models_path` exists but cannot be listed.

        Preconditions:
            `models_path` is a valid filesystem path.
        """
        return os.path.exists(models_path) and any(
            fname.endswith(".pt") for fname in os.listdir(models_path)
        )

    def run_cloud_model_generation(
        self,
        data_folder: str,
        delete_temp_file: bool = False,
        fine_tune: bool = False,
    ) -> None:
        """Run cloud identification/fine-tuning and save returned models locally.

        Args:
            data_folder: Path to calibration data folder (or zip) to upload.
            delete_temp_file: If `True`, delete temporary zip created locally after upload.
            fine_tune: If `True`, trigger fine-tuning job; otherwise trigger identification.

        Returns:
            `None`.

        Side Effects:
            Uploads data to cloud, creates and polls a remote job, downloads model artifact,
            and writes extracted models to local storage.

        Raises:
            FileNotFoundError: If upload input cannot be found.
            ValueError: If cloud payloads are missing required fields.
            RuntimeError: If upload/job/download requests fail.
            TimeoutError: If job does not complete within timeout.
            OSError: If local file operations fail.

        Preconditions:
            `data_folder` exists and contains valid calibration data.
        """
        # Upload compressed data to cloud using self.user_api_token and self.robot_id (have it decompress in the cloud)
        temp_zip_path = self._compress_data(data_folder)
        try:
            data_id, upload_duration = self.upload_data_to_cloud(temp_zip_path)
        finally:
            if delete_temp_file and os.path.exists(temp_zip_path):
                os.remove(temp_zip_path)
        job_kind = "fine-tuning" if fine_tune else "identification"
        print(
            f"Data uploaded in {upload_duration:.2f} seconds. Running {job_kind}... this may take 10-15 minutes depending on server load."
        )

        # Create job to run model generation in the cloud and wait for completion.
        job_id = self._create_identification_job(data_id, fine_tune=fine_tune)
        print(f"Created {job_kind} job {job_id}. Waiting for completion...")

        job_payload = self._wait_for_job_completion(job_id)
        if job_payload.get("status") != "succeeded":
            raise RuntimeError(
                f"Job did not succeed. status={job_payload.get('status')}"
            )

        # Get latest model by robot id
        latest_model = self.get_latest_model_by_robot_id()
        print(
            f"latest model job_id={latest_model.get('job_id')} "
            f"download_path={latest_model.get('download_path')}"
        )

        # Save model in the robot models directory,
        # Appending the date/original name of the data file
        artifact_dir = os.path.join(
            ALL_MODELS_PATH, data_folder.split("/")[-1].split(".")[0]
        )
        downloaded_model_zip = self._download_file(job_id, artifact_dir)
        print(
            f"Downloaded model artifact to {downloaded_model_zip}. Extracting to {ROBOT_MODELS_PATH}..."
        )

        # The API returns a compressed model artifact; reuse local save helper to unzip into current models path.
        self._save_models_locally(downloaded_model_zip)
