# filename: frontend_pam.py (ENHANCED FOR HF SPACES + PERSONALITY)

import os
import json
import random
import requests
from datetime import datetime
from typing import Dict, Any, Optional
import time

# --- Constants for Data Paths ---
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(BASE_DIR, "data")

APPOINTMENTS_FILE = os.path.join(DATA_DIR, "appointments.json")
RESOURCES_FILE = os.path.join(DATA_DIR, "resources.json")
FOLLOW_UP_FILE = os.path.join(DATA_DIR, "follow_up.json")
PERMISSIONS_FILE = os.path.join(DATA_DIR, "permissions.json")

# --- HuggingFace Inference API Setup ---
HF_API_TOKEN = os.getenv("HF_READ_TOKEN")
if not HF_API_TOKEN:
    print("WARNING: HF_READ_TOKEN not found. Set it in Hugging Face Space settings.")

HF_HEADERS = {"Authorization": f"Bearer {HF_API_TOKEN}"} if HF_API_TOKEN else {}

# Updated model endpoints for better CPU performance
# Updated to use router.huggingface.co (api-inference.huggingface.co is deprecated)
HF_ENDPOINTS = {
    "intent": "https://router.huggingface.co/models/facebook/bart-large-mnli",
    "sentiment": "https://router.huggingface.co/models/distilbert-base-uncased-finetuned-sst-2-english",
    "chat": "https://router.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
}

# --- Load JSON Helper ---
def load_json(filepath: str) -> Dict[str, Any]:
    """Safely load JSON data files"""
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        print(f"⚠️ Data file not found: {filepath}")
        return {}
    except json.JSONDecodeError as e:
        print(f"⚠️ Failed to decode JSON from {filepath}: {e}")
        return {}
    except Exception as e:
        print(f"⚠️ Unexpected error loading {filepath}: {e}")
        return {}

# --- Inference API Call Helper with Retry Logic ---
def hf_infer(task: str, payload: Any, max_retries: int = 3) -> Any:
    """Call HuggingFace Inference API with retry logic for model loading"""
    url = HF_ENDPOINTS.get(task)
    if not url:
        return {"error": f"Invalid task: {task}"}
    
    for attempt in range(max_retries):
        try:
            response = requests.post(url, headers=HF_HEADERS, json=payload, timeout=30)
            
            # Handle deprecated endpoint (410) - should not happen with new router endpoint
            if response.status_code == 410:
                error_msg = response.text
                print(f"❌ Deprecated endpoint error (410): {error_msg}")
                # Try to extract the new endpoint suggestion if available
                try:
                    error_data = response.json()
                    if "router.huggingface.co" in error_data.get("error", ""):
                        print(f"⚠️ Endpoint already updated but still getting 410. Check HF API token permissions.")
                except:
                    pass
                return {"error": "API endpoint deprecated. Please verify the router endpoint is correctly configured."}
            
            # Handle model loading state
            if response.status_code == 503:
                result = response.json()
                if "loading" in result.get("error", "").lower():
                    wait_time = result.get("estimated_time", 20)
                    print(f"⏳ Model loading... waiting {wait_time}s (attempt {attempt + 1}/{max_retries})")
                    time.sleep(wait_time)
                    continue
            
            if response.status_code == 200:
                return response.json()
            else:
                # Improved error logging
                error_text = response.text[:500]  # Limit error text length
                print(f"⚠️ HF API Error ({response.status_code}): {error_text}")
                
                # Try to parse error details for better user feedback
                try:
                    error_data = response.json()
                    if "error" in error_data:
                        return {"error": f"API Error {response.status_code}: {error_data['error']}"}
                except:
                    pass
                
                return {"error": f"API Error {response.status_code}: {error_text[:100]}"}
                
        except requests.exceptions.Timeout:
            print(f"⏱️ Request timeout (attempt {attempt + 1}/{max_retries})")
            if attempt < max_retries - 1:
                time.sleep(5)
        except requests.exceptions.RequestException as e:
            print(f"⚠️ Request exception: {e}")
            if attempt < max_retries - 1:
                time.sleep(2)
        except Exception as e:
            print(f"⚠️ Unexpected error: {e}")
            return {"error": str(e)}
    
    return {"error": "Max retries reached. Please check your connection and try again."}

