"""
A/B Testing
Runs experiments comparing different strategies or models.
"""
import json
import random
from pathlib import Path
from typing import Dict, Any, List

class ABTest:
    def __init__(self, data_dir: Path, test_name: str, variant_a: Dict[str, Any], variant_b: Dict[str, Any]):
        self.data_dir = data_dir / "learning"
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.test_name = test_name
        self.variant_a = variant_a
        self.variant_b = variant_b
        self.log_file = self.data_dir / "ab_tests.jsonl"

    def assign(self) -> str:
        variant = 'A' if random.random() < 0.5 else 'B'
        with open(self.log_file, "a", encoding="utf-8") as f:
            f.write(json.dumps({"test_name": self.test_name, "event": "assign", "variant": variant}) + "\n")
        return variant

    def record_outcome(self, variant: str, latency_ms: int, success: bool, tokens: int):
        with open(self.log_file, "a", encoding="utf-8") as f:
            f.write(json.dumps({
                "test_name": self.test_name,
                "event": "outcome",
                "variant": variant,
                "latency_ms": latency_ms,
                "success": success,
                "tokens": tokens
            }) + "\n")

    def results(self) -> Dict[str, Any]:
        if not self.log_file.exists():
            return {"test_name": self.test_name, "variant_a": {}, "variant_b": {}, "winner": None, "confidence": 0.0}
            
        data = {'A': {'count': 0, 'latency_sum': 0, 'success_sum': 0, 'tokens_sum': 0},
                'B': {'count': 0, 'latency_sum': 0, 'success_sum': 0, 'tokens_sum': 0}}
                
        with open(self.log_file, "r", encoding="utf-8") as f:
            for line in f:
                try:
                    entry = json.loads(line)
                    if entry.get("test_name") == self.test_name and entry.get("event") == "outcome":
                        v = entry["variant"]
                        data[v]['count'] += 1
                        data[v]['latency_sum'] += entry.get("latency_ms", 0)
                        data[v]['success_sum'] += 1 if entry.get("success", False) else 0
                        data[v]['tokens_sum'] += entry.get("tokens", 0)
                except Exception:
                    pass
                    
        res = {"test_name": self.test_name}
        for v in ('A', 'B'):
            count = data[v]['count']
            if count > 0:
                res[f"variant_{v.lower()}"] = {
                    "count": count,
                    "avg_latency": data[v]['latency_sum'] / count,
                    "success_rate": data[v]['success_sum'] / count,
                    "avg_tokens": data[v]['tokens_sum'] / count
                }
            else:
                res[f"variant_{v.lower()}"] = {"count": 0, "avg_latency": 0, "success_rate": 0, "avg_tokens": 0}
                
        sa = res["variant_a"]["success_rate"]
        sb = res["variant_b"]["success_rate"]
        winner = 'A' if sa > sb else ('B' if sb > sa else None)
        
        res["winner"] = winner
        res["confidence"] = min(1.0, (data['A']['count'] + data['B']['count']) / 100.0) # naive confidence
        return res

class ABTestManager:
    def __init__(self, data_dir: Path):
        self.data_dir = data_dir / "learning"
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.configs_file = self.data_dir / "ab_configs.json"
        if not self.configs_file.exists():
            self.configs_file.write_text("{}", encoding="utf-8")

    def _read(self) -> Dict[str, Any]:
        try:
            return json.loads(self.configs_file.read_text(encoding="utf-8"))
        except Exception:
            return {}

    def _write(self, data: Dict[str, Any]):
        self.configs_file.write_text(json.dumps(data, indent=2), encoding="utf-8")

    def create_test(self, name: str, variant_a: Dict[str, Any], variant_b: Dict[str, Any]) -> ABTest:
        configs = self._read()
        configs[name] = {"variant_a": variant_a, "variant_b": variant_b, "status": "running"}
        self._write(configs)
        return ABTest(self.data_dir.parent, name, variant_a, variant_b)

    def list_tests(self) -> List[Dict[str, Any]]:
        configs = self._read()
        return [{"name": k, "status": v["status"], "sample_count": 0} for k, v in configs.items()]

    def conclude_test(self, name: str) -> Dict[str, Any]:
        configs = self._read()
        if name not in configs:
            return {}
        test = ABTest(self.data_dir.parent, name, configs[name]["variant_a"], configs[name]["variant_b"])
        res = test.results()
        configs[name]["status"] = "concluded"
        configs[name]["winner"] = res["winner"]
        self._write(configs)
        return configs[name]["variant_a"] if res["winner"] == 'A' else configs[name]["variant_b"]

    def auto_promote(self, name: str):
        winner_cfg = self.conclude_test(name)
        if winner_cfg:
            from .preferences import Preferences
            prefs = Preferences(self.data_dir.parent)
            for k, v in winner_cfg.items():
                prefs.save(k, v)

    def delete_test(self, name: str) -> bool:
        configs = self._read()
        if name in configs:
            del configs[name]
            self._write(configs)
            return True
        return False
