"""
Conversation Test API Routes
- POST /api/test/upload     — Upload Excel/CSV file, parse conversations
- POST /api/test/run        — Run test batch with fake device_ids
- GET  /api/test/results    — Get stored test results

Tool to test chatbot responses across multiple conversation scenarios × versions.
Each conversation_id × version gets a unique fake device_id.
"""

import asyncio
import io
import json
import logging
import time
import uuid
from collections import defaultdict
from typing import Any

import httpx
from fastapi import APIRouter, UploadFile, File
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

logger = logging.getLogger(__name__)
router = APIRouter(tags=["Conversation Test"])

# In-memory store for test results (keyed by test_id)
_test_results: dict[str, Any] = {}


# ─────────────────────────────────────────────────────────
# Models
# ─────────────────────────────────────────────────────────
class TestRunRequest(BaseModel):
    endpoint_url: str = "http://172.16.2.207:5000/api/agent/chat-dev"
    conversations: dict[str, list[str]]  # conv_id -> [messages]
    num_versions: int = 3
    delay_ms: int = 1000
    selected_conv_ids: list[str] | None = None  # If None, run all


# ─────────────────────────────────────────────────────────
# 1. Upload & Parse
# ─────────────────────────────────────────────────────────
@router.post("/api/test/upload", summary="Upload Excel/CSV file for conversation test")
async def upload_test_file(file: UploadFile = File(...)):
    """
    Parse an uploaded Excel (.xlsx) or CSV file.
    Expects columns: message_content, conversation_id_test
    Returns grouped conversations.
    """
    try:
        content = await file.read()
        filename = file.filename or ""

        conversations: dict[str, list[str]] = defaultdict(list)

        if filename.endswith(".xlsx"):
            import openpyxl

            wb = openpyxl.load_workbook(io.BytesIO(content), read_only=True)

            # Try to find 'conversation' sheet, else use active
            if "conversation" in wb.sheetnames:
                ws = wb["conversation"]
            else:
                ws = wb.active

            # Find header row
            headers = {}
            for col_idx in range(1, (ws.max_column or 10) + 1):
                val = ws.cell(1, col_idx).value
                if val:
                    headers[str(val).strip().lower()] = col_idx

            msg_col = headers.get("message_content", 1)
            conv_col = headers.get("conversation_id_test", 2)

            for row_idx in range(2, (ws.max_row or 2) + 1):
                msg = ws.cell(row_idx, msg_col).value
                conv_id = ws.cell(row_idx, conv_col).value

                if msg and conv_id is not None:
                    conversations[str(int(conv_id))].append(str(msg))

            wb.close()

        elif filename.endswith(".csv"):
            import csv

            text = content.decode("utf-8-sig")
            reader = csv.DictReader(io.StringIO(text))
            for row in reader:
                msg = row.get("message_content", "").strip()
                conv_id = row.get("conversation_id_test", "").strip()
                if msg and conv_id:
                    conversations[conv_id].append(msg)
        else:
            return {"status": "error", "message": f"Unsupported file type: {filename}. Use .xlsx or .csv"}

        # Sort by conv_id numerically
        sorted_convs = dict(
            sorted(conversations.items(), key=lambda x: int(x[0]) if x[0].isdigit() else x[0])
        )

        return {
            "status": "success",
            "total_conversations": len(sorted_convs),
            "total_messages": sum(len(msgs) for msgs in sorted_convs.values()),
            "conversations": sorted_convs,
        }

    except Exception as e:
        logger.error(f"Error parsing test file: {e}", exc_info=True)
        return {"status": "error", "message": str(e)}


