"""
🔬 CANIFA Chatbot - STRESS PROFILER (Clear Cache + Multi-User)
===============================================================
Test bottleneck kỹ càng:
  1. Clear Redis cache (cold start thật)
  2. Test 1 user tuần tự (baseline)
  3. Test multi-user đồng thời (concurrent)
  4. So sánh Cold vs Warm performance

Dùng:
  python tests/profiler_stress.py                  # Full test (clear cache + 1 user + 3 users)
  python tests/profiler_stress.py --users 5        # 5 users đồng thời
  python tests/profiler_stress.py --skip-clear     # Không clear cache
  python tests/profiler_stress.py --warm-only      # Chỉ test warm (có cache)
"""

import argparse
import asyncio
import json
import logging
import os
import platform
import statistics
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

import httpx

try:
    import redis.asyncio as aioredis
except ImportError:
    aioredis = None
    print("⚠️  redis package không có, skip cache clear")

try:
    import psutil
except ImportError:
    psutil = None

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger("stress_profiler")

# =====================================================================
# CONFIG
# =====================================================================
API_URL = "http://localhost:5000"
CHAT_ENDPOINT = "/api/agent/chat-dev"

# Redis config (from config.py defaults)
REDIS_HOST = os.getenv("REDIS_CACHE_URL", "172.16.2.192")
REDIS_PORT = int(os.getenv("REDIS_CACHE_PORT", "6379"))
REDIS_DB = int(os.getenv("REDIS_CACHE_DB", "2"))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", None)
REDIS_USERNAME = os.getenv("REDIS_USERNAME", None)

TEST_QUERIES = [
    "Có áo polo nam màu xanh navy không?",
    "Tìm váy liền cho bé gái 5 tuổi",
    "Gợi ý outfit đi biển mùa hè",
    "Áo khoác gió nam size L giá dưới 500k",
    "Quần jeans nữ ống rộng màu đen",
]


# =====================================================================
# DATA
# =====================================================================
@dataclass
class QueryMetric:
    query: str
    user_id: str
    total_time: float = 0.0
    status_code: int = 0
    error: str = ""
    has_products: bool = False
    tokens: int = 0
    phase: str = ""  # "cold" or "warm"


