from __future__ import annotations

import os
from enum import IntEnum
from typing import Dict, Sequence

from ramalama.common import generate_sha256_binary, perror
from ramalama.http_client import download_file
from ramalama.logger import logger


class SnapshotFileType(IntEnum):
    GGUFModel = 1
    ChatTemplate = 2
    Other = 3
    Mmproj = 4
    SafetensorModel = 5


class SnapshotFile:
    def __init__(
        self,
        url: str,
        header: Dict,
        hash: str,
        name: str,
        type: SnapshotFileType,
        should_show_progress: bool = False,
        should_verify_checksum: bool = False,
        required: bool = True,
    ):
        self.url: str = url
        self.header: Dict = header
        self.hash: str = hash
        self.name: str = name
        self.type: SnapshotFileType = type
        self.should_show_progress: bool = should_show_progress
        self.should_verify_checksum: bool = should_verify_checksum
        self.required: bool = required

    def download(self, blob_file_path: str, snapshot_dir: str) -> str:
        if not os.path.exists(blob_file_path):
            if self.should_show_progress:
                perror(f"Downloading {self.name}")
            download_file(
                url=self.url,
                headers=self.header,
                dest_path=blob_file_path,
                show_progress=self.should_show_progress,
            )
        else:
            logger.debug(f"Using cached blob for {self.name} ({os.path.basename(blob_file_path)})")
        prefix = os.path.dirname(self.name)
        if prefix:
            snapshot_dir = os.path.join(snapshot_dir, prefix)
        return os.path.relpath(blob_file_path, start=snapshot_dir)


class LocalSnapshotFile(SnapshotFile):
    def __init__(
        self,
        content: bytes,
        name: str,
        type: SnapshotFileType,
        should_show_progress: bool = False,
        should_verify_checksum: bool = False,
        required: bool = True,
    ):
        super().__init__(
            "",
            {},
            generate_sha256_binary(content),
            name,
            type,
            should_show_progress,
            should_verify_checksum,
            required,
        )
        self.content = content

    def download(self, blob_file_path, snapshot_dir):
        with open(blob_file_path, "wb") as file:
            file.write(self.content)
            file.flush()
        return os.path.relpath(blob_file_path, start=snapshot_dir)


def validate_snapshot_files(snapshot_files: Sequence[SnapshotFile]):
    chat_template_files = []
    mmproj_files = []
    gguf_files = []
    safetensor_files = []
    for file in snapshot_files:
        if file.type == SnapshotFileType.ChatTemplate:
            chat_template_files.append(file)
        if file.type == SnapshotFileType.Mmproj:
            mmproj_files.append(file)
        if file.type == SnapshotFileType.GGUFModel:
            gguf_files.append(file)
        if file.type == SnapshotFileType.SafetensorModel:
            safetensor_files.append(file)

    if len(gguf_files) > 1 and len(safetensor_files) > 1:
        raise ValueError("Snapshot must contain only .gguf or .safetensors model files. Found both.")
    if len(chat_template_files) > 1:
        raise ValueError(f"Only one chat template supported, got {len(chat_template_files)}: {chat_template_files}")
    if len(mmproj_files) > 1:
        raise ValueError(f"Only one mmproj supported, got {len(mmproj_files)}: {mmproj_files}")
