"""Normalize LangGraph events to Arzule TraceEvent format."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Dict, List, Optional

from ..ids import new_span_id
from ..sanitize import sanitize, truncate_string

if TYPE_CHECKING:
    from ..run import ArzuleRun


def _safe_getattr(obj: Any, attr: str, default: Any = None) -> Any:
    """Safely get attribute from object."""
    try:
        return getattr(obj, attr, default)
    except Exception:
        return default


def _base(
    run: "ArzuleRun",
    *,
    span_id: Optional[str],
    parent_span_id: Optional[str],
) -> Dict[str, Any]:
    """Build base event fields."""
    return {
        "schema_version": "trace_event.v0_1",
        "run_id": run.run_id,
        "tenant_id": run.tenant_id,
        "project_id": run.project_id,
        "trace_id": run.trace_id,
        "span_id": span_id or new_span_id(),
        "parent_span_id": parent_span_id,
        "seq": run.next_seq(),
        "ts": run.now(),
        "workstream_id": None,
        "task_id": None,
        "raw_ref": {"storage": "inline"},
    }


# =============================================================================
# Graph Events
# =============================================================================


def evt_graph_start(
    run: "ArzuleRun",
    span_id: str,
    parent_span_id: Optional[str],
    graph_name: str,
    input_data: Any,
    metadata: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Create event for graph execution start."""
    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "graph.start",
        "status": "ok",
        "summary": f"graph execution: {graph_name}",
        "attrs_compact": {
            "graph_name": graph_name,
            "metadata": metadata,
        },
        "payload": {
            "input": sanitize(input_data),
        },
    }


def evt_graph_end(
    run: "ArzuleRun",
    span_id: Optional[str],
    parent_span_id: Optional[str],
    graph_name: str,
    output_data: Any,
    error: Optional[Exception] = None,
) -> Dict[str, Any]:
    """Create event for graph execution end."""
    status = "error" if error else "ok"
    summary = f"graph completed: {graph_name}"
    if error:
        summary = f"graph error: {graph_name} - {str(error)[:100]}"

    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "graph.end",
        "status": status,
        "summary": summary,
        "attrs_compact": {
            "graph_name": graph_name,
            "error": str(error) if error else None,
        },
        "payload": {
            "output": sanitize(output_data) if not error else None,
            "error": str(error) if error else None,
        },
    }


# =============================================================================
# Node Events
# =============================================================================


def evt_node_start(
    run: "ArzuleRun",
    span_id: str,
    parent_span_id: Optional[str],
    node_name: str,
    input_data: Any,
    metadata: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Create event for node execution start."""
    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": node_name,  # Map node to agent
        "event_type": "agent.start",
        "status": "ok",
        "summary": f"node start: {node_name}",
        "attrs_compact": {
            "node_name": node_name,
            "metadata": metadata,
        },
        "payload": {
            "input": sanitize(input_data),
        },
    }


def evt_node_end(
    run: "ArzuleRun",
    span_id: Optional[str],
    parent_span_id: Optional[str],
    node_name: str,
    output_data: Any,
    error: Optional[Exception] = None,
) -> Dict[str, Any]:
    """Create event for node execution end."""
    status = "error" if error else "ok"
    summary = f"node completed: {node_name}"
    if error:
        summary = f"node error: {node_name} - {str(error)[:100]}"

    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": node_name,
        "event_type": "agent.end",
        "status": status,
        "summary": summary,
        "attrs_compact": {
            "node_name": node_name,
            "error": str(error) if error else None,
        },
        "payload": {
            "output": sanitize(output_data) if not error else None,
            "error": str(error) if error else None,
        },
    }


# =============================================================================
# Checkpoint Events
# =============================================================================


def evt_checkpoint_save(
    run: "ArzuleRun",
    span_id: str,
    parent_span_id: Optional[str],
    checkpoint_id: str,
    state: Any,
    metadata: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Create event for checkpoint save."""
    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "checkpoint.save",
        "status": "ok",
        "summary": f"checkpoint saved: {checkpoint_id[:16]}",
        "attrs_compact": {
            "checkpoint_id": checkpoint_id,
            "metadata": metadata,
        },
        "payload": {
            "state": sanitize(state),
        },
    }


def evt_checkpoint_load(
    run: "ArzuleRun",
    span_id: str,
    parent_span_id: Optional[str],
    checkpoint_id: str,
    state: Any,
) -> Dict[str, Any]:
    """Create event for checkpoint load."""
    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "checkpoint.load",
        "status": "ok",
        "summary": f"checkpoint loaded: {checkpoint_id[:16]}",
        "attrs_compact": {
            "checkpoint_id": checkpoint_id,
        },
        "payload": {
            "state": sanitize(state),
        },
    }


