import time
import logging
import os
import evaluate
from huggingface_hub import InferenceClient

logger = logging.getLogger(__name__)

# YOUR EXACT PROMPT
JUDGE_PROMPT = """Grade the system's answer.
Question: {q}
Correct answer: {correct}
System answer: {answer}

Reply with only PASS or FAIL.
PASS = the system answer correctly addresses the question with no major errors.
FAIL = the answer is wrong, missing, or contradicts the correct answer."""

class MetricsService:
    def __init__(self, model_name="gemini-1.5-flash"):
        self.model_name = model_name.lower()
        self.hf_token = os.environ.get("HF_TOKEN")
        
        # Setup as per your implementation
        self.client = InferenceClient(
            model="meta-llama/Llama-3.1-8B-Instruct",
            token=self.hf_token
        )
        self.bertscore = evaluate.load("bertscore")

        self.tiers = {
            "flash": {"input": 0.075, "output": 0.30},
            "pro": {"input": 3.50, "output": 10.50},
            "groq_large": {"input": 0.59, "output": 0.79},
            "groq_small": {"input": 0.05, "output": 0.08},
        }

    def warmup(self):
        """Proactively load BERTScore and check Judge connectivity"""
        logger.info("PRE-LOADING METRICS MODELS (BERTScore)...")
        try:
            # Dummy compute to force model download/load into memory/MPS
            self.bertscore.compute(
                predictions=["warmup"],
                references=["warmup"],
                lang="en"
            )
            logger.info("METRICS MODELS LOADED SUCCESSFULLY.")
        except Exception as e:
            logger.warning(f"Metrics warmup failed: {e}")

    def calculate_cost(self, input_tokens, output_tokens, model_override=None):
        m_name = (model_override or self.model_name).lower()
        if "70b" in m_name or "versatile" in m_name:
            prices = self.tiers["groq_large"]
        elif "8b" in m_name or "instant" in m_name:
            prices = self.tiers["groq_small"]
        elif "pro" in m_name:
            prices = self.tiers["pro"]
        else:
            prices = self.tiers["flash"]
        return ((input_tokens / 1_000_000) * prices["input"]) + ((output_tokens / 1_000_000) * prices["output"])

    def process_metrics(self, client, query, answer, context, usage_metadata, start_time, abstracts_list=None, ground_truth=None, model_name=None):
        """EXACT implementation with NO truncation."""
        # 1. Prepare "Correct Answer"
        # If ground_truth is provided (from benchmark), use it. Otherwise fallback to abstracts/context.
        final_ground_truth = ground_truth or ("\n\n".join(abstracts_list) if abstracts_list else context)
        
        # 2. LLM-as-a-Judge — strip citation markers and disclaimer sentences before judging
        import re as _re
        clean_for_judge = _re.sub(r'\[Paper \d+\]', '', answer).strip()
        # Strip sentences where the model hedges about missing context — these cause false judge FAILs
        _DISCLAIMER = _re.compile(
            r'[^.!?]*\b(not explicitly (mentioned|stated|provided)|'
            r'not (mentioned|provided|specified) in (the )?(provided |given )?context|'
            r'is not clear from|cannot be (determined|found) from|'
            r'no (explicit |direct )?mention)\b[^.!?]*[.!?]?',
            _re.IGNORECASE
        )
        clean_for_judge = _DISCLAIMER.sub('', clean_for_judge).strip()
        prompt = JUDGE_PROMPT.format(
            q=query,
            correct=final_ground_truth,
            answer=clean_for_judge
        )
        
        try:
            verdict = self.client.chat_completion(
                [{"role": "user", "content": prompt}],
                max_tokens=20, 
                temperature=0.0
            )
            judge_pass = "PASS" in verdict.choices[0].message.content.upper()
        except Exception as e:
            logger.error(f"Judge error: {e}")
            judge_pass = False

        # 3. BERTScore — reuse the already-cleaned answer from above
        clean_answer = clean_for_judge
        try:
            bert_results = self.bertscore.compute(
                predictions=[clean_answer],
                references=[final_ground_truth],
                lang="en",
                rescale_with_baseline=True
            )
            f1_score = float(bert_results["f1"][0])
        except Exception as e:
            logger.error(f"BERTScore error: {e}")
            f1_score = 0.0

        # 4. Latency and Cost
        end_time = time.time()
        input_tokens = usage_metadata.prompt_token_count if usage_metadata else 0
        output_tokens = usage_metadata.candidates_token_count if usage_metadata else 0
        
        return {
            "tokens": {"input": input_tokens, "output": output_tokens, "total": input_tokens + output_tokens},
            "latency_seconds": round(end_time - start_time, 2),
            "cost_usd": float(self.calculate_cost(input_tokens, output_tokens, model_override=model_name)),
            "accuracy": {
                "llm_judge": "PASS" if judge_pass else "FAIL",
                "bert_score": f1_score
            }
        }