import logging
import json
import hashlib
import redis
from datetime import datetime, timedelta, timezone
from dateutil.relativedelta import relativedelta
from app.services.infra.database import (
    fetch_tenant_limits, 
    get_catalogue_product_count,
    get_ticket_count, 
    get_vector_db_size, 
    get_total_token_usage
)
# Firestore is imported at the point of use, not here. It builds a client on
# import and so needs Google credentials to exist, yet it serves exactly two
# metrics -- ai_conversations and ai_responses -- both of them Aivora's. A
# service that only caps product counts would otherwise have to carry the
# chatbot's credentials purely to load this module.
from app.core.llm_client import client as openai_client
from app.core.config import settings
from app.core.prompts import QUOTA_MAPPING_PROMPT
import httpx
import asyncio
import jwt
import time
from typing import Optional

logger = logging.getLogger("quotas")

def _quota_redis_enabled():
    env_name = (getattr(settings, "ENV", "") or "").strip().lower()
    if env_name in {"local", "dev", "development", "test"}:
        return False
    return bool(getattr(settings, "QUOTA_REDIS_CACHE_ENABLED", True))

# Initialize Redis Client
redis_client = None
if _quota_redis_enabled():
    try:
        redis_client = redis.Redis(
            host=settings.REDIS_HOST,
            port=settings.REDIS_PORT,
            username=settings.REDIS_USER,
            password=settings.REDIS_PASSWORD,
            ssl=settings.REDIS_SSL,
            decode_responses=True,
            socket_timeout=5
        )
        logger.info("Redis client for quotas initialized successfully.")
    except Exception as e:
        logger.error(f"Failed to initialize Redis client: {e}")
else:
    logger.info("Redis quota cache disabled for this environment.")

class QuotaExceededError(Exception):
    """Exception raised when a tenant has exceeded its resource quota."""
    def __init__(self, message: str):
        self.message = message
        super().__init__(self.message)

class QuotaValidationError(QuotaExceededError):
    """Exception raised when quota limits or usage cannot be validated."""
    pass

CORE_REQUIRED_LIMIT_KEYS = {
    "ai_tokens",
    "total_tickets",
    "storage_vector_capacity",
    "ai_conversations",
}
OPTIONAL_LIMIT_DEFAULTS = {
    "ai_responses": -1,
}
REQUIRED_LIMIT_KEYS = CORE_REQUIRED_LIMIT_KEYS.union(OPTIONAL_LIMIT_DEFAULTS.keys())

# Every product on a plan measures something different, so each is validated
# against its own keys rather than Aivora's. BrandForge caps how many products
# a catalogue may hold and does not meter tokens at all -- its plans carry a
# Product count and no token feature, so demanding Aivora's four keys of it
# would reject a perfectly valid plan.
AIVORA = "Aivora"
BRANDFORGE = "BrandForge"

REQUIRED_BY_PRODUCT = {
    AIVORA: CORE_REQUIRED_LIMIT_KEYS,
    BRANDFORGE: {"product_count"},
}
OPTIONAL_BY_PRODUCT = {
    AIVORA: OPTIONAL_LIMIT_DEFAULTS,
    BRANDFORGE: {},
}

# A metric the product does not define is not a refusal. BrandForge has no
# token allowance, so a catalogue build asking about tokens is asking about
# something its plan never limited -- that is unmetered, not forbidden.
UNMETERED_IS_ALLOWED = {BRANDFORGE}

QUOTA_VALIDATION_MESSAGE = "We’re unable to validate your service limits right now. Please try again shortly."
QUOTA_CACHE_VERSION = "v2"

