import hashlib
import hmac
import html
import json
import logging
import os
import secrets
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional
from uuid import uuid4

from quota_fallback import (
    GRANT,
    REFUND,
    embed_parent_from,
    embed_parents,
    grant_request_nonce,
    is_eligible_embed,
    issue_ticket,
    request_header,
    request_host,
    safe_message,
    signing_secret,
    verify_grant,
)

logger = logging.getLogger(__name__)

FREE_LIMIT_REASON = "LABS_3DAI_FREE_LIMIT"
GRANT_TTL_SECONDS = 180
SETTLED_GRANT_TTL_SECONDS = 3600
MAX_GRANT_ENTRIES = 512
MAX_GRANTS_PER_SESSION = 8
SIGNUP_URL = "https://app.skyebrowse.com/register"
EMBED_LIMIT_MESSAGE = (
    "You have used your free 3D AI generations. Upgrade to a Commercial plan to keep generating."
)
EMBED_UNAVAILABLE_MESSAGE = "3D AI could not confirm your SkyeBrowse plan. Please try again."
PUBLIC_LIMIT_MESSAGE = "You have used your {limit} free generations."
PUBLIC_QUOTA_MESSAGE = "Free generations are busy right now."


def env_flag(name: str) -> bool:
    return os.environ.get(name, "").strip().lower() in {"1", "true", "yes"}


def env_int(name: str, default: int) -> int:
    try:
        return max(0, int(os.environ.get(name, "").strip() or default))
    except ValueError:
        return default


def grant_parents() -> Optional[frozenset]:
    raw = os.environ.get("LABS_EMBED_GRANTS", "").strip()
    if env_flag("LABS_EMBED_GRANTS"):
        return None
    return frozenset(item.strip() for item in raw.split(",") if item.strip())


def grants_required(request: Any) -> bool:
    if not is_eligible_embed(request):
        return False
    parents = grant_parents()
    if parents is None:
        return True
    return embed_parent_from(request_header(request, "referer"), request_host(request)) in parents


def is_browser_visit(request: Any) -> bool:
    referer = request_header(request, "referer")
    return referer.startswith(("https://", "http://")) and not is_eligible_embed(request)


def visitor_ip(request: Any) -> str:
    return request_header(request, "x-forwarded-for").split(",")[0].strip()


@dataclass
class GrantEntry:
    nonce: str
    session_hash: str
    issued_at: float
    settled_at: Optional[float] = None
    mode: Optional[str] = None
    notice: Optional[str] = None


@dataclass
class Access:
    mode: str = "native"
    nonce: Optional[str] = None
    visitor_key: Optional[str] = None
    notice: Optional[str] = None