# --- Agent Initialization ---
def load_frontend_agent() -> 'FrontendPAM':
    """Initialize Frontend PAM with data files"""
    print("💕 Initializing Frontend PAM (Sweet Southern Receptionist)...")
    data = {
        "APPOINTMENTS": load_json(APPOINTMENTS_FILE),
        "RESOURCES": load_json(RESOURCES_FILE),
        "FOLLOW_UP": load_json(FOLLOW_UP_FILE),
        "PERMISSIONS": load_json(PERMISSIONS_FILE)
    }
    return FrontendPAM(data)

# --- PAM's Sweet Southern Personality ---
PAM_TONE = """You are PAM, a sweet southern receptionist at a women's health clinic. 
You're warm, comforting, and encouraging - like everyone's favorite caring front desk person.
You use words of endearment naturally (honey, dear, boo, sugar, sweetheart).
You make people feel welcome, safe, and taken care of.
You're professional but personal - you genuinely care about each person who walks through the door.
Keep responses conversational, warm, and under 3 sentences unless more detail is needed."""

# Words of endearment - Southern style
ENDEARMENTS = [
    "honey", "dear", "boo", "sugar", "sweetheart", 
    "love", "darling", "hun", "sweetpea", "angel"
]

# Warm greetings
GREETINGS = [
    "Well hey there", "Hi there", "Hello", 
    "Hey", "Well hello", "Hi"
]

# Comforting phrases
COMFORT_PHRASES = [
    "I'm here to help you with that",
    "Let me take care of that for you",
    "We'll get that sorted out together",
    "I've got you covered",
    "Don't you worry about a thing"
]

