"""Utilities for working with SimpleLLMFunc React event streams."""

from __future__ import annotations

from dataclasses import is_dataclass
from datetime import datetime
from enum import Enum
import json
from typing import Any, Dict

from SimpleLLMFunc.base.post_process import extract_content_from_stream_response
from SimpleLLMFunc.hooks.events import CustomEvent
from SimpleLLMFunc.hooks.stream import (
    EventYield,
    ReactOutput,
    ResponseYield,
    is_event_yield,
    is_response_yield,
)
from SimpleLLMFunc.utils.tui.formatters import extract_reasoning_delta


def normalize_for_json(value: Any) -> Any:
    """Recursively normalize complex event payloads into JSON-safe values."""
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, datetime):
        return value.isoformat()
    if isinstance(value, Enum):
        return value.value
    if isinstance(value, BaseException):
        return {
            "type": value.__class__.__name__,
            "message": str(value),
        }
    if isinstance(value, dict):
        return {str(key): normalize_for_json(item) for key, item in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [normalize_for_json(item) for item in value]
    if is_dataclass(value):
        dataclass_fields = getattr(value, "__dataclass_fields__", {})
        return normalize_for_json(
            {name: getattr(value, name) for name in dataclass_fields}
        )
    if hasattr(value, "model_dump"):
        try:
            return normalize_for_json(value.model_dump())
        except Exception:
            pass
    if hasattr(value, "dict"):
        try:
            return normalize_for_json(value.dict())
        except Exception:
            pass
    if hasattr(value, "as_dict"):
        try:
            return normalize_for_json(value.as_dict())
        except Exception:
            pass
    if hasattr(value, "__dict__") and not isinstance(value, type):
        try:
            return normalize_for_json(vars(value))
        except Exception:
            pass
    return str(value)


def extract_response_text(response: Any, consumer: str = "agent_stream") -> str:
    """Extract text from a ResponseYield payload or raw provider chunk."""
    if isinstance(response, str):
        return response

    if hasattr(response, "model_dump"):
        try:
            dumped = response.model_dump()
            text = extract_response_text(dumped, consumer)
            if text:
                return text
        except Exception:
            pass

    try:
        return extract_content_from_stream_response(response, consumer) or ""
    except Exception:
        pass

    try:
        choices = None
        if isinstance(response, dict):
            choices = response.get("choices")
        elif hasattr(response, "choices"):
            choices = response.choices

        if isinstance(choices, list) and choices:
            first = choices[0]
            if isinstance(first, dict):
                delta = first.get("delta")
                if isinstance(delta, dict):
                    content = delta.get("content")
                    return content if isinstance(content, str) else ""
                message = first.get("message")
                if isinstance(message, dict):
                    content = message.get("content")
                    return content if isinstance(content, str) else ""
            else:
                delta = getattr(first, "delta", None)
                content = getattr(delta, "content", None) if delta is not None else None
                if isinstance(content, str):
                    return content
                message = getattr(first, "message", None)
                content = (
                    getattr(message, "content", None) if message is not None else None
                )
                if isinstance(content, str):
                    return content
        return ""
    except Exception:
        return ""


def extract_output_text(output: ReactOutput, consumer: str = "agent_stream") -> str:
    """Extract plain text from a ReactOutput if available."""
    if is_response_yield(output):
        return extract_response_text(output.response, consumer)
    return ""


def extract_response_reasoning(response: Any) -> str:
    """Extract provider reasoning text from a response chunk or full response."""
    if isinstance(response, str):
        return ""

    if hasattr(response, "model_dump"):
        try:
            dumped = response.model_dump()
            reasoning = extract_response_reasoning(dumped)
            if reasoning:
                return reasoning
        except Exception:
            pass

    try:
        if isinstance(response, dict):
            choices = response.get("choices")
            if isinstance(choices, list) and choices:
                first = choices[0]
                if isinstance(first, dict):
                    delta = first.get("delta")
                    if isinstance(delta, dict):
                        for key in ("reasoning", "reasoning_content", "reasoning_text"):
                            value = delta.get(key)
                            if isinstance(value, str) and value:
                                return value
                        details = delta.get("reasoning_details")
                        if isinstance(details, list):
                            collected: list[str] = []
                            for detail in details:
                                if not isinstance(detail, dict):
                                    continue
                                detail_type = str(detail.get("type", ""))
                                if detail_type.endswith("encrypted"):
                                    continue
                                text = detail.get("text") or detail.get("data")
                                if isinstance(text, str) and text:
                                    collected.append(text)
                            if collected:
                                return "".join(collected)

        reasoning = extract_reasoning_delta(response)
        if reasoning:
            return reasoning
    except Exception:
        pass

    try:
        choices = (
            response.get("choices") if isinstance(response, dict) else response.choices
        )
        if not isinstance(choices, list) or not choices:
            return ""
        first = choices[0]
        message = (
            first.get("message")
            if isinstance(first, dict)
            else getattr(first, "message", None)
        )
        if message is None:
            return ""

        for attr in ("reasoning", "reasoning_content", "reasoning_text"):
            value = (
                message.get(attr)
                if isinstance(message, dict)
                else getattr(message, attr, None)
            )
            if isinstance(value, str) and value:
                return value

        details = (
            message.get("reasoning_details")
            if isinstance(message, dict)
            else getattr(message, "reasoning_details", None)
        )
        if isinstance(details, list):
            collected: list[str] = []
            for detail in details:
                if isinstance(detail, dict):
                    detail_type = str(detail.get("type", ""))
                    if detail_type.endswith("encrypted"):
                        continue
                    text = detail.get("text") or detail.get("data")
                else:
                    detail_type = str(getattr(detail, "type", ""))
                    if detail_type.endswith("encrypted"):
                        continue
                    text = getattr(detail, "text", None) or getattr(
                        detail, "data", None
                    )
                if isinstance(text, str) and text:
                    collected.append(text)
            return "".join(collected)
    except Exception:
        return ""

    return ""


def serialize_react_output(
    output: ReactOutput,
    delta_consumer: str = "agent_stream",
) -> Dict[str, Any]:
    """Serialize ResponseYield/EventYield into a stable JSON payload."""
    if is_response_yield(output):
        return {
            "type": "response",
            "delta_text": extract_response_text(output.response, delta_consumer),
            "delta_reasoning": extract_response_reasoning(output.response),
            "response": normalize_for_json(output.response),
            "messages": normalize_for_json(output.messages),
        }

    assert is_event_yield(output)
    event_type = getattr(output.event.event_type, "value", output.event.event_type)
    if isinstance(output.event, CustomEvent) and output.event.event_name:
        event_type = output.event.event_name
    return {
        "type": "event",
        "event_type": str(event_type),
        "event": normalize_for_json(output.event),
        "origin": normalize_for_json(output.origin),
    }


def event_name_for_output(output: ReactOutput) -> str:
    if is_response_yield(output):
        return "response"
    assert is_event_yield(output)
    if isinstance(output.event, CustomEvent) and output.event.event_name:
        return output.event.event_name
    event_type = getattr(output.event.event_type, "value", output.event.event_type)
    return str(event_type)


def format_sse(event: str, data: Dict[str, Any]) -> str:
    payload = json.dumps(data, ensure_ascii=False)
    return f"event: {event}\ndata: {payload}\n\n"


def project_response_to_oai_chunk(
    output: ResponseYield,
    request_id: str,
    model: str,
    created_time: int,
    sent_role: bool,
) -> Dict[str, Any] | None:
    """Project a response yield into an OpenAI-compatible streaming chunk."""
    raw_response = normalize_for_json(output.response)
    if isinstance(raw_response, dict):
        if raw_response.get("object") == "chat.completion.chunk":
            return raw_response
        if raw_response.get("object") == "chat.completion":
            choices = raw_response.get("choices")
            if isinstance(choices, list) and choices:
                first = choices[0] if isinstance(choices[0], dict) else {}
                message = first.get("message") if isinstance(first, dict) else {}
                if isinstance(message, dict):
                    delta: Dict[str, Any] = {
                        "role": message.get("role", "assistant"),
                        "content": message.get("content"),
                    }
                    tool_calls = message.get("tool_calls")
                    if tool_calls is not None:
                        delta["tool_calls"] = tool_calls
                    return {
                        "id": request_id,
                        "object": "chat.completion.chunk",
                        "created": created_time,
                        "model": model,
                        "choices": [
                            {
                                "index": 0,
                                "delta": delta,
                                "finish_reason": first.get("finish_reason"),
                            }
                        ],
                    }

    delta_text = extract_response_text(output.response, "agent_stream")
    tool_calls = None
    try:
        if isinstance(raw_response, dict):
            choices = raw_response.get("choices")
            if isinstance(choices, list) and choices:
                first_item = choices[0]
                delta_value: Any = {}
                if isinstance(first_item, dict):
                    delta_value = first_item.get("delta")
                if isinstance(delta_value, dict):
                    tool_calls = delta_value.get("tool_calls")
    except Exception:
        tool_calls = None

    if not delta_text and not tool_calls:
        return None

    delta: Dict[str, Any] = {}
    if not sent_role:
        delta["role"] = "assistant"
    if delta_text:
        delta["content"] = delta_text
    if tool_calls is not None:
        delta["tool_calls"] = tool_calls

    return {
        "id": request_id,
        "object": "chat.completion.chunk",
        "created": created_time,
        "model": model,
        "choices": [
            {
                "index": 0,
                "delta": delta,
                "finish_reason": None,
            }
        ],
    }


__all__ = [
    "EventYield",
    "ResponseYield",
    "event_name_for_output",
    "extract_output_text",
    "extract_response_reasoning",
    "extract_response_text",
    "format_sse",
    "is_event_yield",
    "is_response_yield",
    "normalize_for_json",
    "project_response_to_oai_chunk",
    "serialize_react_output",
]