FEATURE_KEYWORDS = {
    # Listed before the token entry: a feature named "Product count" must not
    # be read as anything else, and dict order decides which keyword wins.
    "product_count": ("product count", "product_count", "products", "catalogue",
                      "catalog", "sku", "skus"),
    "ai_tokens": ("token", "tokens", "character", "characters"),
    "total_tickets": ("ticket", "tickets", "support", "helpdesk", "customer query", "customer queries"),
    "storage_vector_capacity": ("storage", "vector", "knowledge", "kb", "database", "file", "files", "indexing"),
    "ai_conversations": ("conversation", "conversations", "thread", "threads", "chat session", "chat sessions"),
    "ai_responses": ("response", "responses", "assistant message", "assistant messages", "ai reply", "ai replies"),
}

def _coerce_limit(metric_type: str, raw_limit):
    if raw_limit == -1 or raw_limit == "-1":
        return -1

    if raw_limit is None:
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

    if isinstance(raw_limit, bool):
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

    if isinstance(raw_limit, str):
        cleaned = raw_limit.strip().replace(",", "")
        if cleaned.lower() == "unlimited":
            return -1
        if not cleaned.isdigit():
            raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)
        raw_limit = cleaned

    try:
        limit = int(raw_limit)
    except (TypeError, ValueError):
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

    if limit < -1:
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

    return limit

def _validate_mapped_limits(mapping, product_name=AIVORA):
    if not isinstance(mapping, dict):
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

    required = REQUIRED_BY_PRODUCT.get(product_name, CORE_REQUIRED_LIMIT_KEYS)
    optional = OPTIONAL_BY_PRODUCT.get(product_name, OPTIONAL_LIMIT_DEFAULTS)

    missing = required.difference(mapping.keys())
    if missing:
        logger.error(f"Quota mapping for {product_name} missing required keys: "
                     f"{sorted(missing)}")
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

    mapping = {**optional, **mapping}
    keys = required.union(optional.keys())
    return {key: _coerce_limit(key, mapping.get(key)) for key in keys}

def _extract_feature_text(feature):
    if not isinstance(feature, dict):
        return ""

    text_parts = []
    for key in ("name", "featureName", "title", "label", "description", "key", "resourceName"):
        value = feature.get(key)
        if value:
            text_parts.append(str(value))
    return " ".join(text_parts).lower()

def _extract_capping_value(feature):
    if not isinstance(feature, dict):
        return None

    for key in ("cappingValue", "capping_value", "limit", "value", "quota", "cap"):
        if key in feature:
            return feature.get(key)
    return None

def _map_limits_from_features(product_features, product_name=AIVORA):
    required = REQUIRED_BY_PRODUCT.get(product_name, CORE_REQUIRED_LIMIT_KEYS)
    optional = OPTIONAL_BY_PRODUCT.get(product_name, OPTIONAL_LIMIT_DEFAULTS)
    mapping = {key: None for key in required.union(optional.keys())}

    for feature in product_features:
        feature_text = _extract_feature_text(feature)
        raw_limit = _extract_capping_value(feature)
        if raw_limit is None:
            continue

        for metric_type, keywords in FEATURE_KEYWORDS.items():
            if metric_type not in mapping:
                continue
            if any(keyword in feature_text for keyword in keywords):
                mapping[metric_type] = _coerce_limit(metric_type, raw_limit)
                break

    if all(mapping.get(key) is not None for key in required):
        for key, default in optional.items():
            if mapping.get(key) is None:
                mapping[key] = default
        return mapping

    missing = [key for key in required if mapping.get(key) is None]
    logger.warning(f"Could not deterministically map all {product_name} quota "
                   f"limits. Missing: {missing}")
    return None

