from collections import defaultdict
from typing import List, Optional, Tuple

import click
import llm
import ollama
from pydantic import Field


@llm.hookimpl
def register_commands(cli):
    @cli.group(name="ollama")
    def ollama_group():
        "Commands for working with models hosted on Ollama"

    @ollama_group.command(name="list-models")
    def list_models():
        """List models that are available locally on Ollama server."""
        for model in ollama.list()["models"]:
            click.echo(model["name"])


@llm.hookimpl
def register_models(register):
    models = defaultdict(list)
    for model in ollama.list()["models"]:
        models[model["digest"]].append(model["name"])
        if model["name"].endswith(":latest"):
            models[model["digest"]].append(model["name"][: -len(":latest")])
    for names in models.values():
        name, aliases = _pick_primary_name(names)
        register(Ollama(name), aliases=aliases)


class Ollama(llm.Model):
    can_stream: bool = True

    class Options(llm.Options):
        temperature: Optional[float] = Field(
            description=(
                "The temperature of the model. Increasing the temperature will make the model answer more creatively."
            ),
            ge=0,
            le=1,
            default=0.8,
        )
        # TODO: Implement more options

    def __init__(
        self,
        model_id: str,
    ) -> None:
        self.model_id = model_id

    def __str__(self) -> str:
        return f"Ollama: {self.model_id}"

    def execute(
        self,
        prompt: llm.Prompt,
        stream: bool,
        response: llm.Response,
        conversation=None,
    ):
        messages = self.build_messages(prompt, conversation)
        response._prompt_json = {"messages": messages}
        options = {}
        if prompt.options.temperature:
            options["temperature"] = prompt.options.temperature

        if stream:
            response_stream = ollama.chat(
                model=self.model_id,
                messages=messages,
                stream=True,
                options=options,
            )
            for chunk in response_stream:
                yield chunk["message"]["content"]
        else:
            response.response_json = ollama.chat(
                model=self.model_id,
                messages=messages,
                options=options,
            )
            yield response.response_json["message"]["content"]

    def build_messages(self, prompt, conversation):
        messages = []
        if not conversation:
            if prompt.system:
                messages.append({"role": "system", "content": prompt.system})
            messages.append({"role": "user", "content": prompt.prompt})
            return messages
        current_system = None
        for prev_response in conversation.responses:
            if (
                prev_response.prompt.system
                and prev_response.prompt.system != current_system
            ):
                messages.append(
                    {"role": "system", "content": prev_response.prompt.system}
                )
                current_system = prev_response.prompt.system
            messages.append({"role": "user", "content": prev_response.prompt.prompt})
            messages.append({"role": "assistant", "content": prev_response.text()})
        if prompt.system and prompt.system != current_system:
            messages.append({"role": "system", "content": prompt.system})
        messages.append({"role": "user", "content": prompt.prompt})
        return messages


def _pick_primary_name(names: List[str]) -> Tuple[str, List[str]]:
    """Pick the primary model name from a list of names.

    The picking algorithm prefers names with the most specific tag, e.g. "llama2:7b"
    over "llama2:latest" over "llama2".

    Parameters
    ----------
    names : list[str]
        A non-empty list of model names.

    Returns
    -------
    tuple[str, list[str, ...]]
        The primary model name and a list with the secondary names.
    """
    if len(names) == 1:
        return names[0], ()
    sorted_names = sorted(
        names,
        key=lambda name: (
            ":" not in name,
            name.endswith(":latest"),
            name,
        ),
    )
    return sorted_names[0], tuple(sorted_names[1:])
