"""
LangGraph-to-WebSocket streaming bridge.

Transforms LangGraph astream() state updates into WebSocket events
that the dashboard can consume. Also provides console output formatting.
"""

import aiohttp
from typing import Any, Dict, List, Union
from datetime import datetime, timezone
from src.utils.logging import get_logger
from src.utils.config import settings

logger = get_logger(__name__)


# Map LangGraph node names to dashboard event types
NODE_EVENT_MAP: Dict[str, str] = {
    "reconnaissance_search": "reconnaissance_started",
    "coordinate_agents": "agent_consultation",
    "generate_attack": "attack_generated",
    "execute_attack": "response_received",
    "analyze_response": "finding_detected",
    "record_metrics": "attack_graph_updated",
    "summarize_if_needed": "summarization_check",
    "update_session_summary": "session_summary_updated",
    "calculate_score": "campaign_phase_update",
    "critical_finding_handler": "critical_finding_review",
    "generate_report": "test_completed",
}


def extract_event_data(
    node_name: str, state_update: Dict[str, Any]
) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
    """
    Extract relevant data from a state update for WebSocket transmission.

    Filters out non-serializable and oversized fields so only
    meaningful data reaches the dashboard.

    Args:
        node_name: The LangGraph node that produced this update
        state_update: The state diff returned by the node

    Returns:
        Filtered event data suitable for JSON serialization
    """
    if node_name == "coordinate_agents":
        consultation = state_update.get("agent_consultation", {})
        vote_exp = consultation.get("vote_explanation") or {}
        winner_info = vote_exp.get("winner") or {}
        competition = vote_exp.get("competition") or {}

        reasoning = winner_info.get("reasoning", "")
        if not reasoning:
            votes = consultation.get("all_votes", [])
            chosen = consultation.get("chosen_agent", "")
            for v in votes:
                v_dict = v if isinstance(v, dict) else v.dict()
                if v_dict.get("agent_name") == chosen:
                    reasoning = v_dict.get("reasoning", "")
                    break

        all_scores = vote_exp.get("all_scores") or []

        return {
            "winner": consultation.get("chosen_agent", "unknown"),
            "score": consultation.get("winning_score", 0),
            "confidence": consultation.get("confidence", 0),
            "method": consultation.get("method", "unknown"),
            "reasoning": str(reasoning)[:500],
            "pattern": winner_info.get("pattern", ""),
            "margin": competition.get("margin", ""),
            "second_place": competition.get("second_place", ""),
            "second_score": competition.get("second_score", 0),
            "all_scores": all_scores[:6],
        }

    elif node_name == "generate_attack":
        attacks = state_update.get("attack_attempts", [])
        if attacks:
            attack = attacks[-1]
            agent_name = attack.get("agent_name", "unknown")
            return {
                "pattern": attack.get("metadata", {}).get("pattern", "unknown"),
                "query": attack.get("query", ""),
                "query_preview": attack.get("query", "")[:150],
                "agent": agent_name,
                "agent_name": agent_name,  # Dashboard checks both agent_name and agent
                "attack_type": attack.get("attack_type", "unknown"),
            }
        return {}

    elif node_name == "execute_attack":
        responses = state_update.get("target_responses", [])
        if responses:
            resp = responses[-1]
            content = resp.get("content", "")
            ts = resp.get("timestamp")
            return {
                "content": content,
                "content_preview": content[:400],
                "timestamp": (
                    ts.isoformat() if ts is not None and hasattr(ts, "isoformat") else str(ts or "")
                ),
            }
        return {"content": state_update.get("last_response", "")[:400]}

    elif node_name == "analyze_response":
        findings = state_update.get("security_findings", [])
        # Dashboard expects data to be an array of finding objects (or a single object).
        # It does: Array.isArray(data.data) ? data.data : [data.data]
        # Each finding needs: severity, category, description, evidence,
        #                     attack_query (optional), target_response (optional)
        if findings:
            return [
                {
                    "severity": f.get("severity", "info"),
                    "category": f.get("category", "unknown"),
                    "description": f.get("description", ""),
                    "evidence": f.get("evidence", ""),
                    "confidence": f.get("confidence", 0),
                    "attack_query": f.get("attack_query", ""),
                    "target_response": f.get("target_response", ""),
                    "agent_name": f.get("agent_name", ""),
                }
                for f in findings
            ]
        # No findings — return empty array so dashboard skips gracefully
        return []

    elif node_name == "calculate_score":
        return {
            "vulnerability_score": state_update.get("vulnerability_score", 0),
            "campaign_phase": state_update.get("campaign_phase"),
        }

    elif node_name == "generate_report":
        ca = state_update.get("completed_at")
        findings = state_update.get("security_findings", [])
        return {
            "status": state_update.get("test_status", "completed"),
            "vulnerability_score": state_update.get("vulnerability_score", 0),
            "total_findings": len(findings),
            "completed_at": (
                ca.isoformat() if ca is not None and hasattr(ca, "isoformat") else str(ca or "")
            ),
        }

    elif node_name == "record_metrics":
        graph = state_update.get("simple_attack_graph")
        if graph:
            return _format_graph_for_visjs(graph)
        return {"nodes": [], "edges": []}

    # Default: return a safe subset of keys
    safe_keys = {"test_status", "vulnerability_score", "current_attempt", "campaign_phase"}
    return {k: v for k, v in state_update.items() if k in safe_keys}