# ─────────────────────────────────────────────────────────
# 2. Run Test (SSE streaming)
# ─────────────────────────────────────────────────────────
@router.post("/api/test/run", summary="Run conversation test batch")
async def run_test_batch(req: TestRunRequest):
    """
    Run test: for each conversation_id × version, generate a fake device_id,
    send messages sequentially to the target endpoint, and stream results via SSE.
    """
    test_id = str(uuid.uuid4())[:8]

    # Filter conversations if selected
    conv_ids = req.selected_conv_ids or list(req.conversations.keys())
    conversations = {cid: req.conversations[cid] for cid in conv_ids if cid in req.conversations}

    total_tasks = len(conversations) * req.num_versions
    _test_results[test_id] = {
        "test_id": test_id,
        "status": "running",
        "total_tasks": total_tasks,
        "completed": 0,
        "results": {},
    }

    async def event_stream():
        """SSE event stream for real-time progress."""
        # Send initial event
        yield _sse_event("start", {
            "test_id": test_id,
            "total_conversations": len(conversations),
            "num_versions": req.num_versions,
            "total_tasks": total_tasks,
        })

        completed = 0

        async with httpx.AsyncClient(timeout=120.0) as client:
            for conv_id in sorted(conversations.keys(), key=lambda x: int(x) if x.isdigit() else x):
                messages = conversations[conv_id]

                for version in range(1, req.num_versions + 1):
                    device_id = f"test-c{conv_id}-v{version}-{uuid.uuid4().hex[:6]}"
                    result_key = f"{conv_id}_v{version}"

                    try:
                        version_responses = []
                        total_time = 0.0

                        for msg_idx, message in enumerate(messages):
                            start_time = time.time()

                            # Send to target endpoint with fake device_id
                            tagged_message = f"{message} [v{version}]"
                            response = await client.post(
                                req.endpoint_url,
                                json={
                                    "user_query": tagged_message,
                                    "device_id": device_id,
                                },
                                headers={
                                    "Content-Type": "application/json",
                                    "device_id": device_id,
                                },
                            )

                            elapsed = time.time() - start_time
                            total_time += elapsed

                            msg_result = None
                            if response.status_code == 200:
                                data = response.json()
                                msg_result = {
                                    "message": message,
                                    "ai_response": data.get("ai_response", ""),
                                    "product_ids": data.get("product_ids", []),
                                    "response_time": round(elapsed, 2),
                                }
                            else:
                                msg_result = {
                                    "message": message,
                                    "ai_response": f"[ERROR {response.status_code}] {response.text[:200]}",
                                    "product_ids": [],
                                    "response_time": round(elapsed, 2),
                                }

                            version_responses.append(msg_result)

                            # Stream per-message update immediately
                            yield _sse_event("message_done", {
                                "conv_id": conv_id,
                                "version": version,
                                "msg_index": msg_idx,
                                "ai_response": msg_result["ai_response"],
                                "response_time": msg_result["response_time"],
                            })

                            # Delay between messages
                            if msg_idx < len(messages) - 1 and req.delay_ms > 0:
                                await asyncio.sleep(req.delay_ms / 1000)

                        # Store result for conv×version
                        result = {
                            "conv_id": conv_id,
                            "version": version,
                            "device_id": device_id,
                            "messages": messages,
                            "responses": version_responses,
                            "total_time": round(total_time, 2),
                            "final_response": version_responses[-1]["ai_response"] if version_responses else "",
                        }

                        _test_results[test_id]["results"][result_key] = result
                        completed += 1
                        _test_results[test_id]["completed"] = completed

                        # Stream conv×version completion
                        yield _sse_event("progress", {
                            "result_key": result_key,
                            "conv_id": conv_id,
                            "version": version,
                            "device_id": device_id,
                            "completed": completed,
                            "total_tasks": total_tasks,
                            "total_time": result["total_time"],
                            "responses": version_responses,
                            "num_messages": len(messages),
                        })

                    except Exception as e:
                        completed += 1
                        _test_results[test_id]["completed"] = completed
                        error_result = {
                            "conv_id": conv_id,
                            "version": version,
                            "device_id": device_id,
                            "messages": messages,
                            "responses": [],
                            "total_time": 0,
                            "final_response": f"[EXCEPTION] {str(e)}",
                            "error": str(e),
                        }
                        _test_results[test_id]["results"][result_key] = error_result

                        yield _sse_event("error", {
                            "result_key": result_key,
                            "conv_id": conv_id,
                            "version": version,
                            "error": str(e),
                            "completed": completed,
                            "total_tasks": total_tasks,
                        })

        _test_results[test_id]["status"] = "completed"
        yield _sse_event("complete", {
            "test_id": test_id,
            "total_completed": completed,
        })

    return StreamingResponse(event_stream(), media_type="text/event-stream")


# ─────────────────────────────────────────────────────────
# 3. Get Results
# ─────────────────────────────────────────────────────────
@router.get("/api/test/results/{test_id}", summary="Get test results")
async def get_test_results(test_id: str):
    """Retrieve stored test results by test_id."""
    if test_id not in _test_results:
        return {"status": "error", "message": f"Test ID {test_id} not found"}
    return _test_results[test_id]


@router.get("/api/test/results", summary="List all test results")
async def list_test_results():
    """List all stored test results (summary only)."""
    return {
        "tests": [
            {
                "test_id": tid,
                "status": data["status"],
                "total_tasks": data["total_tasks"],
                "completed": data["completed"],
            }
            for tid, data in _test_results.items()
        ]
    }


# ─────────────────────────────────────────────────────────
# 4. Clear Results
# ─────────────────────────────────────────────────────────
@router.delete("/api/test/results/clear", summary="Clear all test results")
async def clear_test_results():
    """Clear all stored test results from memory."""
    count = len(_test_results)
    _test_results.clear()
    return {"status": "success", "cleared": count, "message": f"Cleared {count} test result(s)"}


# ─────────────────────────────────────────────────────────
# Helper
# ─────────────────────────────────────────────────────────
def _sse_event(event_type: str, data: dict) -> str:
    """Format SSE event string."""
    return f"event: {event_type}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
