import base64
import copy
import json
import os
import subprocess
import tempfile
from datetime import datetime, timezone
from pathlib import Path

TIER_ORDER = ("starter", "basic", "pro", "max")
TIER_ALIASES = {
    "entry": "basic",
    "entry tier": "basic",
    "mid": "pro",
    "mid tier": "pro",
    "advanced": "max",
    "advanced tier": "max",
    "starter tier": "starter",
    "basic tier": "basic",
    "pro tier": "pro",
    "max tier": "max",
    # Server S/M/L/Max commercial labels — synonyms for the same four Server
    # tiers/ranks above (starter/basic/pro/max). Additive only; TIER_ORDER and
    # rank comparisons are unchanged.
    "server s": "starter",
    "server s tier": "starter",
    "server m": "basic",
    "server m tier": "basic",
    "server l": "pro",
    "server l tier": "pro",
    "server max": "max",
    "server max tier": "max",
}
STAGING_MAX_TIER = "basic"
SUPPORTED_SIGNATURE_ALGS = {"rsa-sha256", "sha256-rsa"}

# ── Product category: Workstation (standalone, NOT part of the Server ladder) ──
# Nexus One AI Workstation is a personal AI appliance (1-3 active users, 5 max),
# distinct from the Server S/M/L/Max ladder above (TIER_ORDER). It is additive:
# it does not appear in TIER_ORDER and must never be passed through tier_rank(),
# tier_lte(), min_tier(), evaluate_license()/build_tier_options() ladder logic —
# those functions and STAGING_MAX_TIER/TIER_ALIASES above are unchanged and only
# ever resolve to one of TIER_ORDER's four Server tiers. Workstation entitlement
# is looked up independently via TIER_MATRIX["workstation"] in main.py.
WORKSTATION_TIER = "workstation"
CATEGORY_SERVER = "server"
CATEGORY_WORKSTATION = "workstation"
WORKSTATION_ALIASES = {"workstation tier", "personal", "personal ai", "nexus one ai workstation"}


def is_workstation_tier(value):
    """True if `value` refers to the Workstation category (never a Server tier)."""
    return (value or "").strip().lower() in ({WORKSTATION_TIER} | WORKSTATION_ALIASES)


def tier_category(value):
    """Return CATEGORY_WORKSTATION or CATEGORY_SERVER for a tier/category string.
    Server tiers (starter/basic/pro/max — displayed as Server S/M/L/Max) rank via
    TIER_ORDER; Workstation has no rank and is resolved independently."""
    return CATEGORY_WORKSTATION if is_workstation_tier(value) else CATEGORY_SERVER


def utcnow_iso():
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def parse_time(value):
    text = (value or "").strip()
    if not text:
        return None
    if len(text) == 10:
        text += "T00:00:00Z"
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        return datetime.fromisoformat(text)
    except Exception:
        return None


def normalize_tier(value, default="basic"):
    raw = (value or "").strip().lower()
    if raw in TIER_ORDER:
        return raw
    return TIER_ALIASES.get(raw, default)


def tier_rank(tier):
    return TIER_ORDER.index(normalize_tier(tier))


def tier_lte(left, right):
    return tier_rank(left) <= tier_rank(right)


def min_tier(*tiers):
    items = [normalize_tier(t) for t in tiers if t]
    if not items:
        return STAGING_MAX_TIER
    return sorted(items, key=tier_rank)[0]


def read_json_file(path):
    try:
        data = json.loads(Path(path).read_text())
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}


def write_json_file(path, payload):
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(payload, indent=2))


def _canonical_payload(payload):
    clean = copy.deepcopy(payload or {})
    for key in ("signature", "signature_status", "verification_status"):
        clean.pop(key, None)
    return json.dumps(clean, sort_keys=True, separators=(",", ":"), ensure_ascii=True)


