import os
import re
from typing import List, Dict

from groq import Groq
from sentence_transformers import SentenceTransformer

from agents.agent import Agent
from free_config import GROQ_MODEL


class FrontierAgent(Agent):
    name = "Frontier Agent"
    color = Agent.BLUE

    def __init__(self, collection):
        """
        Set up this instance by connecting to Groq, the Chroma datastore,
        and the local sentence-transformer embedding model.
        """
        self.log("Initializing Frontier Agent")
        self.client = Groq(api_key=os.environ["GROQ_API_KEY"])
        self.model = GROQ_MODEL
        self.log(f"Frontier Agent is setting up with Groq ({self.model})")
        self.collection = collection
        self.encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
        self.log("Frontier Agent is ready")

    def make_context(self, similars: List[str], prices: List[float]) -> str:
        """
        Create context that can be inserted into the prompt
        """
        message = "To provide some context, here are some other items that might be similar to the item you need to estimate.\n\n"
        for similar, price in zip(similars, prices):
            message += f"Potentially related product:\n{similar}\nPrice is ${price:.2f}\n\n"
        return message

    def messages_for(
        self, description: str, similars: List[str], prices: List[float]
    ) -> List[Dict[str, str]]:
        """
        Create the message list for the Groq chat completion call
        """
        message = f"Estimate the price of this product. Respond with the price only, no explanation.\n\n{description}\n\n"
        message += self.make_context(similars, prices)
        return [{"role": "user", "content": message}]

    def find_similars(self, description: str):
        """
        Return a list of items similar to the given one by looking in the Chroma datastore
        """
        self.log(
            "Frontier Agent is performing a RAG search of the Chroma datastore to find 5 similar products"
        )
        vector = self.encoder.encode([description])
        results = self.collection.query(query_embeddings=vector.astype(float).tolist(), n_results=5)
        documents = results["documents"][0][:]
        prices = [m["price"] for m in results["metadatas"][0][:]]
        self.log("Frontier Agent has found similar products")
        return documents, prices

    def get_price(self, s) -> float:
        """
        A utility that plucks a floating point number out of a string
        """
        s = s.replace("$", "").replace(",", "")
        match = re.search(r"[-+]?\d*\.\d+|\d+", s)
        return float(match.group()) if match else 0.0

    def price(self, description: str) -> float:
        """
        Call Groq to estimate the price using RAG context from similar products.
        """
        documents, prices = self.find_similars(description)
        self.log(
            f"Frontier Agent is about to call Groq ({self.model}) with context including 5 similar products"
        )
        response = self.client.chat.completions.create(
            model=self.model,
            messages=self.messages_for(description, documents, prices),
            temperature=0,
        )
        reply = response.choices[0].message.content
        result = self.get_price(reply)
        self.log(f"Frontier Agent completed - predicting ${result:.2f}")
        return result