import os
import re

from agents.agent import Agent
from free_config import GROQ_MODEL, USE_MODAL_SPECIALIST

QUESTION = "What does this cost to the nearest dollar?"
PREFIX = "Price is $"


class SpecialistAgent(Agent):
    """
    Prices products using either:
    1. Fine-tuned Llama on Modal (best quality, uses GPU credits), or
    2. Groq fallback (free — same prompt format as the fine-tuned model)

    Set USE_MODAL_SPECIALIST=true only if you have Modal credits to spare.
    """

    name = "Specialist Agent"
    color = Agent.RED

    def __init__(self):
        self.pricer = None
        self.groq_client = None
        self.mode = "off"
        self.model = GROQ_MODEL

        if USE_MODAL_SPECIALIST:
            self._try_modal()

        if self.mode == "off" and os.getenv("GROQ_API_KEY"):
            from groq import Groq

            self.groq_client = Groq(api_key=os.environ["GROQ_API_KEY"])
            self.mode = "groq"
            self.log(f"Specialist Agent using Groq fallback ({self.model}) — no Modal GPU needed")

        if self.mode == "off":
            self.log("Specialist Agent disabled — ensemble will use Frontier + Neural Network only")

    def _try_modal(self):
        try:
            import modal

            self.log("Specialist Agent is initializing - connecting to Modal")
            Pricer = modal.Cls.from_name("pricer-service", "Pricer")
            self.pricer = Pricer()
            self.mode = "modal"
            self.log("Specialist Agent connected to Modal pricer-service")
        except Exception as exc:
            self.log(f"Modal unavailable ({exc}) — will try Groq fallback if configured")

    @property
    def available(self) -> bool:
        return self.mode != "off"

    @staticmethod
    def _parse_price(text: str) -> float:
        if PREFIX in text:
            text = text.split(PREFIX, 1)[1]
        text = text.replace("$", "").replace(",", "")
        match = re.search(r"[-+]?\d*\.\d+|\d+", text)
        return float(match.group()) if match else 0.0

    def _price_with_groq(self, description: str) -> float:
        prompt = f"{QUESTION}\n\n{description}\n\n{PREFIX}"
        response = self.groq_client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            max_tokens=8,
        )
        reply = response.choices[0].message.content or ""
        return self._parse_price(reply)

    def price(self, description: str) -> float:
        if self.mode == "modal":
            self.log("Specialist Agent is calling remote fine-tuned model on Modal")
            result = self.pricer.price.remote(description)
            self.log(f"Specialist Agent completed - predicting ${result:.2f}")
            return result

        if self.mode == "groq":
            self.log(f"Specialist Agent is calling Groq ({self.model})")
            result = self._price_with_groq(description)
            self.log(f"Specialist Agent completed - predicting ${result:.2f}")
            return result

        return 0.0