def verify_signature(payload, public_key_path):
    if not payload:
        return False, "missing"
    signature = (payload.get("signature") or "").strip()
    if not signature:
        return False, "missing_signature"
    alg = (payload.get("signature_alg") or "").strip().lower()
    if alg not in SUPPORTED_SIGNATURE_ALGS:
        return False, "unsupported_signature_alg"
    key_path = Path(public_key_path)
    if not key_path.exists():
        return False, "missing_public_key"
    canonical = _canonical_payload(payload).encode()
    try:
        sig_bytes = base64.b64decode(signature, validate=True)
    except Exception:
        return False, "bad_signature_encoding"

    with tempfile.NamedTemporaryFile(delete=False) as data_file, tempfile.NamedTemporaryFile(delete=False) as sig_file:
        data_file.write(canonical)
        sig_file.write(sig_bytes)
        data_file.flush()
        sig_file.flush()
        data_path = data_file.name
        sig_path = sig_file.name
    try:
        result = subprocess.run(
            ["openssl", "dgst", "-sha256", "-verify", str(key_path), "-signature", sig_path, data_path],
            capture_output=True,
            text=True,
            check=False,
        )
        return result.returncode == 0, "valid" if result.returncode == 0 else "invalid_signature"
    finally:
        for tmp in (data_path, sig_path):
            try:
                os.unlink(tmp)
            except OSError:
                pass


def public_license_record(record):
    public = {k: v for k, v in (record or {}).items() if k != "signature"}
    if record.get("license_key"):
        public["license_key_prefix"] = str(record.get("license_key"))[:12]
        public["license_key_present"] = True
    else:
        public["license_key_present"] = False
    return public


def evaluate_license(payload, public_key_path, machine_id=None, now=None):
    now_dt = parse_time(now) if isinstance(now, str) else now or datetime.now(timezone.utc)
    base = {
        "schema": "cezen.license.v2",
        "status": "missing",
        "valid": False,
        "allowed_tier": STAGING_MAX_TIER,
        "feature_overrides": {},
        "install_type": "field-staging",
        "needs_activation": True,
        "license_record": {},
        "notes": [],
    }
    if not payload:
        return base

    record = payload if isinstance(payload, dict) else {}
    base["license_record"] = public_license_record(record)
    base["install_type"] = (record.get("install_type") or "licensed").strip() or "licensed"
    base["allowed_tier"] = normalize_tier(record.get("allowed_tier") or record.get("tier"), STAGING_MAX_TIER)
    overrides = record.get("feature_overrides") or {}
    base["feature_overrides"] = overrides if isinstance(overrides, dict) else {}

    ok, sig_status = verify_signature(record, public_key_path)
    if not ok:
        base["status"] = "invalid_signature" if sig_status != "missing_signature" else "missing"
        base["notes"].append(sig_status)
        return base

    issued_at = parse_time(record.get("issued_at"))
    if issued_at and issued_at > now_dt:
        base["status"] = "not_yet_valid"
        base["notes"].append("issued_at_in_future")
        return base

    expires_at = parse_time(record.get("expires_at"))
    if expires_at and expires_at < now_dt:
        base["status"] = "expired"
        base["notes"].append("license_expired")
        return base

    machine_binding = record.get("machine_binding") or {}
    if machine_id and isinstance(machine_binding, dict):
        bound_id = (machine_binding.get("machine_id") or "").strip()
        if bound_id and bound_id != machine_id:
            base["status"] = "machine_mismatch"
            base["notes"].append("machine_binding_mismatch")
            return base

    base["status"] = "valid"
    base["valid"] = True
    base["needs_activation"] = False
    return base


def evaluate_override(payload, public_key_path, now=None):
    if not payload:
        return {"status": "missing", "valid": False, "allow_hardware_mismatch": False, "max_override_tier": None}
    record = payload if isinstance(payload, dict) else {}
    ok, sig_status = verify_signature(record, public_key_path)
    if not ok:
        return {
            "status": "invalid_signature" if sig_status != "missing_signature" else "missing",
            "valid": False,
            "allow_hardware_mismatch": False,
            "max_override_tier": None,
        }
    expires_at = parse_time(record.get("expires_at"))
    now_dt = parse_time(now) if isinstance(now, str) else now or datetime.now(timezone.utc)
    if expires_at and expires_at < now_dt:
        return {"status": "expired", "valid": False, "allow_hardware_mismatch": False, "max_override_tier": None}
    return {
        "status": "valid",
        "valid": True,
        "allow_hardware_mismatch": bool(record.get("allow_hardware_mismatch")),
        "max_override_tier": normalize_tier(record.get("max_override_tier"), "max"),
        "override_record": public_license_record(record),
    }


