from __future__ import annotations

import importlib
from dataclasses import dataclass
from pathlib import Path


PATCH_MARKER = "AIRFRANS_SELECTED_OFFER_ID_PATCH"
SDK_API_KEY_MARKER = "AIRFRANS_VAST_SDK_API_KEY_PATCH"


@dataclass(frozen=True)
class PatchStatus:
    path: Path | None
    installed: bool
    patched: bool
    message: str


def locate_vast_utils() -> Path:
    try:
        module = importlib.import_module("sky.provision.vast.utils")
    except ModuleNotFoundError as exc:
        raise RuntimeError("SkyPilot is not importable in this Python environment") from exc
    path = getattr(module, "__file__", None)
    if not path:
        raise RuntimeError("Could not locate sky.provision.vast.utils source file")
    return Path(path)


def patch_status() -> PatchStatus:
    try:
        path = locate_vast_utils()
    except RuntimeError as exc:
        return PatchStatus(path=None, installed=False, patched=False, message=str(exc))
    text = path.read_text()
    patched = PATCH_MARKER in text and SDK_API_KEY_MARKER in text
    if patched:
        return PatchStatus(path=path, installed=True, patched=True, message="SkyPilot Vast selected-offer patch is installed")
    return PatchStatus(path=path, installed=True, patched=False, message="SkyPilot Vast selected-offer patch is missing")


def apply_patch() -> PatchStatus:
    path = locate_vast_utils()
    text = path.read_text()
    already_selected = PATCH_MARKER in text
    already_api_key = SDK_API_KEY_MARKER in text
    if already_selected and already_api_key:
        return PatchStatus(path=path, installed=True, patched=True, message="SkyPilot Vast selected-offer patch already installed")

    if not already_selected:
        old = '''    instance_list = vast.vast().search_offers(query=query_str)\n\n    if isinstance(instance_list, int) or len(instance_list) == 0:\n        raise RuntimeError('Failed to create instances, could not find an '\n                           'offer that satisfies the requirements '\n                           f'"{query_str}".')\n\n    instance_touse = instance_list[0]\n\n    # Start with user-provided kwargs as the base\n    launch_params: Dict[str, Any] = dict(create_instance_kwargs or {})\n    # Remove None values to avoid overriding defaults\n    launch_params = {k: v for k, v in launch_params.items() if v is not None}\n'''
        new = f'''    # {PATCH_MARKER}: allow callers to bypass SkyPilot's Vast offer\n    # search after selecting a vetted Vast offer themselves. This keeps\n    # SkyPilot's create/wait/bootstrap/run lifecycle while preserving\n    # external marketplace-quality filtering.\n    launch_params: Dict[str, Any] = dict(create_instance_kwargs or {{}})\n    launch_params = {{k: v for k, v in launch_params.items() if v is not None}}\n    selected_offer_id = launch_params.pop('selected_offer_id', None)\n\n    if selected_offer_id is not None:\n        logger.info(f'Using externally selected Vast offer {{selected_offer_id}}.')\n        try:\n            instance_touse = {{'id': int(selected_offer_id)}}\n        except (TypeError, ValueError) as e:\n            raise RuntimeError(\n                f'Invalid selected_offer_id for Vast: {{selected_offer_id!r}}'\n            ) from e\n    else:\n        instance_list = vast.vast().search_offers(query=query_str)\n\n        if isinstance(instance_list, int) or len(instance_list) == 0:\n            raise RuntimeError('Failed to create instances, could not find an '\n                               'offer that satisfies the requirements '\n                               f'"{{query_str}}".')\n\n        instance_touse = instance_list[0]\n\n'''
        if old not in text:
            raise RuntimeError(
                f"Could not apply SkyPilot Vast selected-offer patch; expected source block not found in {path}"
            )
        text = text.replace(old, new, 1)

    if not already_api_key:
        old_key = "        f'echo \"{vast.vast().client.api_key}\" > ~/.vast_api_key',"
        new_key = (
            f"        # {SDK_API_KEY_MARKER}: vastai-sdk exposes api_key directly in current releases.\n"
            "        f'echo \"{vast.vast().api_key}\" > ~/.vast_api_key',"
        )
        if old_key not in text:
            raise RuntimeError(
                f"Could not apply SkyPilot Vast API-key patch; expected source line not found in {path}"
            )
        text = text.replace(old_key, new_key, 1)

    path.write_text(text)
    return PatchStatus(path=path, installed=True, patched=True, message="SkyPilot Vast selected-offer patch installed")


def require_patch() -> Path:
    status = patch_status()
    if not status.installed:
        raise RuntimeError(status.message)
    if not status.patched:
        raise RuntimeError(f"{status.message}; run `remote-run doctor --apply-skypilot-patch`")
    assert status.path is not None
    return status.path