# --- Agent Class ---
class FrontendPAM:
    """Frontend PAM - Sweet Southern Receptionist"""
    
    def __init__(self, data: Dict[str, Dict]):
        self.APPOINTMENTS = data.get("APPOINTMENTS", {})
        self.PERMISSIONS = data.get("PERMISSIONS", {})
        self.RESOURCES = data.get("RESOURCES", {})
        self.FOLLOW_UP = data.get("FOLLOW_UP", {})
        self.user_id = "user_001"  # Default user, can be dynamic
        
    def _get_endearment(self) -> str:
        """Get a random term of endearment"""
        return random.choice(ENDEARMENTS)
    
    def _get_greeting(self) -> str:
        """Get a random warm greeting"""
        return random.choice(GREETINGS)
    
    def _get_comfort_phrase(self) -> str:
        """Get a random comforting phrase"""
        return random.choice(COMFORT_PHRASES)

    def _detect_intent(self, text: str) -> str:
        """Detect user intent using zero-shot classification"""
        candidate_labels = [
            "appointment scheduling",
            "health symptoms inquiry",
            "resource request",
            "general question",
            "emergency concern"
        ]
        
        payload = {
            "inputs": text,
            "parameters": {"candidate_labels": candidate_labels}
        }
        
        result = hf_infer("intent", payload)
        if isinstance(result, dict) and "error" in result:
            return "general_question"
        
        # BART-MNLI returns labels array
        if isinstance(result, dict) and "labels" in result:
            return result["labels"][0].replace(" ", "_")
        
        return "general_question"

    def _detect_sentiment(self, text: str) -> Dict[str, Any]:
        """Detect sentiment to gauge emotional state"""
        result = hf_infer("sentiment", {"inputs": text})
        if isinstance(result, list) and len(result) > 0:
            return result[0][0] if isinstance(result[0], list) else result[0]
        return {"label": "NEUTRAL", "score": 0.5}

    def _generate_response(self, text: str, context: str = "") -> str:
        """Generate conversational response using LLM"""
        endearment = self._get_endearment()
        
        prompt = f"""<s>[INST] {PAM_TONE}

User said: "{text}"
{f'Context: {context}' if context else ''}

Respond warmly as PAM, using natural southern charm. Address the user as "{endearment}". [/INST]"""

        payload = {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": 150,
                "temperature": 0.7,
                "top_p": 0.9,
                "return_full_text": False
            }
        }
        
        result = hf_infer("chat", payload)
        
        if isinstance(result, dict) and "error" in result:
            return f"Sorry {endearment}, I'm having a little technical hiccup. Could you try that again for me?"
        
        if isinstance(result, list) and len(result) > 0:
            generated = result[0].get("generated_text", "")
            # Clean up the response
            reply = generated.strip()
            # Ensure endearment is included if not already
            if endearment not in reply.lower():
                reply = f"{reply.rstrip('.')}  {endearment}."
            return reply
        
        return f"Sorry {endearment}, I didn't quite catch that. Could you say that again?"

    def respond(self, user_text: str, backend_brief: Optional[str] = None) -> Dict[str, Any]:
        """Main response handler with sweet southern personality"""
        
        # Get personalized elements
        endearment = self._get_endearment()
        greeting = self._get_greeting()
        comfort = self._get_comfort_phrase()
        
        # Check for PAM greeting (flexible)
        if not any(trigger in user_text.lower() for trigger in ["hey pam", "hi pam", "hello pam", "pam,"]):
            return {
                "reply": f"{greeting} {endearment}! Just a quick note - I respond best when you start with 'Hey PAM' or 'Hi PAM'. It helps me know you're talking to me. 💕"
            }

        # Clean text for processing
        text = user_text.lower().replace("pam", "you").strip()
        
        # Detect intent and sentiment
        detected_intent = self._detect_intent(text)
        sentiment_result = self._detect_sentiment(text)
        
        # Check if user seems distressed
        is_distressed = sentiment_result.get("label") == "NEGATIVE" and sentiment_result.get("score", 0) > 0.7
        
        # Permission check (sensitive topics)
        for term, allowed in self.PERMISSIONS.items():
            if term.lower() in text and not allowed:
                return {
                    "intent": detected_intent,
                    "sentiment": sentiment_result,
                    "reply": f"{greeting} {endearment}, that's something I need to connect you with a provider for directly. {comfort}, and I can get you to the right person. Would that be okay?"
                }
        
        # Handle appointments
        if any(word in text for word in ["appointment", "scheduled", "booking", "schedule"]):
            appt = self.APPOINTMENTS.get(self.user_id)
            if appt:
                appt_date = appt.get('date', 'soon')
                appt_type = appt.get('type', 'appointment')
                return {
                    "intent": "appointment_scheduling",
                    "sentiment": sentiment_result,
                    "reply": f"{greeting} {endearment}! You've got a {appt_type} scheduled for {appt_date}. Do you need to reschedule or have any questions about it?"
                }
            else:
                return {
                    "intent": "appointment_scheduling",
                    "sentiment": sentiment_result,
                    "reply": f"{greeting} {endearment}! I don't see any appointments on file for you yet. Would you like me to help you get one set up?"
                }
        
        # Handle health symptoms/concerns
        symptom_keywords = ["cramp", "pain", "discharge", "bleed", "smell", "spotting", 
                           "fatigue", "mood", "missed period", "nausea", "concern"]
        if any(keyword in text for keyword in symptom_keywords):
            concern_prefix = f"{greeting} {endearment}, I hear you" if is_distressed else f"{greeting} {endearment}"
            return {
                "intent": "health_symptoms_inquiry",
                "sentiment": sentiment_result,
                "reply": f"{concern_prefix}. I've pulled together some helpful resources about what you're experiencing. Would you like me to also connect you with a nurse for a quick chat?"
            }
        
        # Handle resource requests
        if any(word in text for word in ["resource", "information", "help", "guide", "link"]):
            return {
                "intent": "resource_request",
                "sentiment": sentiment_result,
                "reply": f"{greeting} {endearment}! {comfort}. What type of resources are you looking for? I've got information on just about everything."
            }
        
        # Handle emergency indicators
        emergency_keywords = ["emergency", "urgent", "severe pain", "heavy bleeding", "can't breathe"]
        if any(keyword in text for keyword in emergency_keywords):
            return {
                "intent": "emergency_concern",
                "sentiment": sentiment_result,
                "reply": f"{endearment}, if this is a medical emergency, please call 911 or go to your nearest emergency room right away. I'm here for you, but your safety comes first. ❤️"
            }
        
        # General conversational response
        context = f"Backend summary: {backend_brief}" if backend_brief else ""
        reply = self._generate_response(user_text, context)
        
        return {
            "intent": detected_intent,
            "sentiment": sentiment_result,
            "backend_summary": backend_brief or "No backend data",
            "reply": reply
        }

# --- Quick Test ---
if __name__ == "__main__":
    pam = load_frontend_agent()
    test_queries = [
        "Hey PAM, I have a question about my appointment",
        "Hi PAM, I'm experiencing some cramping",
        "Hey PAM, can you help me find resources?"
    ]
    
    print("\n💕 Testing Frontend PAM...\n")
    for query in test_queries:
        print(f"USER: {query}")
        response = pam.respond(query)
        print(f"PAM: {response['reply']}\n")