import os
import time
import logging
from functools import wraps
import numpy as np
import firebase_admin
from firebase_admin import credentials, firestore
from datetime import datetime, timezone, timedelta

# Configure logger
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Enforce IST timezone globally
IST = timezone(timedelta(hours=5, minutes=30))

db = None
_embeddings_cache = {}  # { division: (timestamp, profiles_dict) }
_global_embeddings = {}  # { roll_no: { name, division, embedding } }
_global_embeddings_timestamp = 0.0
CACHE_TTL = 300  # 5 minutes in seconds

def get_firestore_db():
    """Initializes and returns the Firestore client singleton."""
    global db
    if db is not None:
        return db
    
    if not firebase_admin._apps:
        try:
            # Try to build credential dict from env variables
            private_key = os.getenv("FIREBASE_PRIVATE_KEY")
            if private_key:
                # Replace literal escaped newlines with actual newlines first
                private_key = private_key.replace('\\n', '\n')
                
                # Foolproof PEM reconstruction: extract base64 payload and clean it
                header = "-----BEGIN PRIVATE KEY-----"
                footer = "-----END PRIVATE KEY-----"
                
                start_idx = private_key.upper().find(header)
                end_idx = private_key.upper().find(footer)
                
                if start_idx != -1 and end_idx != -1:
                    base64_body = private_key[start_idx + len(header) : end_idx]
                    # First split on whitespace
                    base64_body = "".join(base64_body.split())
                    
                    # Find and log any non-base64 characters found in the payload
                    invalid_chars = [c for c in base64_body if c not in 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=']
                    if invalid_chars:
                        logger.warning(f"Non-base64 characters stripped from body: {invalid_chars!r}")
                    
                    # Keep ONLY valid base64 characters to guarantee a clean PEM payload
                    base64_body = "".join([c for c in base64_body if c in 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/='])
                    
                    # Format to standard 64-character lines
                    chunks = [base64_body[i:i+64] for i in range(0, len(base64_body), 64)]
                    private_key = header + "\n" + "\n".join(chunks) + "\n" + footer
                else:
                    # Fallback to general cleaning if headers are missing
                    logger.warning("PEM headers not found! Falling back to standard cleaning...")
                    private_key = private_key.replace('\r', '')
                    lines = [line.strip() for line in private_key.split('\n') if line.strip()]
                    private_key = '\n'.join(lines)
            
            project_id = os.getenv("FIREBASE_PROJECT_ID")
            client_email = os.getenv("FIREBASE_CLIENT_EMAIL")
            
            if project_id and client_email and private_key:
                cred_dict = {
                    "type": "service_account",
                    "project_id": project_id,
                    "private_key_id": os.getenv("FIREBASE_PRIVATE_KEY_ID"),
                    "private_key": private_key,
                    "client_email": client_email,
                    "client_id": os.getenv("FIREBASE_CLIENT_ID"),
                    "auth_uri": os.getenv("FIREBASE_AUTH_URI", "https://accounts.google.com/o/oauth2/auth"),
                    "token_uri": os.getenv("FIREBASE_TOKEN_URI", "https://oauth2.googleapis.com/token"),
                }
                cred = credentials.Certificate(cred_dict)
                firebase_admin.initialize_app(cred)
                logger.info("Firebase initialized using environment credentials.")
            else:
                # Log what was missing from environment
                missing = []
                if not project_id: missing.append("FIREBASE_PROJECT_ID")
                if not client_email: missing.append("FIREBASE_CLIENT_EMAIL")
                if not private_key: missing.append("FIREBASE_PRIVATE_KEY")
                if missing:
                    logger.warning(f"Missing environment variables for Firebase initialization: {', '.join(missing)}")
                
                # Fallback to local credentials file
                cred_file = os.getenv("FIREBASE_CREDENTIALS_PATH", "firebase_credentials.json")
                if os.path.exists(cred_file):
                    cred = credentials.Certificate(cred_file)
                    firebase_admin.initialize_app(cred)
                    logger.info(f"Firebase initialized using credentials file: {cred_file}")
                else:
                    logger.warning(f"Credentials file {cred_file} not found. Attempting default initialization...")
                    firebase_admin.initialize_app()
                    logger.info("Firebase initialized with default credentials.")
        except Exception as e:
            logger.error(f"Failed to initialize Firebase Admin SDK: {e}")
            raise e
            
    db = firestore.client()
    return db

def db_retry(attempts=3, delay=1):
    """Decorator to retry Firestore operations with linear backoff."""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for i in range(attempts):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    logger.warning(f"Database error in {func.__name__} (attempt {i+1}/{attempts}): {e}")
                    if i < attempts - 1:
                        time.sleep(delay)
            logger.error(f"Database operation {func.__name__} failed after {attempts} attempts.")
            return None
        return wrapper
    return decorator

@db_retry(attempts=3, delay=1)
def check_firebase_connection():
    """Checks the database connection by requesting camera status metadata."""
    client = get_firestore_db()
    # Simple query to verify connection is alive
    client.collection("camera").document("status").get(timeout=3)
    return True

@db_retry(attempts=3, delay=1)
def load_all_embeddings(force_refresh=False):
    """Loads all student profiles with face embeddings from Firestore in-memory.
    Caches them for 5 minutes (300 seconds)."""
    global _global_embeddings, _global_embeddings_timestamp
    now = time.time()
    if not force_refresh and _global_embeddings and (now - _global_embeddings_timestamp < CACHE_TTL):
        logger.info("Returning cached global student embeddings.")
        return _global_embeddings

    client = get_firestore_db()
    # Query all students where embeddingStatus == complete
    docs = client.collection("students").where("embeddingStatus", "==", "complete").stream()
    profiles = {}
    for doc in docs:
        data = doc.to_dict()
        roll_no = doc.id
        embedding = data.get("faceEmbedding")
        name = data.get("name")
        div = data.get("division", "Unknown")
        if embedding:
            profiles[roll_no] = {
                "name": name,
                "division": div,
                "embedding": np.array(embedding, dtype=np.float32)
            }
            
    _global_embeddings = profiles
    _global_embeddings_timestamp = now
    logger.info(f"Loaded and cached {len(profiles)} global student embeddings.")
    return profiles

@db_retry(attempts=3, delay=1)
def get_all_student_embeddings(division):
    """Fetches students in a division from the cached global embeddings."""
    all_profiles = load_all_embeddings()
    if not all_profiles:
        return {}
    filtered = {
        roll_no: {
            "name": data["name"],
            "embedding": data["embedding"]
        }
        for roll_no, data in all_profiles.items()
        if data["division"] == division
    }
    logger.info(f"Filtered {len(filtered)} embeddings for division {division}")
    return filtered

@db_retry(attempts=3, delay=1)
def save_attendance(student_id, name, period, subject, division, date=None, time_val=None):
    """Saves attendance details for a student. Enforces IST timestamp."""
    client = get_firestore_db()
    now_ist = datetime.now(IST)
    
    if not date:
        date = now_ist.strftime("%Y-%m-%d")
        
    timestamp_val = now_ist
    
    # Format period to period-N format
    if isinstance(period, int) or (isinstance(period, str) and period.isdigit()):
        period_id = f"period-{period}"
    else:
        period_id = period
        
    doc_ref = client.collection("attendance") \
                    .document(date) \
                    .collection(division) \
                    .document(period_id) \
                    .collection("students") \
                    .document(student_id)
                    
    doc_ref.set({
        "name": name,
        "rollNo": student_id,
        "status": "present",
        "time": timestamp_val,
        "method": "AI-Camera"
    }, merge=True)
    logger.info(f"Saved attendance for {student_id} on {date} for {period_id}")
    return True

@db_retry(attempts=3, delay=1)
def update_enrollment_status(roll_no, status):
    """Updates the status of enrollment queue for a roll number."""
    client = get_firestore_db()
    now_ist = datetime.now(IST)
    doc_ref = client.collection("enrollmentQueue").document(roll_no)
    
    doc_ref.set({
        "status": status,
        "processedAt": now_ist
    }, merge=True)
    logger.info(f"Updated enrollment status for {roll_no} to {status}")
    return True

@db_retry(attempts=3, delay=1)
def save_face_embedding(roll_no, embedding, photo_count):
    """Saves the generated averaged face embedding to a student profile."""
    client = get_firestore_db()
    now_ist = datetime.now(IST)
    doc_ref = client.collection("students").document(roll_no)
    
    doc_ref.set({
        "faceEmbedding": embedding,
        "photoCount": photo_count,
        "embeddingStatus": "complete",
        "updatedAt": now_ist
    }, merge=True)
    logger.info(f"Saved face embedding for student {roll_no}")
    return True

@db_retry(attempts=3, delay=1)
def update_camera_status(is_online, current_period, division):
    """Updates the camera online/offline and active scan status."""
    client = get_firestore_db()
    now_ist = datetime.now(IST)
    doc_ref = client.collection("camera").document("status")
    
    if isinstance(current_period, int) or (isinstance(current_period, str) and current_period.isdigit()):
        period_str = f"period-{current_period}"
    else:
        period_str = current_period
        
    doc_ref.set({
        "isOnline": is_online,
        "lastScan": now_ist,
        "currentPeriod": period_str,
        "division": division
    }, merge=True)
    logger.info(f"Updated camera status: Online={is_online}, Period={period_str}, Div={division}")
    return True

@db_retry(attempts=3, delay=1)
def get_period_info():
    """Reads period and class division info from the camera status doc."""
    client = get_firestore_db()
    doc_ref = client.collection("camera").document("status")
    doc_snap = doc_ref.get()
    if doc_snap.exists:
        data = doc_snap.to_dict()
        return {
            "currentPeriod": data.get("currentPeriod", "period-1"),
            "division": data.get("division", "BCA-A")
        }
    return None

@db_retry(attempts=3, delay=1)
def delete_student_embedding(roll_no):
    """Deletes face embedding fields from student document and deletes queue item."""
    client = get_firestore_db()
    now_ist = datetime.now(IST)
    doc_ref = client.collection("students").document(roll_no)
    
    doc_ref.update({
        "faceEmbedding": firestore.DELETE_FIELD,
        "photoCount": firestore.DELETE_FIELD,
        "embeddingStatus": firestore.DELETE_FIELD,
        "updatedAt": now_ist
    })
    
    # Clean up enrollmentQueue if exists
    queue_ref = client.collection("enrollmentQueue").document(roll_no)
    if queue_ref.get().exists:
        queue_ref.delete()
        
    logger.info(f"Deleted face embedding fields for {roll_no}")
    return True