#  Copyright © 2025 Dr.-Ing. Paul Wilhelm <paul@wilhelm.dev>
#  This file is part of Archive Agent. See LICENSE for details.

import json
import hashlib
import re
import time
from logging import Logger
from abc import ABC, abstractmethod
from typing import Callable, Optional, cast

from archive_agent.ai.AiResult import AiResult
from archive_agent.core.CacheManager import CacheManager

from archive_agent.ai_provider.AiProviderParams import AiProviderParams


class AiProvider(ABC):
    """
    AI provider.
    """

    AI_REQUEST_TIMEOUT_S = 120

    def __init__(
            self,
            logger: Logger,
            cache: CacheManager,
            invalidate_cache: bool,
            params: AiProviderParams,
            server_url: str,
    ):
        """
        Initialize AI provider.
        :param logger: Logger.
        :param cache: Cache manager.
        :param invalidate_cache: Invalidate cache if enabled, probe cache otherwise.
        :param params: AI provider parameters.
        :param server_url: Server URL.
        """
        self.logger = logger

        self.cache = cache
        self.invalidate_cache = invalidate_cache

        self.params = params

        self.server_url = server_url

        self.supports_vision = self.params.model_vision != ""

        self._last_cache_key: Optional[str] = None

    @staticmethod
    def _sanitize_json(json_raw: str) -> str:
        """
        Strip control characters that break JSON parsing.
        AI models may echo control characters from garbage input (e.g., OCR artifacts).
        Preserves newline, carriage return, and tab (valid JSON whitespace).
        :param json_raw: Raw JSON string from AI response.
        :return: Sanitized JSON string.
        """
        return re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f]', '', json_raw)

    def _handle_cached_request(
            self,
            cache_key_prefix: str,
            callback: Callable,
            callback_kwargs: dict,
    ) -> AiResult:
        """
        Handle cached request.
        :param cache_key_prefix: Cache key prefix.
        :param callback: Callback to execute on a cache miss.
        :param callback_kwargs: Keyword arguments for the callback.
        :return: AI result.
        """
        callback_kwargs_str = json.dumps(callback_kwargs, sort_keys=True)
        params_str = self.params.get_static_cache_key()
        cache_str = f"{cache_key_prefix}:{callback_kwargs_str}:{params_str}"
        cache_key = hashlib.sha256(cache_str.encode('utf-8')).hexdigest()
        self._last_cache_key = cache_key

        cached_result = self.cache.get(key=cache_key, display_key=cache_key_prefix)
        if cached_result is not None:
            ai_result: AiResult = cast(AiResult, cached_result)
            ai_result.total_tokens = 0  # Cached result consumed no tokens
            return ai_result

        t0 = time.monotonic()
        result: AiResult = callback(**callback_kwargs)
        elapsed = time.monotonic() - t0
        self.logger.info(f"API call '{cache_key_prefix}' completed in {elapsed:.1f}s")

        # Cache write.
        self.cache[cache_key] = result

        return result

    def invalidate_last_cached(self) -> None:
        """
        Delete the last cached result by its specific key.
        Thread-safe: each worker thread owns a dedicated AiProvider instance,
        so _last_cache_key is never shared between threads.
        """
        if self._last_cache_key is not None:
            try:
                del self.cache[self._last_cache_key]
            except KeyError:
                pass

    @abstractmethod
    def _perform_chunk_callback(self, prompt) -> AiResult:
        """
        Perform chunk callback.
        :param prompt: Prompt.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        raise NotImplementedError

    def chunk_callback(self, prompt: str) -> AiResult:
        """
        Chunk callback with caching.
        :param prompt: Prompt.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        return self._handle_cached_request(
            cache_key_prefix="chunk_callback",
            callback=self._perform_chunk_callback,
            callback_kwargs=dict(prompt=prompt),
        )

    @abstractmethod
    def _perform_embed_callback(self, text: str) -> AiResult:
        """
        Perform embed callback.
        :param text: Text.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        raise NotImplementedError

    def embed_callback(self, text: str) -> AiResult:
        """
        Embed callback with caching.
        :param text: Text.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        return self._handle_cached_request(
            cache_key_prefix="embed_callback",
            callback=self._perform_embed_callback,
            callback_kwargs=dict(text=text),
        )

    @abstractmethod
    def _perform_rerank_callback(self, prompt: str) -> AiResult:
        """
        Perform rerank callback.
        :param prompt: Prompt.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        raise NotImplementedError

    def rerank_callback(self, prompt: str) -> AiResult:
        """
        Rerank callback.
        :param prompt: Prompt.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        return self._handle_cached_request(
            cache_key_prefix="rerank_callback",
            callback=self._perform_rerank_callback,
            callback_kwargs=dict(prompt=prompt),
        )

    @abstractmethod
    def _perform_query_callback(self, prompt: str) -> AiResult:
        """
        Perform query callback.
        :param prompt: Prompt.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        raise NotImplementedError

    def query_callback(self, prompt: str) -> AiResult:
        """
        Query callback.
        NOTE: This call is NOT cached, as the user expects a novel answer on each call.
        :param prompt: Prompt.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        return self._perform_query_callback(prompt=prompt)

    @abstractmethod
    def _perform_vision_callback(self, prompt: str, image_base64: str) -> AiResult:
        """
        Perform vision callback.
        :param prompt: Prompt.
        :param image_base64: Image as UTF-8 encoded Base64 string.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        raise NotImplementedError

    def vision_callback(self, prompt: str, image_base64: str) -> AiResult:
        """
        Vision callback with caching.
        :param prompt: Prompt.
        :param image_base64: Image as UTF-8 encoded Base64 string.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        return self._handle_cached_request(
            cache_key_prefix="vision_callback",
            callback=self._perform_vision_callback,
            callback_kwargs=dict(prompt=prompt, image_base64=image_base64),
        )
