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

import json
import traceback
from logging import Logger

from openai import OpenAI

from archive_agent.ai_provider.AiProvider import AiProvider
from archive_agent.ai_provider.AiProviderError import AiProviderError
from archive_agent.ai.AiResult import AiResult
from archive_agent.ai_provider.AiProviderParams import AiProviderParams

from archive_agent.ai.chunk.AiChunk import ChunkSchema
from archive_agent.ai.rerank.AiRerank import RerankSchema
from archive_agent.ai.query.AiQuery import QuerySchema
from archive_agent.ai.vision.AiVisionSchema import VisionSchema

from archive_agent.core.CacheManager import CacheManager


class LMStudioProvider(AiProvider):
    """
    LM Studio provider with structured output and optional vision support.
    """

    def __init__(
            self,
            logger: Logger,
            cache: CacheManager,
            invalidate_cache: bool,
            params: AiProviderParams,
            server_url: str,
    ):
        """
        Initialize LM Studio 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.
        """
        AiProvider.__init__(
            self,
            logger=logger,
            cache=cache,
            invalidate_cache=invalidate_cache,
            params=params,
            server_url=server_url,
        )

        self.client = OpenAI(base_url=self.server_url, api_key="lm-studio", timeout=AiProvider.AI_REQUEST_TIMEOUT_S)

    def _perform_chunk_callback(self, prompt: str) -> AiResult:
        """
        Chunk callback.
        :param prompt: Prompt.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        # noinspection PyTypeChecker
        response = self.client.chat.completions.create(
            model=self.params.model_chunk,
            temperature=0.0,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": prompt,
                        },
                    ],
                },
            ],
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": ChunkSchema.__name__,
                    "schema": ChunkSchema.model_json_schema(),
                    "strict": True,
                },
            },
        )

        formatted_response = json.dumps(response.__dict__, indent=2, default=str)

        json_raw = response.choices[0].message.content
        if json_raw is None:
            raise AiProviderError(f"Missing JSON\n{formatted_response}")
        try:
            parsed_schema = ChunkSchema.model_validate_json(self._sanitize_json(json_raw))
        except Exception as e:
            raise AiProviderError(f"Invalid JSON:\n{json_raw}\n{e}\n{formatted_response}")

        return AiResult(
            total_tokens=response.usage.total_tokens if response.usage else 0,
            output_text=json_raw,
            parsed_schema=parsed_schema,
        )

    def _perform_embed_callback(self, text: str) -> AiResult:
        """
        Embed callback.
        :param text: Text.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        try:
            response = self.client.embeddings.create(
                input=text,
                model=self.params.model_embed
            )

            return AiResult(
                total_tokens=response.usage.total_tokens,
                embedding=response.data[0].embedding
            )

        except Exception as e:
            text_preview = text[:200] + "..." if len(text) > 200 else text
            tb = traceback.format_exc()
            raise AiProviderError(
                f"Embedding failed ({len(text)} chars): {type(e).__name__}: {e}\n"
                f"Text: {text_preview}\n"
                f"Traceback:\n{tb}"
            )

    def _perform_rerank_callback(self, prompt: str) -> AiResult:
        """
        Rerank callback.
        :param prompt: Prompt.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        # noinspection PyTypeChecker
        response = self.client.chat.completions.create(
            model=self.params.model_rerank,
            temperature=0.0,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": prompt,
                        },
                    ],
                },
            ],
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": RerankSchema.__name__,
                    "schema": RerankSchema.model_json_schema(),
                    "strict": True,
                },
            },
        )

        formatted_response = json.dumps(response.__dict__, indent=2, default=str)

        json_raw = response.choices[0].message.content
        if json_raw is None:
            raise AiProviderError(f"Missing JSON\n{formatted_response}")
        try:
            parsed_schema = RerankSchema.model_validate_json(self._sanitize_json(json_raw))
        except Exception as e:
            raise AiProviderError(f"Invalid JSON:\n{json_raw}\n{e}\n{formatted_response}")

        return AiResult(
            total_tokens=response.usage.total_tokens if response.usage else 0,
            output_text=json_raw,
            parsed_schema=parsed_schema,
        )

    def _perform_query_callback(self, prompt: str) -> AiResult:
        """
        Query callback.
        :param prompt: Prompt.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        # noinspection PyTypeChecker
        response = self.client.chat.completions.create(
            model=self.params.model_query,
            temperature=self.params.temperature_query,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": prompt,
                        },
                    ],
                },
            ],
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": QuerySchema.__name__,
                    "schema": QuerySchema.model_json_schema(),
                    "strict": True,
                },
            },
        )

        formatted_response = json.dumps(response.__dict__, indent=2, default=str)

        json_raw = response.choices[0].message.content
        if json_raw is None:
            raise AiProviderError(f"Missing JSON\n{formatted_response}")
        try:
            parsed_schema = QuerySchema.model_validate_json(self._sanitize_json(json_raw))
        except Exception as e:
            raise AiProviderError(f"Invalid JSON:\n{json_raw}\n{e}\n{formatted_response}")

        return AiResult(
            total_tokens=response.usage.total_tokens if response.usage else 0,
            output_text=json_raw,
            parsed_schema=parsed_schema,
        )

    def _perform_vision_callback(self, prompt: str, image_base64: str) -> AiResult:
        """
        Vision callback.
        :param prompt: Prompt.
        :param image_base64: Image as UTF-8 encoded Base64 string.
        :return: AI result.
        :raises AiProviderError: On error.
        """
        if not self.params.model_vision:
            raise AiProviderError("Vision model is not configured.")

        # noinspection PyTypeChecker
        response = self.client.chat.completions.create(
            model=self.params.model_vision,
            temperature=0.0,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": prompt,
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_base64}"
                            },
                        },
                    ],
                },
            ],
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": VisionSchema.__name__,
                    "schema": VisionSchema.model_json_schema(),
                    "strict": True,
                },
            },
        )

        formatted_response = json.dumps(response.__dict__, indent=2, default=str)

        json_raw = response.choices[0].message.content
        if json_raw is None:
            raise AiProviderError(f"Missing JSON\n{formatted_response}")
        try:
            parsed_schema = VisionSchema.model_validate_json(self._sanitize_json(json_raw))
        except Exception as e:
            raise AiProviderError(f"Invalid JSON:\n{json_raw}\n{e}\n{formatted_response}")

        return AiResult(
            total_tokens=response.usage.total_tokens if response.usage else 0,
            output_text=json_raw,
            parsed_schema=parsed_schema,
        )