@dataclass
class PhaseResult:
    phase: str
    metrics: list[QueryMetric] = field(default_factory=list)

    @property
    def success_times(self):
        return [m.total_time for m in self.metrics if m.status_code == 200]

    @property
    def avg(self):
        t = self.success_times
        return statistics.mean(t) if t else 0

    @property
    def p50(self):
        t = sorted(self.success_times)
        return t[len(t) // 2] if t else 0

    @property
    def p95(self):
        t = sorted(self.success_times)
        return t[int(len(t) * 0.95)] if t else 0

    @property
    def min_t(self):
        t = self.success_times
        return min(t) if t else 0

    @property
    def max_t(self):
        t = self.success_times
        return max(t) if t else 0

    @property
    def success_count(self):
        return sum(1 for m in self.metrics if m.status_code == 200)

    @property
    def error_count(self):
        return sum(1 for m in self.metrics if m.error)


# =====================================================================
# CACHE MANAGEMENT
# =====================================================================
async def clear_redis_cache() -> dict:
    """Clear toàn bộ Redis cache (resp_cache + emb_cache)."""
    if not aioredis:
        return {"status": "skipped", "reason": "redis package not installed"}

    try:
        conn_kwargs = {
            "host": REDIS_HOST,
            "port": REDIS_PORT,
            "db": REDIS_DB,
            "decode_responses": True,
            "socket_connect_timeout": 5,
        }
        if REDIS_PASSWORD:
            conn_kwargs["password"] = REDIS_PASSWORD
        if REDIS_USERNAME:
            conn_kwargs["username"] = REDIS_USERNAME

        client = aioredis.Redis(**conn_kwargs)
        await client.ping()

        # Count keys before
        resp_keys = await client.keys("resp_cache:*")
        emb_keys = await client.keys("emb_cache:*")

        counts = {
            "resp_cache_before": len(resp_keys),
            "emb_cache_before": len(emb_keys),
        }

        # Delete all cache keys (but keep other keys like prompt_version)
        deleted = 0
        if resp_keys:
            deleted += await client.delete(*resp_keys)
        if emb_keys:
            deleted += await client.delete(*emb_keys)

        counts["deleted"] = deleted

        # Verify
        resp_after = await client.keys("resp_cache:*")
        emb_after = await client.keys("emb_cache:*")
        counts["resp_cache_after"] = len(resp_after)
        counts["emb_cache_after"] = len(emb_after)

        await client.aclose()

        return {"status": "cleared", **counts}

    except Exception as e:
        return {"status": "error", "error": str(e)}


async def get_cache_stats() -> dict:
    """Lấy thông tin cache hiện tại."""
    if not aioredis:
        return {"status": "unavailable"}

    try:
        conn_kwargs = {
            "host": REDIS_HOST,
            "port": REDIS_PORT,
            "db": REDIS_DB,
            "decode_responses": True,
            "socket_connect_timeout": 5,
        }
        if REDIS_PASSWORD:
            conn_kwargs["password"] = REDIS_PASSWORD
        if REDIS_USERNAME:
            conn_kwargs["username"] = REDIS_USERNAME

        client = aioredis.Redis(**conn_kwargs)
        await client.ping()

        resp_keys = await client.keys("resp_cache:*")
        emb_keys = await client.keys("emb_cache:*")

        await client.aclose()

        return {
            "resp_cache_keys": len(resp_keys),
            "emb_cache_keys": len(emb_keys),
        }
    except Exception as e:
        return {"error": str(e)}


# =====================================================================
# QUERY RUNNER
# =====================================================================
async def run_query(client: httpx.AsyncClient, query: str, user_id: str, phase: str) -> QueryMetric:
    """Chạy 1 query, đo thời gian."""
    metric = QueryMetric(query=query, user_id=user_id, phase=phase)

    payload = {"user_query": query, "user_id": user_id}
    start = time.perf_counter()

    try:
        response = await client.post(
            f"{API_URL}{CHAT_ENDPOINT}",
            json=payload,
            timeout=180.0,
        )
        metric.total_time = time.perf_counter() - start
        metric.status_code = response.status_code

        if response.status_code == 200:
            body = response.text
            metric.has_products = "product_ids" in body.lower()
            metric.tokens = len(body.split())

    except httpx.TimeoutException:
        metric.total_time = time.perf_counter() - start
        metric.error = "TIMEOUT (>180s)"
    except Exception as e:
        metric.total_time = time.perf_counter() - start
        metric.error = str(e)

    return metric


# =====================================================================
# TEST PHASES
# =====================================================================
async def run_sequential(queries: list[str], user_id: str, phase: str) -> PhaseResult:
    """Chạy queries tuần tự - 1 user."""
    result = PhaseResult(phase=phase)

    async with httpx.AsyncClient() as client:
        for i, query in enumerate(queries, 1):
            logger.info(f"   [{phase.upper()}] {i}/{len(queries)}: {query[:40]}...")
            metric = await run_query(client, query, user_id, phase)
            result.metrics.append(metric)

            status = "✅" if not metric.error else "❌"
            logger.info(f"   {status} {metric.total_time:.2f}s | Products: {metric.has_products}")

            # Nghỉ 1s giữa queries
            if i < len(queries):
                await asyncio.sleep(1)

    return result


async def run_concurrent(queries: list[str], num_users: int, phase: str) -> PhaseResult:
    """Chạy queries đồng thời - nhiều users."""
    result = PhaseResult(phase=phase)

    async def user_worker(user_idx: int, query: str):
        user_id = f"stress_user_{user_idx}"
        async with httpx.AsyncClient() as client:
            metric = await run_query(client, query, user_id, phase)
            return metric

    # Mỗi user gửi 1 query cùng lúc
    tasks = []
    for i in range(num_users):
        query = queries[i % len(queries)]
        tasks.append(user_worker(i + 1, query))

    logger.info(f"   [{phase.upper()}] Launching {num_users} concurrent users...")
    start = time.perf_counter()
    metrics = await asyncio.gather(*tasks, return_exceptions=True)
    total = time.perf_counter() - start
    logger.info(f"   [{phase.upper()}] All {num_users} users finished in {total:.1f}s")

    for m in metrics:
        if isinstance(m, Exception):
            result.metrics.append(QueryMetric(
                query="error", user_id="unknown", error=str(m), phase=phase
            ))
        else:
            result.metrics.append(m)

    return result


# =====================================================================
# REPORT
# =====================================================================
def print_phase_table(phase: PhaseResult):
    """In bảng kết quả cho 1 phase."""
    print(f"\n   {'User':<15} {'Query':<35} {'Time':>8} {'Status':>8}")
    print(f"   {'─'*68}")
    for m in phase.metrics:
        q = m.query[:33] + ".." if len(m.query) > 33 else m.query
        status = "✅" if not m.error else f"❌ {m.error[:10]}"
        print(f"   {m.user_id:<15} {q:<35} {m.total_time:>7.2f}s {status:>8}")


def print_comparison(phases: list[PhaseResult]):
    """So sánh giữa các phases."""
    print(f"\n   {'Phase':<25} {'Avg':>8} {'P50':>8} {'P95':>8} {'Min':>8} {'Max':>8} {'OK':>5} {'Err':>5}")
    print(f"   {'─'*78}")
    for p in phases:
        if p.success_times:
            print(
                f"   {p.phase:<25} {p.avg:>7.2f}s {p.p50:>7.2f}s {p.p95:>7.2f}s "
                f"{p.min_t:>7.2f}s {p.max_t:>7.2f}s {p.success_count:>5} {p.error_count:>5}"
            )
        else:
            print(f"   {p.phase:<25} {'N/A':>8} {'N/A':>8} {'N/A':>8} {'N/A':>8} {'N/A':>8} {0:>5} {p.error_count:>5}")


def print_full_report(phases: list[PhaseResult], cache_before: dict, cache_after: dict):
    """In báo cáo tổng hợp."""
    print("\n" + "=" * 75)
    print("🔬 BÁO CÁO STRESS TEST - CANIFA CHATBOT")
    print("=" * 75)

    # System
    print(f"\n📦 System: {platform.platform()}")
    print(f"   Python: {platform.python_version()} | CPUs: {os.cpu_count()}")
    if psutil:
        vm = psutil.virtual_memory()
        print(f"   RAM: {vm.available / (1024**3):.1f}GB free / {vm.total / (1024**3):.1f}GB total")

    # Cache info
    print(f"\n📦 Redis Cache:")
    print(f"   Before: resp={cache_before.get('resp_cache_keys', '?')}, emb={cache_before.get('emb_cache_keys', '?')}")
    print(f"   After:  resp={cache_after.get('resp_cache_keys', '?')}, emb={cache_after.get('emb_cache_keys', '?')}")

    # Per-phase details
    for phase in phases:
        print(f"\n{'─'*75}")
        print(f"📊 Phase: {phase.phase}")
        print_phase_table(phase)

    # Comparison
    print(f"\n{'='*75}")
    print("📊 SO SÁNH CÁC PHASE")
    print(f"{'='*75}")
    print_comparison(phases)

    # Bottleneck analysis
    print(f"\n{'='*75}")
    print("🎯 PHÂN TÍCH BOTTLENECK")
    print(f"{'='*75}")

    cold = next((p for p in phases if "cold" in p.phase.lower()), None)
    warm = next((p for p in phases if "warm" in p.phase.lower()), None)
    concurrent = next((p for p in phases if "concurrent" in p.phase.lower()), None)

    if cold and warm and cold.success_times and warm.success_times:
        speedup = cold.avg / warm.avg if warm.avg > 0 else 0
        cache_impact = cold.avg - warm.avg
        print(f"\n   🧊 Cold (no cache) Avg: {cold.avg:.2f}s")
        print(f"   🔥 Warm (cached)   Avg: {warm.avg:.2f}s")
        print(f"   ⚡ Cache speedup:       {speedup:.1f}x ({cache_impact:.1f}s faster)")

        if cache_impact > 3:
            print(f"   🔴 Cache impact LỚN ({cache_impact:.1f}s): Embedding API là bottleneck chính!")
            print(f"      → Khuyến nghị: Cache warming, local embedding model")
        elif cache_impact > 1:
            print(f"   🟡 Cache impact VỪA ({cache_impact:.1f}s): Có cải thiện đáng kể")
        else:
            print(f"   🟢 Cache impact NHỎ ({cache_impact:.1f}s): Bottleneck ở chỗ khác (LLM hoặc DB)")

    if concurrent and concurrent.success_times:
        baseline_avg = warm.avg if (warm and warm.success_times) else (cold.avg if cold and cold.success_times else 0)
        if baseline_avg > 0:
            slowdown = concurrent.avg / baseline_avg
            print(f"\n   👥 Concurrent ({len(concurrent.metrics)} users) Avg: {concurrent.avg:.2f}s")
            print(f"   👤 Single user Avg: {baseline_avg:.2f}s")
            print(f"   📈 Concurrency penalty: {slowdown:.1f}x")

            if slowdown > 2:
                print(f"   🔴 CRITICAL: {slowdown:.1f}x chậm hơn khi multi-user!")
                print(f"      → Có thể do: Connection pool exhaustion, GIL, rate limit LLM API")
            elif slowdown > 1.3:
                print(f"   🟡 Moderate penalty: Hệ thống chịu tải OK nhưng chậm hơn")
            else:
                print(f"   🟢 Minimal penalty: Hệ thống scale tốt!")

    # Overall verdict
    all_times = []
    for p in phases:
        all_times.extend(p.success_times)

    if all_times:
        overall_avg = statistics.mean(all_times)
        print(f"\n   📊 Overall Avg: {overall_avg:.2f}s across {len(all_times)} queries")

        if overall_avg > 15:
            print("   🔴 VERDICT: CHẬM - Cần tối ưu urgently")
        elif overall_avg > 10:
            print("   🟡 VERDICT: CHẤP NHẬN ĐƯỢC - Multi-tool + LLM flow")
        elif overall_avg > 5:
            print("   🟢 VERDICT: TỐT cho agentic chatbot")
        else:
            print("   ⭐ VERDICT: XUẤT SẮC!")

    print("\n" + "=" * 75)


# =====================================================================
# MAIN
# =====================================================================
async def main():
    parser = argparse.ArgumentParser(description="CANIFA Chatbot Stress Profiler")
    parser.add_argument("--users", type=int, default=3, help="Concurrent users (default: 3)")
    parser.add_argument("--queries", type=int, default=5, help="Queries per phase (default: 5)")
    parser.add_argument("--skip-clear", action="store_true", help="Skip cache clearing")
    parser.add_argument("--warm-only", action="store_true", help="Only test warm (skip cold)")
    parser.add_argument("--output", type=str, help="Save JSON report")
    args = parser.parse_args()

    queries = TEST_QUERIES[:args.queries]

    print("\n" + "=" * 60)
    print("🔬 CANIFA CHATBOT - STRESS PROFILER")
    print("=" * 60)
    print(f"🌐 Target:        {API_URL}")
    print(f"📝 Queries/phase: {len(queries)}")
    print(f"👥 Concurrent:    {args.users} users")
    print(f"🧹 Clear cache:   {'No' if args.skip_clear else 'Yes'}")
    print(f"{'='*60}\n")

    # Health check
    async with httpx.AsyncClient() as client:
        try:
            r = await client.get(f"{API_URL}/health", timeout=5.0)
            if r.status_code != 200:
                logger.error(f"❌ Server unhealthy: {r.status_code}")
                return
            logger.info(f"✅ Server healthy: {API_URL}")
        except Exception as e:
            logger.error(f"❌ Không kết nối được: {e}")
            return

    phases = []

    # ── Phase 0: Cache before ──
    cache_before = await get_cache_stats()
    logger.info(f"📦 Cache trước test: {cache_before}")

    # ── Phase 1: COLD TEST (clear cache first) ──
    if not args.warm_only:
        if not args.skip_clear:
            logger.info("\n🧹 CLEARING CACHE...")
            clear_result = await clear_redis_cache()
            logger.info(f"   Result: {clear_result}")
            await asyncio.sleep(1)

        logger.info("\n🧊 PHASE 1: COLD TEST (1 user, no cache)")
        cold_result = await run_sequential(queries, "cold_tester", "🧊 Cold (1 user)")
        phases.append(cold_result)

        await asyncio.sleep(3)  # Let system settle

    # ── Phase 2: WARM TEST (cache populated from Phase 1) ──
    logger.info("\n🔥 PHASE 2: WARM TEST (1 user, with cache)")
    warm_result = await run_sequential(queries, "warm_tester", "🔥 Warm (1 user)")
    phases.append(warm_result)

    await asyncio.sleep(3)

    # ── Phase 3: CONCURRENT TEST ──
    logger.info(f"\n👥 PHASE 3: CONCURRENT TEST ({args.users} users)")
    concurrent_result = await run_concurrent(queries, args.users, f"👥 Concurrent ({args.users} users)")
    phases.append(concurrent_result)

    # ── Cache after ──
    cache_after = await get_cache_stats()

    # ── Report ──
    print_full_report(phases, cache_before, cache_after)

    # ── Save JSON ──
    if args.output:
        output_data = {
            "cache_before": cache_before,
            "cache_after": cache_after,
            "phases": [
                {
                    "phase": p.phase,
                    "avg": p.avg,
                    "p50": p.p50,
                    "p95": p.p95,
                    "min": p.min_t,
                    "max": p.max_t,
                    "success": p.success_count,
                    "errors": p.error_count,
                    "queries": [
                        {
                            "query": m.query,
                            "user": m.user_id,
                            "time": round(m.total_time, 3),
                            "status": m.status_code,
                            "error": m.error,
                            "products": m.has_products,
                        }
                        for m in p.metrics
                    ],
                }
                for p in phases
            ],
        }
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(output_data, f, indent=2, ensure_ascii=False)
        print(f"\n📄 JSON report: {args.output}")

    print("\n✅ Stress test hoàn tất!\n")


if __name__ == "__main__":
    if platform.system() == "Windows":
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    asyncio.run(main())
