"""
Phase 92: A/B Testing for Platform Features.

Create experiments with 2+ variants, deterministically assign entities
via hash-based bucketing, record conversion events, and compute per-variant
statistics including conversion rates and average values.

Storage:
  state/broker/experiments/experiments.json   -- experiment definitions
  state/broker/experiments/events.jsonl       -- raw conversion events
"""

import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from brynq.server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync

# ── Paths ──────────────────────────────────────────────────────────────────

EXP_DIR = BROKER_DIR / "experiments"
EXP_DIR.mkdir(parents=True, exist_ok=True)

EXPERIMENTS_FILE = EXP_DIR / "experiments.json"
EVENTS_FILE = EXP_DIR / "events.jsonl"

# ── Persistence Helpers ────────────────────────────────────────────────────


def _load_experiments() -> dict:
    try:
        if EXPERIMENTS_FILE.exists():
            return json.loads(EXPERIMENTS_FILE.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        pass
    return {}


def _save_experiments(data: dict) -> None:
    _write_json_sync(EXPERIMENTS_FILE, data)


# ── Public API ─────────────────────────────────────────────────────────────


def create_experiment(name: str, description: str, variants: list[str],
                      traffic_split: Optional[list[float]] = None) -> dict:
    """Create an experiment with 2+ *variants*.

    *traffic_split* is a list of floats summing to 1.0 (one per variant).
    If omitted, traffic is split evenly.
    """
    if len(variants) < 2:
        return {"error": "An experiment needs at least 2 variants"}
    if len(set(variants)) != len(variants):
        return {"error": "Variant names must be unique"}

    if traffic_split is None:
        traffic_split = [round(1.0 / len(variants), 4)] * len(variants)
        # Fix rounding so it sums to 1.0
        traffic_split[-1] = round(1.0 - sum(traffic_split[:-1]), 4)
    else:
        if len(traffic_split) != len(variants):
            return {"error": "traffic_split length must match variants length"}
        if abs(sum(traffic_split) - 1.0) > 0.01:
            return {"error": "traffic_split must sum to 1.0"}

    experiments = _load_experiments()
    if name in experiments:
        return {"error": f"Experiment '{name}' already exists"}

    experiments[name] = {
        "name": name,
        "description": description,
        "variants": variants,
        "traffic_split": traffic_split,
        "status": "active",
        "winner": None,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "concluded_at": None,
    }
    _save_experiments(experiments)
    return {"ok": True, "name": name, "variants": variants}


def assign_variant(experiment_name: str, entity_id: str) -> dict:
    """Deterministically assign *entity_id* to a variant using hash bucketing."""
    experiments = _load_experiments()
    exp = experiments.get(experiment_name)
    if not exp:
        return {"error": f"Experiment '{experiment_name}' not found"}

    # Hash-based deterministic assignment
    key = f"{experiment_name}:{entity_id}"
    h = int(hashlib.sha256(key.encode()).hexdigest(), 16)
    bucket = (h % 10000) / 10000.0  # 0.0 - 0.9999

    cumulative = 0.0
    chosen = exp["variants"][-1]
    for variant, split in zip(exp["variants"], exp["traffic_split"]):
        cumulative += split
        if bucket < cumulative:
            chosen = variant
            break

    return {"experiment": experiment_name, "entity_id": entity_id, "variant": chosen}


def record_event(experiment_name: str, entity_id: str, event_type: str,
                 value: Any = None) -> dict:
    """Record a conversion or metric event for an entity in an experiment."""
    experiments = _load_experiments()
    if experiment_name not in experiments:
        return {"error": f"Experiment '{experiment_name}' not found"}

    assignment = assign_variant(experiment_name, entity_id)
    variant = assignment["variant"]

    entry = {
        "experiment": experiment_name,
        "entity_id": entity_id,
        "variant": variant,
        "event_type": event_type,
        "value": value,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    _append_jsonl_sync(EVENTS_FILE, entry)
    return {"ok": True, **entry}


def get_results(experiment_name: str) -> dict:
    """Compute per-variant results: count, conversion rate, average value."""
    experiments = _load_experiments()
    exp = experiments.get(experiment_name)
    if not exp:
        return {"error": f"Experiment '{experiment_name}' not found"}

    events = _read_jsonl_sync(EVENTS_FILE)
    exp_events = [e for e in events if e.get("experiment") == experiment_name]

    # Unique entities per variant
    variant_entities: dict[str, set] = {v: set() for v in exp["variants"]}
    variant_values: dict[str, list] = {v: [] for v in exp["variants"]}
    variant_conversions: dict[str, set] = {v: set() for v in exp["variants"]}

    for e in exp_events:
        v = e.get("variant")
        if v not in variant_entities:
            continue
        variant_entities[v].add(e["entity_id"])
        if e.get("event_type") == "conversion":
            variant_conversions[v].add(e["entity_id"])
        if e.get("value") is not None:
            try:
                variant_values[v].append(float(e["value"]))
            except (ValueError, TypeError):
                pass

    results = {}
    for v in exp["variants"]:
        entity_count = len(variant_entities[v])
        conversion_count = len(variant_conversions[v])
        values = variant_values[v]
        results[v] = {
            "entities": entity_count,
            "conversions": conversion_count,
            "conversion_rate": round(conversion_count / entity_count, 4) if entity_count else 0.0,
            "avg_value": round(sum(values) / len(values), 4) if values else 0.0,
            "total_events": sum(1 for e in exp_events if e.get("variant") == v),
        }

    return {"experiment": experiment_name, "status": exp["status"],
            "winner": exp.get("winner"), "results": results}


def conclude_experiment(experiment_name: str, winner: Optional[str] = None) -> dict:
    """End an experiment and optionally declare a *winner* variant."""
    experiments = _load_experiments()
    if experiment_name not in experiments:
        return {"error": f"Experiment '{experiment_name}' not found"}
    exp = experiments[experiment_name]
    if winner and winner not in exp["variants"]:
        return {"error": f"Winner '{winner}' is not a valid variant"}

    exp["status"] = "concluded"
    exp["winner"] = winner
    exp["concluded_at"] = datetime.now(timezone.utc).isoformat()
    _save_experiments(experiments)
    return {"ok": True, "experiment": experiment_name, "winner": winner}


def list_experiments(status: Optional[str] = None) -> list[dict]:
    """List all experiments, optionally filtered by status ('active'/'concluded')."""
    experiments = _load_experiments()
    result = list(experiments.values())
    if status:
        result = [e for e in result if e.get("status") == status]
    result.sort(key=lambda e: e.get("created_at", ""))
    return result


def get_experiment(name: str) -> Optional[dict]:
    """Get a single experiment by name."""
    return _load_experiments().get(name)
