from __future__ import annotations

import os
import re
import shutil
from pathlib import Path

from ramalama.common import SPLIT_MODEL_PATH_RE, generate_sha256, is_split_file_model
from ramalama.model_store.snapshot_file import SnapshotFile, SnapshotFileType
from ramalama.path_utils import normalize_host_path_for_container
from ramalama.transports.base import Transport
from ramalama.transports.huggingface import HuggingfaceRepository
from ramalama.transports.modelscope import ModelScopeRepository


class LocalModelFile(SnapshotFile):
    def __init__(
        self,
        url,
        header,
        model_file_hash,
        name,
        should_show_progress=False,
        should_verify_checksum=False,
        required=True,
    ):
        super().__init__(
            url,
            header,
            model_file_hash,
            name,
            SnapshotFileType.GGUFModel,
            should_show_progress,
            should_verify_checksum,
            required,
        )

    def download(self, blob_file_path, snapshot_dir):
        if not os.path.exists(self.url):
            raise FileNotFoundError(f"No such file: '{self.url}'")
        # moving from the local location to blob directory so the model store "owns" the data
        shutil.copy(self.url, blob_file_path)
        return os.path.relpath(blob_file_path, start=snapshot_dir)


class URL(Transport):
    def __init__(self, model, model_store_path, scheme):
        self.type = scheme
        super().__init__(model, model_store_path)

        # Use the URL scheme as model type so we can distinguish
        # between the various types such as http, https and file
        self._model_type = scheme

        split = self.model.rsplit("/", 1)
        self.directory = split[0].removeprefix("/") if len(split) > 1 else ""

    @property
    def model_alias(self):
        if self.type == "file":
            return self.model
        return super().model_alias()

    def extract_model_identifiers(self):
        if self.type == "file":
            return (
                Path(self.model).name,
                "latest",
                normalize_host_path_for_container(str(Path(self.model).parent)).removeprefix("/"),
            )

        model_name, model_tag, model_organization = super().extract_model_identifiers()

        parts = model_organization.split("/")
        if len(parts) > 2 and parts[-2] == "blob":
            model_organization = "/".join(parts[:-2])
            model_tag = parts[-1]

        # handling huggingface specific URLs for more precise identifiers
        if len(parts) > 2 and HuggingfaceRepository.REGISTRY_URL.endswith(parts[0]) and parts[-2] == "resolve":
            model_organization = "/".join(parts[:-2])
            model_tag = parts[-1]

        if len(parts) > 3 and parts[-3] == "file":
            model_organization = "/".join(parts[:-3])
            model_tag = parts[-1]

        # handling modelscope specific URLs for more precise identifiers
        if len(parts) > 2 and ModelScopeRepository.REGISTRY_URL.endswith(parts[0]) and parts[-2] == "resolve":
            model_organization = "/".join(parts[:-2])
            model_tag = parts[-1]

        return model_name, model_tag, model_organization

    def _assemble_split_file_list(self, snapshot_hash: str) -> list[SnapshotFile]:
        files: list[SnapshotFile] = []

        # model is split, lets fetch all files based on the name pattern
        match = re.match(SPLIT_MODEL_PATH_RE, self.model)
        if match is None:
            return files

        path_part = match[1]
        if path_part:
            path_part += '/'
        filename_base = match[2]
        total_parts = int(match[3])

        for i in range(1, total_parts + 1):
            file_name = f"{filename_base}-{i:05d}-of-{total_parts:05d}.gguf"
            url = f"{self.type}://{path_part}{file_name}"
            if self.type == "file":
                files.append(
                    LocalModelFile(
                        url=f"{path_part}{file_name}",
                        header={},
                        model_file_hash=generate_sha256(file_name),
                        name=file_name,
                        required=True,
                    )
                )
            else:
                files.append(
                    SnapshotFile(
                        url=url,
                        header={},
                        hash=generate_sha256(file_name),
                        type=SnapshotFileType.GGUFModel,
                        name=file_name,
                        should_show_progress=True,
                        required=True,
                    )
                )

        return files

    def pull(self, args) -> None:
        name, tag, _ = self.extract_model_identifiers()
        _, _, all_files = self.model_store.get_cached_files(tag)
        if all_files:
            return

        files: list[SnapshotFile] = []
        snapshot_hash = generate_sha256(name)
        if is_split_file_model(self.model):
            files = self._assemble_split_file_list(snapshot_hash)
            self.model_store.new_snapshot(tag, snapshot_hash, files, verify=getattr(args, "verify", True))
            return

        if self.type == "file":
            files.append(
                LocalModelFile(
                    url=self.model,  # model contains the full path here
                    header={},
                    model_file_hash=snapshot_hash,
                    name=name,
                    required=True,
                )
            )
            self.model_store.new_snapshot(tag, snapshot_hash, files, verify=getattr(args, "verify", True))
            return

        files.append(
            SnapshotFile(
                url=f"{self.type}://{self.model}",
                header={},
                hash=snapshot_hash,
                type=SnapshotFileType.GGUFModel,
                name=name,
                should_show_progress=True,
                required=True,
            )
        )
        self.model_store.new_snapshot(tag, snapshot_hash, files, verify=getattr(args, "verify", True))
        return