def get_mapped_limits(product_features, product_name=AIVORA):
    """
    Maps one product's plan features to internal quota categories, using an LLM
    where the keywords cannot settle it. Caches in Redis for 10 minutes.

    The product is a parameter because each one measures something different:
    Aivora counts tokens and conversations, BrandForge counts products. Reading
    Aivora's numbers for a BrandForge tenant is how a customer who never bought
    Aivora ends up unable to build anything.
    """
    if not product_features or not isinstance(product_features, list):
        return {}

    product_data = next((p for p in product_features
                         if p.get("productName") == product_name), None)
    if not product_data or not product_data.get("features"):
        return {}

    product_plan_features = product_data["features"]

    # 2. Check Cache for LLM mapping
    features_json = json.dumps(product_plan_features, sort_keys=True)
    features_hash = hashlib.md5(features_json.encode()).hexdigest()
    # Keyed by product: two products on one plan would otherwise share an entry
    # and serve each other's limits.
    cache_key = f"quota_map:{QUOTA_CACHE_VERSION}:{product_name}:{features_hash}"
    
    if redis_client:
        try:
            cached_val = redis_client.get(cache_key)
            if cached_val:
                logger.debug("Found quota mapping in Redis cache.")
                return _validate_mapped_limits(json.loads(cached_val),
                                               product_name)
        except Exception as e:
            logger.warning(f"Error reading from Redis cache: {e}")

    deterministic_mapping = _map_limits_from_features(product_plan_features,
                                                      product_name)
    if deterministic_mapping is None and product_name in UNMETERED_IS_ALLOWED:
        # No LLM fallback for these. BrandForge's features are a short, literal
        # list -- if "Product count" is not among them the plan simply has no
        # cap, and asking a model to invent one would either hallucinate a
        # number or raise, which is how a plan that merely omits a limit ends
        # up refusing every request.
        logger.info(f"No {product_name} limits found in the plan; unmetered.")
        return {}
    if deterministic_mapping:
        logger.info("Successfully extracted quota limits deterministically from product features.")
        if redis_client:
            try:
                redis_client.setex(cache_key, 600, json.dumps(deterministic_mapping))
                logger.info("Stored deterministic quota mapping in Redis cache (10 min TTL).")
            except Exception as e:
                logger.warning(f"Error writing deterministic quota mapping to Redis cache: {e}")
        return deterministic_mapping

    # 3. LLM Mapping
    logger.info("Quota mapping cache miss. Triggering LLM mapping...")
    try:
        features_list_str = json.dumps(product_plan_features, indent=2)
        prompt = QUOTA_MAPPING_PROMPT.format(features_list=features_list_str)
        
        response = openai_client.chat.completions.create(
            model=settings.LLM_MODEL,
            messages=[{"role": "user", "content": prompt}],
            response_format={"type": "json_object"}
        )
        
        mapping = _validate_mapped_limits(
            json.loads(response.choices[0].message.content), product_name)
        logger.info(f"Successfully extracted quota limits via LLM: {json.dumps(mapping)}")
        
        # 4. Store in Cache (10 minutes TTL)
        if redis_client:
            try:
                redis_client.setex(cache_key, 600, json.dumps(mapping))
                logger.info("Stored new quota mapping in Redis cache (10 min TTL).")
            except Exception as e:
                logger.warning(f"Error writing to Redis cache: {e}")
        
        return mapping
    except Exception as e:
        logger.error(f"Failed to map quotas via LLM: {e}", exc_info=True)
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

def calculate_anniversary_start(baseline: datetime):
    """Calculates the start of the current monthly cycle relative to the baseline date."""
    # Ensure baseline is naive UTC for consistent comparison with datetime.utcnow()
    if baseline and baseline.tzinfo is not None:
        baseline = baseline.astimezone(timezone.utc).replace(tzinfo=None)
        
    now = datetime.utcnow()
    # Total months between the baseline and now
    months_diff = (now.year - baseline.year) * 12 + now.month - baseline.month
    
    # Calculate the candidate start date for this month
    candidate = baseline + relativedelta(months=months_diff)
    
    # If the candidate date is in the future, the cycle started one month earlier
    if candidate > now:
        return baseline + relativedelta(months=months_diff - 1)
    
    return candidate

