from __future__ import annotations

import json
import os
import re
import tempfile
import urllib.request
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Optional

from ramalama.common import (
    SPLIT_MODEL_PATH_RE,
    available,
    exec_cmd,
    generate_sha256,
    is_split_file_model,
    perror,
    run_cmd,
)
from ramalama.endian import EndianMismatchError
from ramalama.logger import logger
from ramalama.model_store.snapshot_file import SnapshotFile, SnapshotFileType
from ramalama.transports.base import Transport


class HFStyleRepoFile(SnapshotFile):
    def __init__(
        self, url, header, hash, name, type, should_show_progress=False, should_verify_checksum=False, required=True
    ):
        super().__init__(url, header, hash, name, type, should_show_progress, should_verify_checksum, required)

    def download(self, blob_file_path, snapshot_dir):
        # moving from the cached temp directory to blob directory
        import shutil

        shutil.move(self.url, blob_file_path)
        return os.path.relpath(blob_file_path, start=snapshot_dir)


def fetch_checksum_from_api_base(checksum_api_url, headers=None, extractor_func=None):
    """
    Base function for fetching checksums from API endpoints.

    Args:
    checksum_api_url (str): The URL of the API endpoint to fetch the checksum from.
    headers (dict, optional): Optional headers to include in the request.
    extractor_func (callable, optional): Optional function to extract the checksum from the response data.

    Returns:
    str: The extracted checksum or the raw response data.

    Raises:
    KeyError: If the API request fails or the checksum cannot be extracted.
    """
    logger.debug(f"Fetching checksum from {checksum_api_url}")
    request = urllib.request.Request(url=checksum_api_url)
    if headers:
        for key, value in headers.items():
            request.add_header(key, value)

    try:
        with urllib.request.urlopen(request) as response:
            data = response.read().decode()

        return extractor_func(data) if extractor_func else data.strip()

    except (json.JSONDecodeError, urllib.error.HTTPError, urllib.error.URLError) as e:
        raise KeyError(f"failed to pull {checksum_api_url}: {str(e).strip()}")


