"""
🔬 CANIFA Chatbot - Bottleneck Profiler
========================================
Script đo bottleneck end-to-end cho chatbot API.

Đo:
  1. API response time (tổng)
  2. Time-to-first-token (TTFT) cho streaming
  3. Tool execution latency
  4. Memory & CPU usage
  5. Async event loop lag

Dùng:
  python tests/profiler_bottleneck.py                         # Chạy mặc định
  python tests/profiler_bottleneck.py --queries 5 --url http://localhost:8001
  python tests/profiler_bottleneck.py --profile cprofile      # Deep profile với cProfile
  python tests/profiler_bottleneck.py --profile pyinstrument  # Flame graph
"""

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

# Add project root
sys.path.insert(0, str(Path(__file__).parent.parent))

try:
    import httpx
except ImportError:
    print("❌ Cần cài httpx: pip install httpx")
    sys.exit(1)

try:
    import psutil
except ImportError:
    psutil = None
    print("⚠️  psutil không có, bỏ qua CPU/Memory monitoring")


# =====================================================================
# CONFIG
# =====================================================================
DEFAULT_URL = "http://localhost:5000"
DEFAULT_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",
]

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


# =====================================================================
# DATA MODELS
# =====================================================================
@dataclass
class QueryResult:
    query: str
    total_time: float = 0.0
    ttft: float = 0.0  # Time to first token
    token_count: int = 0
    status_code: int = 0
    error: str = ""
    has_products: bool = False
    memory_before_mb: float = 0.0
    memory_after_mb: float = 0.0
    cpu_percent: float = 0.0
    chunks_received: int = 0