def estimate_text_tokens(text: str) -> int:
    """Conservative local estimate used for quota preflight before provider usage is known."""
    if not text:
        return 1
    return max(1, len(text) // 4)

def get_quota_start_date(tenant_id: str, current_features: list, db_updated_at: datetime):
    """
    Retrieves the start date for the current monthly quota cycle.
    Logic:
    1. Baseline is db_updated_at.
    2. If productFeatures content changes OR db_updated_at changes, we reset the baseline to 'now'.
    3. The current start_date is the monthly anniversary of that baseline.
    """
    # Normalize db_updated_at to naive UTC for internal consistency
    if db_updated_at and db_updated_at.tzinfo is not None:
        db_updated_at = db_updated_at.astimezone(timezone.utc).replace(tzinfo=None)

    if not redis_client:
        baseline = db_updated_at or datetime.utcnow()
        return calculate_anniversary_start(baseline)

    state_key = f"quota_state_v3:{tenant_id}"
    now = datetime.utcnow()
    
    # 1. Calculate Hash of current features to detect content changes
    features_hash = hashlib.md5(json.dumps(current_features, sort_keys=True).encode()).hexdigest()
    
    try:
        cached_state = redis_client.get(state_key)
        if cached_state:
            state = json.loads(cached_state)
            baseline = datetime.fromisoformat(state["baseline_date"])
            old_hash = state.get("features_hash")
            old_db_updated = datetime.fromisoformat(state.get("db_updated_at")) if state.get("db_updated_at") else None
            
            reset_required = False
            # 2a. Reset if content changed
            if old_hash != features_hash:
                logger.info(f"Quota feature content changed for {tenant_id}. Resetting baseline to NOW.")
                baseline = now
                reset_required = True
            # 2b. Reset if DB timestamp changed (Plan was explicitly updated/switched)
            elif old_db_updated and db_updated_at and db_updated_at > old_db_updated:
                logger.info(f"Plan updatedAt changed in DB for {tenant_id}. Resetting baseline to {db_updated_at}.")
                baseline = db_updated_at
                reset_required = True
                
            if reset_required:
                # Update Redis state with new baseline
                new_state = {
                    "baseline_date": baseline.isoformat(),
                    "features_hash": features_hash,
                    "db_updated_at": db_updated_at.isoformat() if db_updated_at else None,
                    "updated_at": now.isoformat()
                }
                redis_client.set(state_key, json.dumps(new_state))
        else:
            # 3. New tenant state initialization
            logger.info(f"Initializing quota baseline for tenant {tenant_id}")
            baseline = db_updated_at or now
            new_state = {
                "baseline_date": baseline.isoformat(),
                "features_hash": features_hash,
                "db_updated_at": db_updated_at.isoformat() if db_updated_at else None,
                "updated_at": now.isoformat()
            }
            redis_client.set(state_key, json.dumps(new_state))

        # 4. Calculate current start date relative to baseline
        return calculate_anniversary_start(baseline)
        
    except Exception as e:
        logger.error(f"Error managing quota state in Redis for {tenant_id}: {e}")
        baseline = db_updated_at or now
        return calculate_anniversary_start(baseline)

def check_quota(tenant_id: str, metric_type: str, user_id: Optional[str] = None,
                requested_amount: int = 1, product_name: str = AIVORA):
    """
    Checks the current usage against the monthly (or overall) limits.

    product_name selects whose numbers apply. Catalogue work passes BrandForge,
    which caps products and meters nothing else; the chat passes Aivora, which
    meters tokens, conversations and tickets.
    """
    logger.info(f"Checking {metric_type} quota for tenant: {tenant_id} "
                f"(User: {user_id}, Product: {product_name})")

    # Get raw features and updatedAt from master DB
    product_features, db_updated_at = fetch_tenant_limits(tenant_id)
    if not product_features:
        if product_name in UNMETERED_IS_ALLOWED:
            # Same reasoning as an absent cap: no plan data is a limit nobody
            # recorded, not a limit of zero. Blocking here would make catalogue
            # ingestion depend on the control plane being reachable.
            logger.info(f"No plan features for {tenant_id}; {metric_type} is "
                        f"unmetered under {product_name}.")
            return
        logger.error(f"No quota features found for tenant {tenant_id}. Blocking request because limits cannot be validated.")
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

    limits = get_mapped_limits(product_features, product_name)
    if not limits:
        if product_name in UNMETERED_IS_ALLOWED:
            # The plan does not carry this product, or carries it without a
            # cap. That is a limit nobody configured rather than a limit of
            # zero, and refusing here would stop tenants who work today.
            logger.info(f"No {product_name} limits on {tenant_id}'s plan; "
                        f"{metric_type} is unmetered.")
            return
        logger.error(f"No quota mapping generated for tenant {tenant_id}. Blocking request because limits cannot be validated.")
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

    # Handle "unlimited" (-1)
    limit = limits.get(metric_type)
    if limit is None and product_name in UNMETERED_IS_ALLOWED:
        # BrandForge has no token allowance at all, so a build asking about
        # tokens is asking about something its plan never limited.
        logger.info(f"{product_name} does not meter {metric_type}; allowing.")
        return
    limit = _coerce_limit(metric_type, limit)
    if limit == -1:
        logger.debug(f"Quota for {metric_type} is unlimited.")
        return

    # Get the monthly reset start date relative to the updatedAt/Content changed baseline
    start_date = get_quota_start_date(tenant_id, product_features, db_updated_at)
    
    # Usage Calculation
    quota_hit = False
    usage = 0

    try:
        if metric_type == "ai_tokens":
            usage = get_total_token_usage(tenant_id, start_date=start_date)
            if usage + max(0, requested_amount) > limit: quota_hit = True
                    
        elif metric_type == "total_tickets":
            usage = get_ticket_count(tenant_id, start_date=start_date)
            if usage + max(0, requested_amount) > limit: quota_hit = True
                    
        elif metric_type == "storage_vector_capacity":
            # OVERALL Limit - no start_date
            usage_bytes = get_vector_db_size(tenant_id)
            limit_bytes = limit * 1024 * 1024 # Convert MB from plan label to bytes
            usage = usage_bytes
            if usage_bytes + max(0, requested_amount) > limit_bytes: quota_hit = True
                    
        elif metric_type == "ai_conversations":
            from app.services.integrations.firestore import (
                get_conversation_thread_count,
            )
            usage = get_conversation_thread_count(tenant_id, start_date=start_date)
            if usage + max(0, requested_amount) > limit: quota_hit = True

        elif metric_type == "ai_responses":
            from app.services.integrations.firestore import (
                get_ai_response_message_count,
            )
            usage = get_ai_response_message_count(tenant_id, start_date=start_date)
            if usage + max(0, requested_amount) > limit: quota_hit = True

        elif metric_type == "product_count":
            # A total, not a monthly allowance: a merchant owns 100 products,
            # they do not get 100 more each month. No start_date, like storage.
            usage = get_catalogue_product_count(tenant_id)
            # Only growth is refused. A catalogue already past its cap -- a plan
            # downgrade, or a limit introduced after the fact -- keeps working
            # and keeps refreshing; it simply cannot get bigger. Refusing a
            # request that adds nothing would freeze a merchant's prices and
            # stock the moment they reached their limit.
            adding = max(0, requested_amount)
            if adding and usage + adding > limit: quota_hit = True

        else:
            logger.error(f"Unknown quota metric requested: {metric_type}")
            raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)
    except QuotaValidationError:
        raise
    except Exception as e:
        logger.error(f"Could not calculate quota usage for {metric_type} in tenant {tenant_id}: {e}", exc_info=True)
        raise QuotaValidationError(QUOTA_VALIDATION_MESSAGE)

    if quota_hit:
        logger.warning(f"Quota exceeded for {metric_type} (Tenant: {tenant_id}, Usage: {usage}, Limit: {limit})")
        send_quota_notification(tenant_id, metric_type, user_id)
        
        # Metric-Specific Error Messages
        # Mapping the metric_type to user-friendly messages that explain the issue and next steps without exposing internal details.
        msg_map = {
            "ai_tokens": "We’re unable to process your request right now. Please try again shortly.",
            "total_tickets": "We couldn’t create your support request at the moment. Please contact our team through the Contact Us page.",
            "storage_vector_capacity": "That information isn’t available at the moment. Please get in touch with our support team for further assistance.",
            "ai_conversations": "This service is currently unavailable. Please try again in a little while.",
            "ai_responses": "You’ve hit your AI response quota. Upgrade or wait for the reset to continue.",
            # Says the number, unlike the others: this one is the merchant's
            # own doing and they can act on it by removing products or
            # upgrading. "Unavailable" would just look broken.
            "product_count": f"Your plan allows {limit} products and this would "
                             f"take you past that. Remove some products or "
                             f"upgrade your plan to add more.",
        }
        err_msg = msg_map.get(metric_type, "You’ve hit your resource quota. Upgrade or wait for the reset to continue.")
        raise QuotaExceededError(err_msg)

    logger.info(f"Quota check passed for {metric_type} (Tenant: {tenant_id}, Usage: {usage}/{limit}, Start Date: {start_date})")

