import numpy as np
from src.state import AgentState
from langchain_core.messages import HumanMessage, AIMessage

# Anchor sentences — representative of in-scope and out-of-scope topics
IN_SCOPE_ANCHORS = [
    "where is my order",
    "track my delivery",
    "return policy refund",
    "payment method billing",
    "late shipment compensation",
    "seller review rating",
    "order status cancelled",
    "how long does shipping take",
    "I want to cancel my order",
    "product not delivered yet",
    "here is my order id",
    "the order number is",
    "my order id",
    "order reference number",
    "136cce7faa42fdb2cefd53fdc79a6098",  # Example order ID format
]

OUT_OF_SCOPE_ANCHORS = [
    "who is the president of the United States",
    "write me a poem about love",
    "explain quantum physics",
    "what is the recipe for pasta",
    "tell me a joke",
    "history of the Roman Empire",
    "help me write code in Python",
    "what is the weather today",
    "recommend a movie to watch",
    "translate this text to French",
]

# Lazy globals — reuses the same model already loaded by rag_retriever
_scope_model = None
_in_scope_embeddings = None
_out_scope_embeddings = None

def _init_scope_model():
    global _scope_model, _in_scope_embeddings, _out_scope_embeddings
    if _scope_model is not None:
        return
    from sentence_transformers import SentenceTransformer
    # Same model already cached by warmup_models.py — no re-download
    _scope_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
    _in_scope_embeddings = _scope_model.encode(IN_SCOPE_ANCHORS, normalize_embeddings=True)
    _out_scope_embeddings = _scope_model.encode(OUT_OF_SCOPE_ANCHORS, normalize_embeddings=True)

def validate_scope(state: AgentState) -> AgentState:
    """
    Zero-shot semantic scope classification using sentence-transformer embeddings.
    Compares user input against in-scope and out-of-scope anchor embeddings via
    cosine similarity. No LLM API call needed — runs in ~5ms after warmup.
    """
    print(f"[scope_validator] Starting validation for: {state.get('user_input', '')[:50]}...")
    try:
        _init_scope_model()
        print("[scope_validator] Model initialized")
        user_input = state.get("user_input", "")

        # Encode query (normalized = cosine sim is just dot product)
        query_emb = _scope_model.encode([user_input], normalize_embeddings=True)[0]

        # Max cosine similarity to each anchor set
        in_scope_score = float(np.max(_in_scope_embeddings @ query_emb))
        out_scope_score = float(np.max(_out_scope_embeddings @ query_emb))

        # Check if input looks like an order ID (32-char hex)
        import re
        is_order_id = bool(re.search(r'\b[a-f0-9]{32}\b', user_input.lower()))
        
        # If it's an order ID or in_scope_score wins, consider it in scope
        is_in_scope = is_order_id or in_scope_score >= out_scope_score or in_scope_score > 0.30

        if not is_in_scope:
            final_response = (
                "I'm Olist's customer support assistant — I can help with orders, "
                "deliveries, returns, payments, and seller queries. "
                f"Your question doesn't seem related to Olist support "
                f"(confidence: {in_scope_score:.2f}). Could you ask me something "
                "about your Olist experience?"
            )
            print(f"[scope_validator] Out of scope detected (score: {in_scope_score:.2f})")
            return {
                "scope_score": in_scope_score,
                "scope_checked": True,
                "scope_valid": False,
                "final_response": final_response,
                "messages": [
                    HumanMessage(content=state.get("user_input", "")),
                    AIMessage(content=final_response)
                ]
            }
        else:
            print(f"[scope_validator] In scope (score: {in_scope_score:.2f})")
            return {
                "scope_score": in_scope_score,
                "scope_checked": True,
                "scope_valid": True
            }
    except Exception as e:
        # On any failure, default to in-scope (fail open — better UX)
        print(f"[scope_validator] Exception: {str(e)}")
        return {
            "scope_checked": False,
            "scope_valid": True,
            "error_log": state.get("error_log", []) + [f"[scope_validator] Error: {str(e)}"]
        }