class HFStyleRepository(ABC):
    FILE_NAME_CONFIG = "config.json"
    FILE_NAME_GENERATION_CONFIG = "generation_config.json"
    FILE_NAME_TOKENIZER_CONFIG = "tokenizer_config.json"

    def __init__(self, name: str, organization: str, tag: str = 'latest'):
        self.name = name
        self.organization = organization
        self.tag = tag
        self.headers: dict = {}
        self.blob_url: Optional[str] = None
        self.model_filename: Optional[str] = None
        self.model_hash: Optional[str] = None
        self.mmproj_filename: Optional[str] = None
        self.mmproj_hash: Optional[str] = None
        self.other_files: list[dict] = []
        self.additional_safetensor_files: list[dict] = []
        self.safetensors_index_file: Optional[str] = None
        self.fetch_metadata()

    @abstractmethod
    def fetch_metadata(self):
        pass

    def get_file_list(self, cached_files: list[str]) -> list[SnapshotFile]:
        files = []
        if self.model_filename not in cached_files:
            files.append(self.model_file())
        assert self.model_filename
        if is_split_file_model(self.model_filename):
            # If the model is split, we need to add all parts
            match = re.match(SPLIT_MODEL_PATH_RE, self.model_filename)
            if match:
                path_part = match[1]
                if path_part:
                    path_part += '/'
                filename_base = match[2]
                total_parts = int(match[3])
                for i in range(2, total_parts + 1):
                    file_name = f"{path_part}{filename_base}-{i:05d}-of-{total_parts:05d}.gguf"
                    if file_name not in cached_files:
                        # Add the split file part if it is not already cached
                        logger.debug(f"Adding split file part: {file_name}")
                        files.append(
                            SnapshotFile(
                                url=f"{self.blob_url}/{file_name}",
                                header=self.headers,
                                hash=generate_sha256(file_name),
                                type=SnapshotFileType.Other,
                                name=file_name,
                                should_show_progress=True,
                                should_verify_checksum=False,
                            )
                        )

        # Handle additional safetensors files (for sharded models)
        for safetensor_file in self.additional_safetensor_files:
            filename = safetensor_file['filename']
            if filename not in cached_files:
                logger.debug(f"Adding safetensors file: {filename}")
                files.append(
                    SnapshotFile(
                        url=f"{self.blob_url}/{filename}",
                        header=self.headers,
                        hash=f"sha256:{safetensor_file['oid']}",
                        type=SnapshotFileType.SafetensorModel,
                        name=filename,
                        should_show_progress=True,
                        should_verify_checksum=True,
                    )
                )

        # Other files
        for other_file in self.other_files:
            filename = other_file['filename']
            if filename not in cached_files:
                logger.debug(f"Adding other safetensors file: {filename}")
                # The file's oid is either a verifiable SHA256 lfs hash or a SHA1 git
                # hash, in which case generate a non-verifiable SHA256 hash.
                # TODO: Calculate a verifiable content hash
                oid = other_file['oid']
                file_hash = (
                    oid
                    if len(oid) == 64
                    else generate_sha256(f"{self.organization}/{self.name}/{filename}", with_sha_prefix=False)
                )
                files.append(
                    SnapshotFile(
                        url=f"{self.blob_url}/{filename}",
                        header=self.headers,
                        hash=f"sha256:{file_hash}",
                        type=SnapshotFileType.Other,
                        name=filename,
                        should_show_progress=False,
                        should_verify_checksum=len(oid) == 64,
                    )
                )

        if self.mmproj_filename and self.mmproj_filename not in cached_files:
            files.append(self.mmproj_file())

        # Add safetensors index file if present
        if self.safetensors_index_file and self.safetensors_index_file not in cached_files:
            logger.debug(f"Adding safetensors index file: {self.safetensors_index_file}")
            files.append(
                SnapshotFile(
                    url=f"{self.blob_url}/{self.safetensors_index_file}",
                    header=self.headers,
                    hash=generate_sha256(
                        f"{self.organization}/{self.name}/{self.safetensors_index_file}", with_sha_prefix=False
                    ),
                    type=SnapshotFileType.Other,
                    name=self.safetensors_index_file,
                    required=False,
                )
            )

        if self.FILE_NAME_CONFIG not in cached_files:
            files.append(self.config_file())
        if self.FILE_NAME_GENERATION_CONFIG not in cached_files:
            files.append(self.generation_config_file())
        if self.FILE_NAME_TOKENIZER_CONFIG not in cached_files:
            files.append(self.tokenizer_config_file())

        return files

    def model_file(self) -> SnapshotFile:
        assert self.model_filename
        assert self.model_hash

        # Determine file type based on extension
        file_type = SnapshotFileType.GGUFModel
        if Path(self.model_filename).match('*.safetensors'):
            file_type = SnapshotFileType.SafetensorModel

        return SnapshotFile(
            url=f"{self.blob_url}/{self.model_filename}",
            header=self.headers,
            hash=self.model_hash,
            type=file_type,
            name=self.model_filename,
            should_show_progress=True,
            should_verify_checksum=True,
        )

    def mmproj_file(self) -> SnapshotFile:
        assert self.model_filename
        assert self.model_hash
        assert self.mmproj_filename is not None
        assert self.mmproj_hash is not None
        return SnapshotFile(
            url=f"{self.blob_url}/{self.mmproj_filename}",
            header=self.headers,
            hash=self.mmproj_hash,
            type=SnapshotFileType.Mmproj,
            name=self.mmproj_filename,
            required=False,
            should_show_progress=True,
            should_verify_checksum=True,
        )

    def config_file(self) -> SnapshotFile:
        return SnapshotFile(
            url=f"{self.blob_url}/{self.FILE_NAME_CONFIG}",
            header=self.headers,
            hash=generate_sha256(self.FILE_NAME_CONFIG),
            type=SnapshotFileType.Other,
            name=self.FILE_NAME_CONFIG,
            required=False,
        )

    def generation_config_file(self) -> SnapshotFile:
        return SnapshotFile(
            url=f"{self.blob_url}/{self.FILE_NAME_GENERATION_CONFIG}",
            header=self.headers,
            hash=generate_sha256(self.FILE_NAME_GENERATION_CONFIG),
            type=SnapshotFileType.Other,
            name=self.FILE_NAME_GENERATION_CONFIG,
            required=False,
        )

    def tokenizer_config_file(self) -> SnapshotFile:
        return SnapshotFile(
            url=f"{self.blob_url}/{self.FILE_NAME_TOKENIZER_CONFIG}",
            header=self.headers,
            hash=generate_sha256(self.FILE_NAME_TOKENIZER_CONFIG),
            type=SnapshotFileType.Other,
            name=self.FILE_NAME_TOKENIZER_CONFIG,
            required=False,
        )


