from agentic_workflow.config import LLM_MODEL
from langchain_ollama import ChatOllama
import re
import json
from langchain_core.messages import HumanMessage, SystemMessage

llm = ChatOllama(model=LLM_MODEL, temperature=0)

def evaluate_response(query: str, response: str) -> dict:
    """
    Scores the draft response and decides whether to retry.
 
    Returns a dict with keys:
      - score:       int 1-10
      - pass:        bool  → True means use this response
      - reason:      str   → short explanation
      - suggestion:  str   → how to improve (if not passing)
    """
    system = """
        You are a strict response quality evaluator. Given a user query and a draft response, return ONLY valid JSON:
        - "score": integer 1-10 (10 = excellent)
        - "pass": true if score >= 7, false otherwise
        - "reason": one-sentence explanation
        - "suggestion": if not passing, a short instruction to improve the response; else ""
        
        Return ONLY the JSON object. No explanation.
        """
 
    response_eval = llm.invoke([
        SystemMessage(content=system),
        HumanMessage(content=f"Query: {query}\n\nDraft response: {response}")
    ])
    raw = response_eval.content.strip()

    raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw).strip()

    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError:
        parsed = {}

    return {
        "score": 5,
        "pass": False,
        "reason": "Evaluation failed",
        "suggestion": "Retry with more detail",
        **parsed
    }

# result = evaluate_response(
#     query="What is the refund policy?",
#     response="According to the context provided earlier, the cancellation procedure involves a notice period of 30 days. The premium paid by you will be returned on a pro-rata basis or 25% of the annual premium, whichever is higher, will be retained. Any cancellation request sent after 30 days of commencement of the policy will be refunded on a pro-rata basis."
# )
# print("Test 1 (good response):", result)

# result = evaluate_response(
#     query="What is the refund policy?",
#     response="You can get a refund if you contact support."
# )
# print("Test 3 (vague response):", result)