# vis.js color palettes (matching websocket.py graph endpoint)
_NODE_COLORS = {
    "initial": "#00D4FF",
    "probe": "#888888",
    "partial": "#FFB86C",
    "jailbreak": "#FF79C6",
    "exploit": "#FF5555",
    "goal": "#50FA7B",
    "dead_end": "#6272A4",
    "backtrack": "#BD93F9",
}

_EDGE_COLORS = {
    "success": "#50FA7B",
    "partial": "#FFB86C",
    "blocked": "#FF5555",
    "probe": "#888888",
}


def _format_graph_for_visjs(graph: Dict[str, Any]) -> Dict[str, Any]:
    """Convert simple_attack_graph nodes/edges into vis.js format.

    The dashboard's ``updateAttackGraph(data)`` expects ``data.nodes``
    and ``data.edges`` as arrays of vis.js-ready objects with colors,
    shapes, and tooltips.
    """
    vis_nodes = []
    for node in graph.get("nodes", []):
        ntype = node.get("type", "probe")
        bg = _NODE_COLORS.get(ntype, "#666666")
        vis_nodes.append(
            {
                "id": node["id"],
                "label": node.get("label", node["id"]),
                "type": ntype,
                "color": {"background": bg, "border": bg},
                "shape": (
                    "diamond" if ntype == "initial" else ("star" if ntype == "goal" else "dot")
                ),
                "size": 15 + node.get("visit_count", 0) * 2,
                "title": (
                    f"Type: {ntype}\n"
                    f"Visits: {node.get('visit_count', 0)}\n"
                    f"Success: {node.get('success_rate', 0):.0%}"
                ),
            }
        )

    vis_edges = []
    for edge in graph.get("edges", []):
        outcome = edge.get("outcome", "probe")
        vis_edges.append(
            {
                "id": edge.get("id", f"{edge['from']}-{edge['to']}"),
                "from": edge["from"],
                "to": edge["to"],
                "label": edge.get("pattern", edge.get("attack_type", ""))[:20],
                "color": _EDGE_COLORS.get(outcome, "#444444"),
                "arrows": "to",
                "title": (
                    f"Agent: {edge.get('agent', 'unknown')}\n"
                    f"Pattern: {edge.get('pattern', 'unknown')}\n"
                    f"Outcome: {outcome}"
                ),
            }
        )

    return {"nodes": vis_nodes, "edges": vis_edges}