async def _do_notify(payload):
    # Generate system token
    token = generate_system_token()
    
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }
    
    async with httpx.AsyncClient(timeout=10.0) as client:
        try:
            r = await client.post(settings.NOTIFICATION_API_URL, json=payload, headers=headers)
            logger.info(f"Quota notification API result: {r.status_code}")
        except Exception as e:
            logger.error(f"Failed to send quota notification: {e}")

def generate_system_token(expires_in_sec: int = 3600) -> str:
    """
    Generates a system JWT for authenticating with the notification service.
    """
    if not settings.SYSTEM_JWT_KEY:
        logger.warning("SYSTEM_JWT_KEY not configured. Notification may fail.")
        return ""
        
    now = int(time.time())
    payload = {
        "type": "system",
        "service": settings.SYSTEM_SERVICE_NAME,
        "iat": now,
        "exp": now + expires_in_sec,
    }
    token = jwt.encode(payload, settings.SYSTEM_JWT_KEY, algorithm="HS256")
    return token

def send_quota_notification(tenant_id, metric_type, user_id):
    if not user_id:
        logger.debug("No user_id provided for notification alert. Skipping.")
        return

    # Notification Message Mapping
    msg_map = {
        "ai_tokens": ("Alert 🚨 Aivora AI Token Limit", "AI token quota approaching limit — chatbot responses will be paused soon."),
        "total_tickets": ("Alert 🚨 Aivora Ticket Count Limit", "Ticket count limit reached — new submissions are on hold."),
        "storage_vector_capacity": ("Alert 🚨 Aivora DB Storage", "Vector DB Storage at capacity — new training uploads disabled."),
        "ai_conversations": ("Alert 🚨 Aivora Conversation Count Limit", "Conversation limit has exhausted — Live chat temporarily paused."),
        "ai_responses": ("Alert 🚨 Aivora AI Response Limit", "Aivora AI response limit hit. Upgrade or wait for reset.")
    }
    
    title, desc = msg_map.get(metric_type, ("Quota Limit hit", "A resource limit has been reached."))
    
    payload = {
        "title": title,
        "description": f"<p>{desc}</p>",
        "isSelectedAll": True,
        "userIds": [user_id]
    }
    
    logger.info(f"Triggering quota notification for {metric_type} to user {user_id}")
    
    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            loop.create_task(_do_notify(payload))
        else:
            asyncio.run(_do_notify(payload))
    except Exception as e:
        logger.error(f"Error orchestrating notification task: {e}")