class HFStyleRepoModel(Transport, ABC):
    def __init__(self, model, model_store_path):
        super().__init__(model, model_store_path)

    @abstractmethod
    def get_cli_command(self):
        """Return the CLI command name (e.g., 'hf', 'modelscope')"""
        pass

    @abstractmethod
    def get_missing_message(self):
        """Return the missing CLI message"""
        pass

    @abstractmethod
    def get_registry_url(self):
        """Return the registry URL"""
        pass

    @abstractmethod
    def get_accept_header(self):
        """Return the accept header"""
        pass

    @abstractmethod
    def get_repo_type(self):
        """Return the repo type name (e.g., 'huggingface', 'modelscope')"""
        pass

    @abstractmethod
    def fetch_checksum_from_api(self, organization, file):
        """Fetch checksum from API"""
        pass

    @abstractmethod
    def create_repository(self, name, organization, tag):
        """Create repository instance"""
        pass

    @abstractmethod
    def get_cli_download_args(self, directory_path, model):
        """Get CLI download arguments"""
        pass

    @abstractmethod
    def in_existing_cache(self, args, target_path, sha256_checksum):
        """Check if file exists in existing cache"""
        pass

    @abstractmethod
    def _collect_cli_files(self, tempdir: str):
        """Collect files from CLI download"""
        pass

    def get_login_args(self):
        """Return the login subcommand arguments. Can be overridden for different CLI structures."""
        return ["login"]

    def get_logout_args(self):
        """Return the logout subcommand arguments. Can be overridden for different CLI structures."""
        return ["logout"]

    def login(self, args):
        if not available(self.get_cli_command()):
            raise NotImplementedError(self.get_missing_message())
        conman_args = [self.get_cli_command()] + self.get_login_args()
        if args.token:
            conman_args.extend(["--token", args.token])
        self.exec(conman_args, args)

    def logout(self, args):
        if not available(self.get_cli_command()):
            raise NotImplementedError(self.get_missing_message())
        conman_args = [self.get_cli_command()] + self.get_logout_args()
        if args.token:
            conman_args.extend(["--token", args.token])
        self.exec(conman_args, args)

    def pull(self, args):
        name, tag, organization = self.extract_model_identifiers()
        _, cached_files, all = self.model_store.get_cached_files(tag)
        if all:
            if not args.quiet:
                perror(f"Using cached {self.get_repo_type()}://{name}:{tag} ...")
            return

        try:
            if not args.quiet:
                self.print_pull_message(f"{self.get_repo_type()}://{organization}/{name}:{tag}")

            repo = self.create_repository(name, organization, tag)
            snapshot_hash = repo.model_hash
            files = repo.get_file_list(cached_files)
            self.model_store.new_snapshot(tag, snapshot_hash, files, verify=getattr(args, "verify", True))

        except EndianMismatchError:
            # No use pulling again
            raise
        except Exception as e:
            if isinstance(e, OSError) and hasattr(e, "winerror") and e.winerror == 206:
                raise  # Path too long on Windows
            if not available(self.get_cli_command()):
                perror(f"URL pull failed and {self.get_cli_command()} not available")
                raise KeyError(f"Failed to pull model: {str(e)}")

            # Create temporary directory for downloading via CLI
            # Ensure the base store path exists before creating the temp directory inside it
            os.makedirs(self.model_store.base_path, exist_ok=True)
            with tempfile.TemporaryDirectory(prefix="tmp_hfcli_", dir=self.model_store.base_path) as tempdir:
                model = f"{organization}/{name}"
                conman_args = self.get_cli_download_args(tempdir, model)
                run_cmd(conman_args)

                snapshot_hash, files = self._collect_cli_files(tempdir)
                self.model_store.new_snapshot(tag, snapshot_hash, files, verify=getattr(args, "verify", True))

    def exec(self, cmd_args, args):
        try:
            exec_cmd(cmd_args)
        except FileNotFoundError as e:
            perror(f"{str(e).strip()}\n{self.get_missing_message()}")