# =============================================================================
# Task Events (for parallel execution tracking)
# =============================================================================


def evt_task_start(
    run: "ArzuleRun",
    span_id: str,
    parent_span_id: Optional[str],
    task_id: str,
    task_name: str,
    task_input: Any,
) -> Dict[str, Any]:
    """Create event for task start (parallel execution)."""
    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "task.start",
        "status": "ok",
        "summary": f"task start: {task_name}",
        "attrs_compact": {
            "task_id": task_id,
            "task_name": task_name,
        },
        "payload": {
            "input": sanitize(task_input),
        },
    }


def evt_task_end(
    run: "ArzuleRun",
    span_id: Optional[str],
    parent_span_id: Optional[str],
    task_id: str,
    task_name: str,
    result: Any,
    error: Optional[Exception] = None,
) -> Dict[str, Any]:
    """Create event for task end."""
    status = "error" if error else "ok"
    summary = f"task completed: {task_name}"
    if error:
        summary = f"task error: {task_name} - {str(error)[:100]}"

    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "task.end",
        "status": status,
        "summary": summary,
        "attrs_compact": {
            "task_id": task_id,
            "task_name": task_name,
            "error": str(error) if error else None,
        },
        "payload": {
            "result": sanitize(result) if not error else None,
            "error": str(error) if error else None,
        },
    }


# =============================================================================
# LLM Events (from nested LangChain calls)
# =============================================================================


def evt_llm_start(
    run: "ArzuleRun",
    span_id: str,
    parent_span_id: Optional[str],
    model_name: str,
    prompts: Optional[List[str]],
    metadata: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Create event for LLM call start."""
    prompts = prompts or []

    # Truncate prompts for payload
    truncated_prompts = [truncate_string(p, 1000) for p in prompts[:5]]
    if len(prompts) > 5:
        truncated_prompts.append(f"... and {len(prompts) - 5} more prompts")

    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "llm.call.start",
        "status": "ok",
        "summary": f"llm call: {model_name}",
        "attrs_compact": {
            "model": model_name,
            "prompt_count": len(prompts),
        },
        "payload": {
            "prompts": truncated_prompts,
            "metadata": sanitize(metadata) if metadata else {},
        },
    }


def evt_llm_end(
    run: "ArzuleRun",
    span_id: Optional[str],
    parent_span_id: Optional[str],
    model_name: str,
    response: Any,
    error: Optional[Exception] = None,
) -> Dict[str, Any]:
    """Create event for LLM call end."""
    status = "error" if error else "ok"
    summary = f"llm completed: {model_name}"
    if error:
        summary = f"llm error: {model_name} - {str(error)[:100]}"

    # Extract token usage if available
    token_usage = {}
    if hasattr(response, "llm_output") and response.llm_output:
        token_usage = response.llm_output.get("token_usage", {})

    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "llm.call.end",
        "status": status,
        "summary": summary,
        "attrs_compact": {
            "model": model_name,
            "token_usage": token_usage,
            "error": str(error) if error else None,
        },
        "payload": {
            "response": sanitize(response) if not error else None,
            "error": str(error) if error else None,
        },
    }


# =============================================================================
# Tool Events (from nested tool calls)
# =============================================================================


def evt_tool_start(
    run: "ArzuleRun",
    span_id: str,
    parent_span_id: Optional[str],
    tool_name: str,
    tool_input: Any,
) -> Dict[str, Any]:
    """Create event for tool call start."""
    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "tool.call.start",
        "status": "ok",
        "summary": f"tool call: {tool_name}",
        "attrs_compact": {
            "tool": tool_name,
        },
        "payload": {
            "input": sanitize(tool_input),
        },
    }


def evt_tool_end(
    run: "ArzuleRun",
    span_id: Optional[str],
    parent_span_id: Optional[str],
    tool_name: str,
    output: Any,
    error: Optional[Exception] = None,
) -> Dict[str, Any]:
    """Create event for tool call end."""
    status = "error" if error else "ok"
    summary = f"tool completed: {tool_name}"
    if error:
        summary = f"tool error: {tool_name} - {str(error)[:100]}"

    return {
        **_base(run, span_id=span_id, parent_span_id=parent_span_id),
        "agent": None,
        "event_type": "tool.call.end",
        "status": status,
        "summary": summary,
        "attrs_compact": {
            "tool": tool_name,
            "error": str(error) if error else None,
        },
        "payload": {
            "output": sanitize(output) if not error else None,
            "error": str(error) if error else None,
        },
    }