@dataclass
class ProfileReport:
    results: list[QueryResult] = field(default_factory=list)
    total_duration: float = 0.0
    system_info: dict = field(default_factory=dict)

    @property
    def avg_response_time(self) -> float:
        times = [r.total_time for r in self.results if r.status_code == 200]
        return statistics.mean(times) if times else 0

    @property
    def p50_response_time(self) -> float:
        times = sorted([r.total_time for r in self.results if r.status_code == 200])
        return times[len(times) // 2] if times else 0

    @property
    def p95_response_time(self) -> float:
        times = sorted([r.total_time for r in self.results if r.status_code == 200])
        idx = int(len(times) * 0.95)
        return times[min(idx, len(times) - 1)] if times else 0

    @property
    def avg_ttft(self) -> float:
        times = [r.ttft for r in self.results if r.ttft > 0]
        return statistics.mean(times) if times else 0


# =====================================================================
# PROFILER CORE
# =====================================================================
class BottleneckProfiler:
    """Profiler cho CANIFA Chatbot API."""

    def __init__(self, base_url: str, conversation_id: str = "profiler_test"):
        self.base_url = base_url.rstrip("/")
        self.conversation_id = conversation_id
        self.process = psutil.Process() if psutil else None

    def _get_memory_mb(self) -> float:
        if self.process:
            return self.process.memory_info().rss / (1024 * 1024)
        return 0

    def _get_cpu_percent(self) -> float:
        if self.process:
            return self.process.cpu_percent(interval=0.1)
        return 0

    async def profile_streaming_query(self, query: str, client: httpx.AsyncClient) -> QueryResult:
        """Profile một query với streaming response."""
        result = QueryResult(query=query)
        result.memory_before_mb = self._get_memory_mb()

        payload = {
            "user_query": query,
            "user_id": "profiler_bot",
        }

        start = time.perf_counter()
        first_token_time = None

        try:
            async with client.stream(
                "POST",
                f"{self.base_url}/api/agent/chat-dev",
                json=payload,
                timeout=120.0,
            ) as response:
                result.status_code = response.status_code

                if response.status_code != 200:
                    result.error = f"HTTP {response.status_code}"
                    result.total_time = time.perf_counter() - start
                    return result

                full_content = ""
                async for chunk in response.aiter_text():
                    if first_token_time is None and chunk.strip():
                        first_token_time = time.perf_counter()
                        result.ttft = first_token_time - start

                    result.chunks_received += 1
                    full_content += chunk

                result.total_time = time.perf_counter() - start

                # Check for product IDs in response
                result.has_products = "product_ids" in full_content.lower() or "sku" in full_content.lower()
                result.token_count = len(full_content.split())

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

        result.memory_after_mb = self._get_memory_mb()
        result.cpu_percent = self._get_cpu_percent()

        return result

    async def profile_non_streaming_query(self, query: str, client: httpx.AsyncClient) -> QueryResult:
        """Profile một query với non-streaming response (fallback)."""
        result = QueryResult(query=query)
        result.memory_before_mb = self._get_memory_mb()

        payload = {
            "user_query": query,
            "user_id": "profiler_bot",
        }

        start = time.perf_counter()

        try:
            response = await client.post(
                f"{self.base_url}/api/agent/chat-dev",
                json=payload,
                timeout=120.0,
            )
            result.total_time = time.perf_counter() - start
            result.status_code = response.status_code
            result.ttft = result.total_time  # Non-streaming: TTFT = total

            if response.status_code == 200:
                body = response.text
                result.has_products = "product_ids" in body.lower() or "sku" in body.lower()
                result.token_count = len(body.split())

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

        result.memory_after_mb = self._get_memory_mb()
        result.cpu_percent = self._get_cpu_percent()

        return result

    async def run(self, queries: list[str], use_streaming: bool = True) -> ProfileReport:
        """Chạy profiling cho danh sách queries."""
        report = ProfileReport()
        report.system_info = self._collect_system_info()

        overall_start = time.perf_counter()

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

            # Run queries sequentially
            for i, query in enumerate(queries, 1):
                logger.info(f"\n{'='*60}")
                logger.info(f"📝 Query {i}/{len(queries)}: {query}")
                logger.info(f"{'='*60}")

                if use_streaming:
                    result = await self.profile_streaming_query(query, client)
                else:
                    result = await self.profile_non_streaming_query(query, client)

                report.results.append(result)

                # Log kết quả ngay
                status = "✅" if not result.error else "❌"
                logger.info(f"{status} Total: {result.total_time:.2f}s | "
                           f"TTFT: {result.ttft:.2f}s | "
                           f"Chunks: {result.chunks_received} | "
                           f"Tokens: ~{result.token_count}")

                if result.error:
                    logger.info(f"   ⚠️  Error: {result.error}")

                if psutil:
                    mem_delta = result.memory_after_mb - result.memory_before_mb
                    logger.info(f"   💾 Memory: {result.memory_after_mb:.1f}MB "
                               f"(Δ{mem_delta:+.1f}MB) | CPU: {result.cpu_percent:.1f}%")

                # Nghỉ giữa các queries để không quá tải
                if i < len(queries):
                    await asyncio.sleep(2)

        report.total_duration = time.perf_counter() - overall_start
        return report

    def _collect_system_info(self) -> dict:
        info = {
            "platform": platform.platform(),
            "python": platform.python_version(),
            "cpu_count": os.cpu_count(),
        }
        if psutil:
            vm = psutil.virtual_memory()
            info["total_ram_gb"] = round(vm.total / (1024**3), 1)
            info["available_ram_gb"] = round(vm.available / (1024**3), 1)
        return info


# =====================================================================
# EVENT LOOP LAG DETECTOR
# =====================================================================
class EventLoopLagMonitor:
    """Đo async event loop lag - phát hiện blocking code."""

    def __init__(self, threshold_ms: float = 100):
        self.threshold_ms = threshold_ms
        self.lags: list[float] = []
        self._running = False

    async def start(self):
        self._running = True
        while self._running:
            t1 = time.perf_counter()
            await asyncio.sleep(0.01)  # 10ms expected
            t2 = time.perf_counter()
            lag_ms = (t2 - t1 - 0.01) * 1000
            if lag_ms > self.threshold_ms:
                self.lags.append(lag_ms)
                logger.warning(f"⚡ Event loop lag: {lag_ms:.0f}ms (threshold: {self.threshold_ms}ms)")

    def stop(self):
        self._running = False

    @property
    def summary(self) -> str:
        if not self.lags:
            return "✅ Không phát hiện event loop lag"
        return (
            f"⚠️  {len(self.lags)} lần lag > {self.threshold_ms}ms | "
            f"Max: {max(self.lags):.0f}ms | Avg: {statistics.mean(self.lags):.0f}ms"
        )


# =====================================================================
# REPORT GENERATOR
# =====================================================================
def print_report(report: ProfileReport):
    """In báo cáo bottleneck đẹp."""
    print("\n" + "=" * 70)
    print("🔬 BÁO CÁO BOTTLENECK PROFILING - CANIFA CHATBOT")
    print("=" * 70)

    # System info
    si = report.system_info
    print(f"\n📦 System: {si.get('platform', 'N/A')}")
    print(f"   Python: {si.get('python', 'N/A')} | CPUs: {si.get('cpu_count', 'N/A')}")
    if "total_ram_gb" in si:
        print(f"   RAM: {si['available_ram_gb']}GB free / {si['total_ram_gb']}GB total")

    # Per-query results
    print(f"\n{'─'*70}")
    print(f"{'Query':<35} {'Total':>8} {'TTFT':>8} {'Chunks':>8} {'Status':>8}")
    print(f"{'─'*70}")

    for r in report.results:
        status = "✅" if not r.error else "❌"
        q = r.query[:33] + ".." if len(r.query) > 33 else r.query
        print(f"{q:<35} {r.total_time:>7.2f}s {r.ttft:>7.2f}s {r.chunks_received:>7} {status:>8}")

    # Summary stats
    print(f"\n{'='*70}")
    print("📊 TỔNG KẾT")
    print(f"{'='*70}")

    successful = [r for r in report.results if r.status_code == 200]
    failed = [r for r in report.results if r.error]

    print(f"\n   Queries:     {len(report.results)} total | {len(successful)} thành công | {len(failed)} lỗi")
    print(f"   Tổng thời gian: {report.total_duration:.1f}s")

    if successful:
        times = [r.total_time for r in successful]
        print(f"\n   ⏱️  Response Time:")
        print(f"      Average:  {report.avg_response_time:.2f}s")
        print(f"      P50:      {report.p50_response_time:.2f}s")
        print(f"      P95:      {report.p95_response_time:.2f}s")
        print(f"      Min:      {min(times):.2f}s")
        print(f"      Max:      {max(times):.2f}s")

        ttfts = [r.ttft for r in successful if r.ttft > 0]
        if ttfts:
            print(f"\n   🚀 Time-to-First-Token (TTFT):")
            print(f"      Average:  {report.avg_ttft:.2f}s")
            print(f"      Min:      {min(ttfts):.2f}s")
            print(f"      Max:      {max(ttfts):.2f}s")

        if psutil:
            mems = [r.memory_after_mb for r in successful]
            cpus = [r.cpu_percent for r in successful if r.cpu_percent > 0]
            print(f"\n   💾 Resource Usage:")
            print(f"      Memory:   {min(mems):.0f} - {max(mems):.0f} MB")
            if cpus:
                print(f"      CPU:      {min(cpus):.0f} - {max(cpus):.0f}%")

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

    if report.avg_response_time > 15:
        print("\n   🔴 CRITICAL: Response time > 15s")
        print("      → Kiểm tra: LLM model latency, tool execution, network")
    elif report.avg_response_time > 10:
        print("\n   🟡 WARNING: Response time 10-15s")
        print("      → Bình thường cho multi-tool + LLM generation")
        print("      → Tối ưu: cache warming, embedding local, smaller model")
    elif report.avg_response_time > 5:
        print("\n   🟢 GOOD: Response time 5-10s")
        print("      → Tốt cho agentic flow có tool calls")
    else:
        print("\n   ⭐ EXCELLENT: Response time < 5s")

    if report.avg_ttft > 8:
        print("\n   🔴 TTFT > 8s: User phải đợi quá lâu")
        print("      → Xem xét: streaming optimization, early return pattern")
    elif report.avg_ttft > 4:
        print("\n   🟡 TTFT 4-8s: Chấp nhận được nhưng có thể cải thiện")
    else:
        print("\n   🟢 TTFT < 4s: Tốt!")

    # Recommendations
    print(f"\n{'='*70}")
    print("💡 KHUYẾN NGHỊ PYTHON PROFILING TOOLS")
    print(f"{'='*70}")
    print("""
   📦 Đã test xong E2E bottleneck. Để đi sâu hơn, dùng các tools sau:

   1. py-spy (Sampling profiler - KHÔNG cần sửa code)
      pip install py-spy
      py-spy top --pid <PID>           # Real-time CPU profiling
      py-spy record -o flame.svg -- python server.py  # Flame graph

   2. pyinstrument (Call-stack profiler - SIÊU DỄ DÙNG)
      pip install pyinstrument
      # Wrap trong code:
      from pyinstrument import Profiler
      profiler = Profiler(async_mode="enabled")
      profiler.start()
      # ... code ...
      profiler.stop()
      profiler.print()  # Hoặc profiler.output_html()

   3. yappi (Async-aware profiler - TỐT NHẤT cho asyncio)
      pip install yappi
      yappi.set_clock_type("wall")  # Real time, không chỉ CPU
      yappi.start()
      # ... run queries ...
      yappi.stop()
      yappi.get_func_stats().print_all()

   4. scalene (AI-powered profiler - CPU + Memory + GPU)
      pip install scalene
      scalene server.py   # Tự analyze bottleneck bằng AI!

   5. memray (Memory profiler - từ Bloomberg)
      pip install memray
      memray run server.py
      memray flamegraph output.bin  # Memory flame graph

   6. OpenTelemetry (Distributed tracing - ĐÃ CÓ TRONG requirements.txt)
      → Đã có opentelemetry-sdk, chỉ cần enable traces!
    """)

    print("=" * 70)
    print("✅ Profiling hoàn tất!")
    print("=" * 70)


# =====================================================================
# MAIN
# =====================================================================
async def main():
    parser = argparse.ArgumentParser(description="CANIFA Chatbot Bottleneck Profiler")
    parser.add_argument("--url", default=DEFAULT_URL, help="Base URL of the API")
    parser.add_argument("--queries", type=int, default=5, help="Number of queries to test")
    parser.add_argument("--no-stream", action="store_true", help="Disable streaming")
    parser.add_argument("--custom-query", type=str, help="Test with a custom query")
    parser.add_argument("--profile", choices=["cprofile", "pyinstrument", "yappi"],
                       help="Enable deep profiling mode")
    parser.add_argument("--monitor-lag", action="store_true", help="Monitor event loop lag")
    parser.add_argument("--output", type=str, help="Save JSON report to file")

    args = parser.parse_args()

    # Build query list
    if args.custom_query:
        queries = [args.custom_query]
    else:
        queries = DEFAULT_QUERIES[: args.queries]

    print("\n" + "=" * 60)
    print("🔬 CANIFA CHATBOT - BOTTLENECK PROFILER")
    print("=" * 60)
    print(f"🌐 Target:  {args.url}")
    print(f"📝 Queries: {len(queries)}")
    print(f"📡 Mode:    {'Non-streaming' if args.no_stream else 'Streaming'}")
    if args.profile:
        print(f"🔍 Profiler: {args.profile}")
    print(f"{'='*60}\n")

    profiler = BottleneckProfiler(args.url)

    # Optional event loop lag monitor
    lag_monitor = None
    lag_task = None
    if args.monitor_lag:
        lag_monitor = EventLoopLagMonitor(threshold_ms=100)
        lag_task = asyncio.create_task(lag_monitor.start())

    # Optional deep profiling
    deep_profiler = None
    if args.profile == "pyinstrument":
        try:
            from pyinstrument import Profiler as PyProfiler
            deep_profiler = PyProfiler(async_mode="enabled")
            deep_profiler.start()
            logger.info("🔍 pyinstrument profiler started")
        except ImportError:
            logger.warning("⚠️  pyinstrument chưa cài: pip install pyinstrument")
    elif args.profile == "yappi":
        try:
            import yappi
            yappi.set_clock_type("wall")
            yappi.start()
            logger.info("🔍 yappi profiler started")
        except ImportError:
            logger.warning("⚠️  yappi chưa cài: pip install yappi")

    # Run profiling
    report = await profiler.run(queries, use_streaming=not args.no_stream)

    # Stop monitors
    if lag_monitor:
        lag_monitor.stop()
        if lag_task:
            lag_task.cancel()

    # Stop deep profilers
    if args.profile == "pyinstrument" and deep_profiler:
        deep_profiler.stop()
        print("\n" + "=" * 60)
        print("🔍 PYINSTRUMENT RESULTS")
        print("=" * 60)
        deep_profiler.print()
        # Save HTML report
        html = deep_profiler.output_html()
        with open("profiler_report.html", "w") as f:
            f.write(html)
        print("📄 HTML report saved: profiler_report.html")

    elif args.profile == "yappi":
        try:
            import yappi
            yappi.stop()
            print("\n" + "=" * 60)
            print("🔍 YAPPI TOP 30 FUNCTIONS (by total time)")
            print("=" * 60)
            stats = yappi.get_func_stats()
            stats.sort("ttot", "desc")
            stats.print_all(
                columns={
                    0: ("name", 60),
                    1: ("ncall", 8),
                    2: ("ttot", 10),
                    3: ("tsub", 10),
                    4: ("tavg", 10),
                },
                out=sys.stdout,
            )
        except ImportError:
            pass

    # Print report
    print_report(report)

    # Event loop lag summary
    if lag_monitor:
        print(f"\n⚡ Event Loop Lag: {lag_monitor.summary}")

    # Save JSON output
    if args.output:
        output_data = {
            "system": report.system_info,
            "summary": {
                "total_queries": len(report.results),
                "avg_response_time": report.avg_response_time,
                "p50": report.p50_response_time,
                "p95": report.p95_response_time,
                "avg_ttft": report.avg_ttft,
                "total_duration": report.total_duration,
            },
            "results": [
                {
                    "query": r.query,
                    "total_time": r.total_time,
                    "ttft": r.ttft,
                    "tokens": r.token_count,
                    "chunks": r.chunks_received,
                    "status": r.status_code,
                    "error": r.error,
                    "has_products": r.has_products,
                    "memory_mb": r.memory_after_mb,
                    "cpu_percent": r.cpu_percent,
                }
                for r in report.results
            ],
        }
        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 saved: {args.output}")


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