"""
InALign Usage Limiter

Tracks and enforces usage limits per client based on their plan.
"""

import os
from datetime import datetime, timezone
from typing import Optional, Tuple
from dataclasses import dataclass

# Plan limits
PLAN_LIMITS = {
    "starter": {
        "actions_per_month": 1000,
        "retention_days": 7,
        "max_agents": 1,
    },
    "pro": {
        "actions_per_month": 50000,
        "retention_days": 30,
        "max_agents": 10,
    },
    "enterprise": {
        "actions_per_month": float('inf'),  # Unlimited
        "retention_days": 365,
        "max_agents": float('inf'),
    }
}

# In-memory usage tracking (replace with Redis/DB in production)
# Format: {client_id: {"month": "2026-02", "count": 123, "plan": "starter"}}
USAGE_CACHE = {}


@dataclass
class UsageStatus:
    allowed: bool
    current_count: int
    limit: int
    remaining: int
    plan: str
    message: str


def get_current_month() -> str:
    """Get current month as YYYY-MM string."""
    return datetime.now(timezone.utc).strftime("%Y-%m")


def get_usage(client_id: str) -> dict:
    """Get current usage for a client."""
    if client_id not in USAGE_CACHE:
        USAGE_CACHE[client_id] = {
            "month": get_current_month(),
            "count": 0,
            "plan": "starter"
        }

    # Reset if new month
    current_month = get_current_month()
    if USAGE_CACHE[client_id]["month"] != current_month:
        USAGE_CACHE[client_id]["month"] = current_month
        USAGE_CACHE[client_id]["count"] = 0

    return USAGE_CACHE[client_id]


def set_plan(client_id: str, plan: str):
    """Set the plan for a client."""
    usage = get_usage(client_id)
    usage["plan"] = plan


def check_and_increment(client_id: str, plan: str = None) -> UsageStatus:
    """
    Check if client can perform action and increment counter.
    Returns UsageStatus with allowed=True/False.
    """
    usage = get_usage(client_id)

    # Update plan if provided
    if plan:
        usage["plan"] = plan

    current_plan = usage.get("plan", "starter")
    limits = PLAN_LIMITS.get(current_plan, PLAN_LIMITS["starter"])
    limit = limits["actions_per_month"]

    current_count = usage["count"]

    # Check if over limit
    if current_count >= limit:
        return UsageStatus(
            allowed=False,
            current_count=current_count,
            limit=int(limit) if limit != float('inf') else -1,
            remaining=0,
            plan=current_plan,
            message=f"Monthly limit reached ({int(limit)} actions). Upgrade to Pro for more."
        )

    # Increment and allow
    usage["count"] = current_count + 1
    remaining = int(limit - usage["count"]) if limit != float('inf') else -1

    return UsageStatus(
        allowed=True,
        current_count=usage["count"],
        limit=int(limit) if limit != float('inf') else -1,
        remaining=remaining,
        plan=current_plan,
        message="OK"
    )


def get_usage_stats(client_id: str) -> dict:
    """Get usage statistics for a client."""
    usage = get_usage(client_id)
    plan = usage.get("plan", "starter")
    limits = PLAN_LIMITS.get(plan, PLAN_LIMITS["starter"])
    limit = limits["actions_per_month"]

    return {
        "client_id": client_id,
        "plan": plan,
        "month": usage["month"],
        "actions_used": usage["count"],
        "actions_limit": int(limit) if limit != float('inf') else "unlimited",
        "actions_remaining": int(limit - usage["count"]) if limit != float('inf') else "unlimited",
        "retention_days": limits["retention_days"],
        "max_agents": int(limits["max_agents"]) if limits["max_agents"] != float('inf') else "unlimited",
    }


# Integration with payments.py - sync customer plans
def sync_from_payments():
    """Sync plans from payments module."""
    try:
        from .payments import CUSTOMERS
        for email, data in CUSTOMERS.items():
            client_id = data.get("client_id")
            plan = data.get("plan", "starter")
            if client_id:
                set_plan(client_id, plan)
    except ImportError:
        pass