class GrantBook:
    def __init__(
        self,
        ttl: float = GRANT_TTL_SECONDS,
        max_entries: int = MAX_GRANT_ENTRIES,
        settled_ttl: float = SETTLED_GRANT_TTL_SECONDS,
        max_per_session: int = MAX_GRANTS_PER_SESSION,
    ):
        self._lock = threading.Lock()
        self._entries: Dict[str, GrantEntry] = {}
        self._lapsed: Dict[str, List[str]] = {}
        self._ttl = ttl
        self._settled_ttl = settled_ttl
        self._max_entries = max_entries
        self._max_per_session = max_per_session

    def request(self, session_hash: str) -> Optional[Dict[str, str]]:
        if not session_hash:
            return None
        with self._lock:
            self._purge()
            if len(self._entries) >= self._max_entries:
                return None
            if len(self._session_entries(session_hash)) >= self._max_per_session:
                return None
            entry = GrantEntry(uuid4().hex, session_hash, time.time())
            self._entries[entry.nonce] = entry
        ticket = issue_ticket(entry.nonce, GRANT)
        return {"ticket": ticket} if ticket else None

    def settle(self, session_hash: str, ack: Any, receipt: Any) -> None:
        if not isinstance(ack, dict) or not isinstance(receipt, dict) or not receipt.get("ticket"):
            return
        if ack.get("ticket") != receipt.get("ticket"):
            return
        nonce = grant_request_nonce(receipt.get("ticket"))
        with self._lock:
            self._purge()
            entry = self._entries.get(nonce) if nonce else None
            if entry is None or entry.session_hash != session_hash or entry.settled_at is not None:
                return
            entry.settled_at = time.time()
            grant = verify_grant(ack.get("grant"), entry.nonce)
            if grant is not None:
                entry.mode = grant["mode"]
            elif ack.get("reason") == FREE_LIMIT_REASON:
                entry.notice = safe_message(ack.get("error"), EMBED_LIMIT_MESSAGE)
            else:
                entry.notice = safe_message(ack.get("error"), EMBED_UNAVAILABLE_MESSAGE)

    def take(self, session_hash: str) -> Access:
        with self._lock:
            self._purge()
            entry = self._pop_settled(session_hash)
            lapsed = None if entry else self._pop_lapsed(session_hash)
        if lapsed:
            return Access(mode="denied", nonce=lapsed, notice=embed_notice(EMBED_UNAVAILABLE_MESSAGE))
        if entry is None or entry.mode is None:
            notice = entry.notice if entry and entry.notice else EMBED_UNAVAILABLE_MESSAGE
            return Access(mode="denied", notice=embed_notice(notice))
        return Access(nonce=entry.nonce)

    def _pop_settled(self, session_hash: str) -> Optional[GrantEntry]:
        settled = [entry for entry in self._session_entries(session_hash) if entry.settled_at is not None]
        return self._entries.pop(min(settled, key=lambda entry: entry.settled_at).nonce) if settled else None

    def _pop_lapsed(self, session_hash: str) -> Optional[str]:
        lapsed = self._lapsed.get(session_hash)
        if not lapsed:
            return None
        nonce = lapsed.pop(0)
        if not lapsed:
            del self._lapsed[session_hash]
        return nonce

    def _session_entries(self, session_hash: str) -> List[GrantEntry]:
        return [entry for entry in self._entries.values() if entry.session_hash == session_hash]

    def _purge(self) -> None:
        now = time.time()
        for nonce, entry in list(self._entries.items()):
            pending = entry.settled_at is None
            started = entry.issued_at if pending else entry.settled_at
            if started < now - (self._ttl if pending else self._settled_ttl):
                del self._entries[nonce]
                if entry.mode is not None:
                    self._lapse(entry)

    def _lapse(self, entry: GrantEntry) -> None:
        self._lapsed.setdefault(entry.session_hash, []).append(entry.nonce)
        while len(self._lapsed) > self._max_entries:
            del self._lapsed[next(iter(self._lapsed))]


class PublicAllowance:
    def __init__(self, limit: int, window_seconds: float, path: Optional[Path]):
        self.limit = limit
        self._window = window_seconds
        self._path = path
        self._lock = threading.Lock()
        self._salt = (signing_secret() or secrets.token_hex(32)).encode("utf-8")
        self._uses: Dict[str, List[float]] = self._load()

    @property
    def enabled(self) -> bool:
        return self.limit > 0

    def key_for(self, request: Any) -> Optional[str]:
        address = visitor_ip(request)
        if not address:
            logger.warning("Public allowance skipped: request has no X-Forwarded-For")
            return None
        return hmac.new(self._salt, address.encode("utf-8"), hashlib.sha256).hexdigest()[:32]

    def try_consume(self, key: Optional[str]) -> bool:
        if key is None:
            return True
        with self._lock:
            uses = self._recent(key)
            if len(uses) >= self.limit:
                return False
            self._uses[key] = uses + [time.time()]
            self._save()
            return True

    def refund(self, key: Optional[str]) -> None:
        if key is None:
            return
        with self._lock:
            uses = self._recent(key)
            if uses:
                self._uses[key] = uses[:-1]
                self._save()

    def used(self, key: Optional[str]) -> int:
        with self._lock:
            return len(self._recent(key)) if key else 0

    def _recent(self, key: str) -> List[float]:
        cutoff = time.time() - self._window
        return [stamp for stamp in self._uses.get(key, []) if stamp > cutoff]

    def _load(self) -> Dict[str, List[float]]:
        if self._path is None or not self._path.exists():
            return {}
        try:
            data = json.loads(self._path.read_text())
        except (OSError, ValueError):
            return {}
        if not isinstance(data, dict):
            return {}
        return {str(key): [float(stamp) for stamp in value if isinstance(stamp, (int, float))]
                for key, value in data.items() if isinstance(value, list)}

    def _save(self) -> None:
        if self._path is None:
            return
        cutoff = time.time() - self._window
        snapshot = {key: [stamp for stamp in uses if stamp > cutoff] for key, uses in self._uses.items()}
        snapshot = {key: uses for key, uses in snapshot.items() if uses}
        self._uses = snapshot
        try:
            temporary = self._path.with_suffix(".tmp")
            temporary.write_text(json.dumps(snapshot))
            temporary.replace(self._path)
        except OSError as error:
            logger.warning("Public allowance not persisted (%s)", type(error).__name__)