def build_tier_options(license_eval, hardware_tier, override_eval=None):
    override_eval = override_eval or {}
    licensed_cap = normalize_tier(license_eval.get("allowed_tier"), STAGING_MAX_TIER)
    hardware_cap = normalize_tier(hardware_tier, STAGING_MAX_TIER)
    override_cap = normalize_tier(override_eval.get("max_override_tier"), "max") if override_eval.get("allow_hardware_mismatch") else None
    options = []
    for tier in TIER_ORDER:
        state = "enabled"
        reason = ""
        if not tier_lte(tier, licensed_cap):
            state = "disabled_by_license"
            reason = f"Not included in current license ({licensed_cap})."
        elif not tier_lte(tier, hardware_cap):
            if override_cap and tier_lte(tier, override_cap):
                state = "override_required"
                reason = f"Requires Cezen override; hardware recommendation is {hardware_cap}."
            else:
                state = "disabled_by_hardware"
                reason = f"Requires larger hardware; recommendation is {hardware_cap}."
        options.append({"tier": tier, "state": state, "reason": reason, "selectable": state in {"enabled", "override_required"}})
    return options


def provisioned_components_from_record(record):
    record = record or {}
    components = record.get("components") or {}
    if isinstance(components, dict):
        return components
    skip_roles = set(record.get("skip_roles") or [])
    return {
        "ollama": "ollama" not in skip_roles,
        "jupyterlab": "jupyterlab" not in skip_roles,
        "chromadb": "chromadb" not in skip_roles,
        "vllm": "vllm" not in skip_roles,
        "mlflow": "mlflow" not in skip_roles,
        "minio": "minio" not in skip_roles,
        "monitoring": "monitoring" not in skip_roles,
        "k3s": "k3s" not in skip_roles,
    }


def resolve_effective_features(base_features, feature_overrides, provisioned_components, hardware_features):
    features = copy.deepcopy(base_features or {})
    for key, value in (feature_overrides or {}).items():
        features[key] = value

    provisioned_components = provisioned_components or {}
    hardware_features = hardware_features or {}

    if features.get("rag") and not provisioned_components.get("chromadb", True):
        features["rag"] = False
    if features.get("gpu_inference") and not (
        hardware_features.get("ollama_gpu") or hardware_features.get("vllm")
    ):
        features["gpu_inference"] = False
    if features.get("fine_tuning") and not (
        provisioned_components.get("jupyterlab", False) and hardware_features.get("fine_tuning_qlora")
    ):
        features["fine_tuning"] = False
    if features.get("deepspeed") and not hardware_features.get("distributed_training"):
        features["deepspeed"] = False
    return features


def collect_entitlement(tier_matrix, license_record, install_record, feasibility, public_key_path, machine_id=None):
    license_eval = evaluate_license(license_record, public_key_path, machine_id=machine_id)
    override_eval = evaluate_override((install_record or {}).get("override_record") or read_json_file((install_record or {}).get("override_path", "")), public_key_path)
    licensed_tier = normalize_tier(license_eval.get("allowed_tier"), STAGING_MAX_TIER)
    provisioned_tier = normalize_tier((install_record or {}).get("provisioned_tier") or (install_record or {}).get("selected_tier"), STAGING_MAX_TIER)
    hardware_tier = normalize_tier(((feasibility or {}).get("recommendation") or {}).get("recommended_tier"), STAGING_MAX_TIER)
    components = provisioned_components_from_record(install_record)
    tier_defaults = copy.deepcopy((tier_matrix or {}).get(licensed_tier, {}))
    effective_features = resolve_effective_features(
        tier_defaults.get("features", {}),
        license_eval.get("feature_overrides", {}),
        components,
        (feasibility or {}).get("features", {}),
    )

    states = []
    if tier_rank(licensed_tier) > tier_rank(provisioned_tier):
        states.append("licensed_not_provisioned")
    if tier_rank(provisioned_tier) > tier_rank(licensed_tier):
        states.append("provisioned_above_license")
    if tier_rank(provisioned_tier) > tier_rank(hardware_tier) and not override_eval.get("allow_hardware_mismatch"):
        states.append("hardware_below_provisioned")
    if license_eval.get("status") != "valid":
        states.append(license_eval.get("status"))

    return {
        "license_status": license_eval.get("status"),
        "license_valid": license_eval.get("valid"),
        "license_record": public_license_record(license_record),
        "licensed_tier": licensed_tier,
        "provisioned_tier": provisioned_tier,
        "hardware_recommended_tier": hardware_tier,
        "effective_features": effective_features,
        "feature_overrides": license_eval.get("feature_overrides", {}),
        "provisioned_components": components,
        "compliance_state": states or ["licensed_and_provisioned"],
        "override": override_eval,
        "install_type": license_eval.get("install_type") or (install_record or {}).get("install_type") or "field-staging",
    }