# services/llm/ollama_service.py
"""
Ollama service implementation
"""
import ollama
import json
from typing import Dict, List
from .base import LLMService
from utils.logger import setup_logger


class OllamaService(LLMService):
    def __init__(self, model: str = "qwen3:latest", max_tokens: int = 4000):
        self.model = model
        # Explicitly set the host to avoid potential DNS/proxy issues with 'localhost'
        self.client = ollama.Client(host="127.0.0.1:11434")
        self.max_tokens = max_tokens
        self.logger = setup_logger("ollama")

    def generate_text(self, prompt: str) -> str:
        self.logger.info(f"--- PROMPT ---\n{prompt}")
        try:
            options = {"num_predict": self.max_tokens}
            response = self.client.generate(
                model=self.model, prompt=prompt, options=options
            )
            response_content = response["response"]
            self.logger.info(f"--- RESPONSE ---\n{response_content}")
            return response_content
        except Exception as e:
            self.logger.error(f"Error in Ollama API call: {e}")
            print(f"Error in Ollama API call: {e}")
            return None

    def get_similarity_scores(self, texts_pairs: Dict[str, List[str]]) -> List[float]:
        system_prompt = (
            "Evaluate the semantic similarity between the following table of pairs of texts in json format on a scale from 0 to 1. "
            "Return the similarity scores for every row in JSON format as a list of numbers, without any additional text or formatting."
        )

        request_payload = json.dumps(texts_pairs)
        prompt = f"{system_prompt}\n\n{request_payload}"

        try:
            options = {"num_predict": self.max_tokens}
            response = self.client.generate(
                model=self.model, prompt=prompt, options=options
            )

            try:
                scores = json.loads(response["response"].strip())
                if isinstance(scores, dict) and "similarity_scores" in scores:
                    return scores["similarity_scores"]
                elif isinstance(scores, list):
                    return scores
                else:
                    raise ValueError("Unexpected response format")
            except json.JSONDecodeError:
                raise ValueError("Could not decode response as JSON")

        except Exception as e:
            print(f"Error in Ollama similarity calculation: {e}")
            return None