async def emit_ws_event_http(session_id: str, event_type: str, data: Any) -> None:
    """
    Publish an event to the dashboard via the HTTP publish endpoint.

    This works when the streaming bridge runs in a separate process
    from the FastAPI server. Failures are swallowed (debug logged)
    so tests continue.

    Args:
        session_id: Test session identifier
        event_type: WebSocket event type
        data: Event data
    """
    api_port = getattr(settings, "api_port", 8000)
    url = f"http://localhost:{api_port}/api/v1/ws/publish/{session_id}"
    payload = {
        "type": event_type,
        "data": data,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    try:
        timeout = aiohttp.ClientTimeout(total=5)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.post(url, json=payload) as resp:
                await resp.read()
    except Exception as e:
        logger.debug("ws_publish_failed", error=str(e))


async def stream_workflow_to_websocket(
    session_id: str,
    node_name: str,
    state_update: Dict[str, Any],
) -> None:
    """
    Bridge a single LangGraph state update to the WebSocket dashboard.

    Called inside the ``async for chunk in app.astream(...)`` loop
    for each node completion.

    Args:
        session_id: Test session identifier
        node_name: The LangGraph node that just completed
        state_update: The state diff produced by the node
    """
    event_type = NODE_EVENT_MAP.get(node_name, node_name)
    event_data = extract_event_data(node_name, state_update)

    # Skip emitting finding_detected when there are no findings
    # (analyze_response runs every round, but findings are rare)
    if event_type == "finding_detected" and isinstance(event_data, list) and len(event_data) == 0:
        pass  # No findings — don't send empty event
    else:
        await emit_ws_event_http(session_id, event_type, event_data)

    # Cache individual findings for report generation via API
    if node_name == "analyze_response":
        findings = state_update.get("security_findings", [])
        for finding in findings:
            if finding.get("finding_id"):
                try:
                    await _cache_finding(finding)
                except Exception:
                    pass  # Non-critical — dashboard may not be running


async def _cache_finding(finding: Dict[str, Any]) -> None:
    """Cache a single finding via the reports API endpoint."""
    api_port = getattr(settings, "api_port", 8000)
    url = f"http://localhost:{api_port}/api/v1/reports/findings/cache"
    try:
        timeout = aiohttp.ClientTimeout(total=3)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.post(url, json=finding) as resp:
                await resp.read()
    except Exception as e:
        logger.debug("finding_cache_failed", error=str(e))


# ---------------------------------------------------------------------------
# Console output formatting (replaces the print() calls in test_orchestrated)
# ---------------------------------------------------------------------------


def print_node_output(node_name: str, state_update: Dict[str, Any]) -> None:
    """
    Print human-readable console output for a completed LangGraph node.

    Mirrors the rich console output from test_orchestrated.py
    so CLI users get the same experience.

    Args:
        node_name: The LangGraph node that completed
        state_update: The state diff from that node
    """
    if node_name == "reconnaissance_search":
        print("  Tavily reconnaissance complete")

    elif node_name == "coordinate_agents":
        consultation = state_update.get("agent_consultation", {})
        winner = consultation.get("chosen_agent", "unknown")
        score = consultation.get("winning_score", 0)
        print(f"  CONSENSUS: {winner} (score: {score:.2f})")

    elif node_name == "generate_attack":
        attacks = state_update.get("attack_attempts", [])
        if attacks:
            att = attacks[-1]
            pattern = att.get("metadata", {}).get("pattern", "unknown")
            agent = att.get("agent_name", "unknown")
            query = att.get("query", "")
            print(f"  Attack generated: {pattern} (agent: {agent})")
            if query:
                preview = query[:500] + "..." if len(query) > 500 else query
                print(f"  Query ({len(query)} chars):\n    {preview}")

    elif node_name == "execute_attack":
        content = state_update.get("last_response", "")
        preview = content[:500] + "..." if len(content) > 500 else content
        print(f"  Response ({len(content)} chars):\n    {preview}")

    elif node_name == "analyze_response":
        findings = state_update.get("security_findings", [])
        refusal = state_update.get("last_refusal_level", "unknown")
        if findings:
            for f in findings:
                sev = f.get("severity", "?").upper()
                cat = f.get("category", "?")
                print(f"  [{sev}] {cat}: {f.get('description', '')[:80]}")
        else:
            print(f"  No findings (refusal: {refusal})")

    elif node_name == "record_metrics":
        phase = state_update.get("campaign_phase")
        if phase:
            print(f"  Campaign phase: {phase}")

    elif node_name == "calculate_score":
        score = state_update.get("vulnerability_score", 0)
        print(f"  Vulnerability score: {score:.1f}")

    elif node_name == "critical_finding_handler":
        status = state_update.get("test_status")
        if status == "stopped":
            print("  Critical finding review: TEST STOPPED by reviewer")
        else:
            print("  Critical finding review: Continuing to report")

    elif node_name == "generate_report":
        status = state_update.get("test_status", "completed")
        score = state_update.get("vulnerability_score", 0)
        print(f"  Test {status} — final score: {score:.1f}")

    else:
        # Unmapped nodes — show a generic update line
        pass  # Avoid noisy output for utility nodes