def public_usage_path() -> Optional[Path]:
    configured = os.environ.get("LABS_PUBLIC_USAGE_PATH", "").strip()
    if configured:
        return Path(configured)
    base = Path("/data") if Path("/data").is_dir() else Path("/tmp")
    return base / "labs_3dai_public_usage.json"


def public_allowance_from_env() -> PublicAllowance:
    hours = env_int("LABS_PUBLIC_WINDOW_HOURS", 720) or 720
    return PublicAllowance(env_int("LABS_PUBLIC_FREE_LIMIT", 0), hours * 3600, public_usage_path())


def embed_notice(message: str) -> str:
    return f'<div class="labs-access-notice">{html.escape(message)}</div>'


def signup_notice(message: str) -> str:
    return (
        '<div class="labs-access-notice">'
        f"<strong>{html.escape(message)}</strong> "
        "Create a free SkyeBrowse account to keep generating 3D models. "
        f'<a href="{SIGNUP_URL}" target="_blank" rel="noopener">Sign up free</a>'
        "</div>"
    )


def bridge_script(path: Path) -> str:
    return path.read_text().replace("__LABS_EMBED_PARENTS__", json.dumps(list(embed_parents())))


class AccessPolicy:
    def __init__(self, grants: GrantBook, public: PublicAllowance):
        self.grants = grants
        self.public = public
        self._notices: Dict[str, str] = {}
        self._refunds: Dict[str, List[str]] = {}
        self._lock = threading.Lock()

    def authorize(self, request: Any, session_hash: str) -> Access:
        if grants_required(request):
            access = self.grants.take(session_hash)
            if access.mode == "denied":
                self.generation_failed(session_hash, access)
            return access
        if not (self.public.enabled and is_browser_visit(request)):
            return Access()
        key = self.public.key_for(request)
        if not self.public.try_consume(key):
            limit_message = PUBLIC_LIMIT_MESSAGE.format(limit=self.public.limit)
            return Access(mode="denied", notice=signup_notice(limit_message))
        return Access(visitor_key=key)

    def quota_denied(self, access: Access) -> Optional[str]:
        self.public.refund(access.visitor_key)
        return signup_notice(PUBLIC_QUOTA_MESSAGE) if access.visitor_key else None

    def generation_failed(self, session_hash: str, access: Access) -> None:
        self.public.refund(access.visitor_key)
        if access.nonce:
            with self._lock:
                self._refunds.setdefault(session_hash, []).append(access.nonce)

    def refund_receipt(self, session_hash: str) -> Optional[Dict[str, str]]:
        with self._lock:
            pending = self._refunds.get(session_hash) or []
            nonce = pending.pop(0) if pending else None
            if not pending:
                self._refunds.pop(session_hash, None)
        ticket = issue_ticket(nonce, REFUND) if nonce else None
        return {"ticket": ticket} if ticket else None

    def remember(self, session_hash: str, notice: Optional[str]) -> None:
        with self._lock:
            if notice:
                self._notices[session_hash] = notice
            else:
                self._notices.pop(session_hash, None)

    def pop_notice(self, session_hash: str) -> Optional[str]:
        with self._lock:
            return self._notices.pop(session_hash, None)