from ...core.main import Step
from ...core.sdk import ContinueSDK
from ...libs.util.paths import getConfigFilePath
from ...models.filesystem import RangeInFile
from ...models.main import Range

MODEL_CLASS_TO_MESSAGE = {
    "Ollama": "To get started with Ollama, download the app from [ollama.ai](https://ollama.ai/). Once it is downloaded, be sure to pull at least one model and use its name in the model field in config.py (e.g. `model='codellama'`).",
    "LlamaCpp": "To get started with this model, clone the [`llama.cpp` repo](https://github.com/ggerganov/llama.cpp) and follow the instructions to set up the server [here](https://github.com/ggerganov/llama.cpp/blob/master/examples/server/README.md#build). Any of the parameters described in the README can be passed to the `llama_cpp_args` field in the `LlamaCpp` class in `config.py`."
}


class SetupModelStep(Step):
    model_class: str
    name: str = "Setup model in config.py"

    async def run(self, sdk: ContinueSDK):
        await sdk.ide.setFileOpen(getConfigFilePath())
        self.description = MODEL_CLASS_TO_MESSAGE.get(
            self.model_class, "Please finish setting up this model in `config.py`"
        )

        config_contents = await sdk.ide.readFile(getConfigFilePath())
        start = config_contents.find("default=") + len("default=")
        end = config_contents.find("saved=") - 1
        range = Range.from_indices(config_contents, start, end)
        range.end.line -= 1
        await sdk.ide.highlightCode(
            RangeInFile(filepath=getConfigFilePath(), range=range)
        )