"""DSPy-powered engine component that mirrors the design implementation."""

from __future__ import annotations

import json
import os
from collections import OrderedDict, defaultdict
from collections.abc import Iterable, Mapping, Sequence
from contextlib import nullcontext
from typing import Any, Literal

from pydantic import BaseModel, Field

from flock.artifacts import Artifact
from flock.components import EngineComponent
from flock.dashboard.events import StreamingOutputEvent
from flock.logging.logging import get_logger
from flock.registry import type_registry
from flock.runtime import EvalInputs, EvalResult


logger = get_logger(__name__)


_live_patch_applied = False


# T071: Auto-detect test environment for streaming
def _default_stream_value() -> bool:
    """Return default stream value based on environment.

    Returns False in pytest (clean test output), True otherwise (rich streaming).
    """
    import sys

    return "pytest" not in sys.modules


# Apply the Rich Live patch immediately on module import
def _apply_live_patch_on_import() -> None:
    """Apply Rich Live crop_above patch when module is imported."""
    try:
        _ensure_live_crop_above()
    except Exception:
        pass  # Silently ignore if Rich is not available


def _ensure_live_crop_above() -> None:
    """Monkeypatch rich.live_render to support 'crop_above' overflow."""
    global _live_patch_applied
    if _live_patch_applied:
        return
    try:
        from typing import Literal as _Literal

        from rich import live_render as _lr
    except Exception:
        return

    # Extend the accepted literal at runtime so type checks don't block the new option.
    current_args = getattr(_lr.VerticalOverflowMethod, "__args__", ())
    if "crop_above" not in current_args:
        _lr.VerticalOverflowMethod = _Literal["crop", "crop_above", "ellipsis", "visible"]  # type: ignore[assignment]

    if getattr(_lr.LiveRender.__rich_console__, "_flock_crop_above", False):
        _live_patch_applied = True
        return

    Segment = _lr.Segment
    Text = _lr.Text
    loop_last = _lr.loop_last

    def _patched_rich_console(self, console, options):
        renderable = self.renderable
        style = console.get_style(self.style)
        lines = console.render_lines(renderable, options, style=style, pad=False)
        shape = Segment.get_shape(lines)

        _, height = shape
        max_height = options.size.height
        if height > max_height:
            if self.vertical_overflow == "crop":
                lines = lines[:max_height]
                shape = Segment.get_shape(lines)
            elif self.vertical_overflow == "crop_above":
                lines = lines[-max_height:]
                shape = Segment.get_shape(lines)
            elif self.vertical_overflow == "ellipsis" and max_height > 0:
                lines = lines[: (max_height - 1)]
                overflow_text = Text(
                    "...",
                    overflow="crop",
                    justify="center",
                    end="",
                    style="live.ellipsis",
                )
                lines.append(list(console.render(overflow_text)))
                shape = Segment.get_shape(lines)
        self._shape = shape

        new_line = Segment.line()
        for last, line in loop_last(lines):
            yield from line
            if not last:
                yield new_line

    _patched_rich_console._flock_crop_above = True  # type: ignore[attr-defined]
    _lr.LiveRender.__rich_console__ = _patched_rich_console
    _live_patch_applied = True


class DSPyEngine(EngineComponent):
    """Execute a minimal DSPy program backed by a hosted LLM.

    Behavior intentionally mirrors ``design/dspy_engine.py`` so that orchestration
    relies on the same model resolution, signature preparation, and result
    normalization logic.
    """

    name: str | None = "dspy"
    model: str | None = None
    instructions: str | None = None
    temperature: float = 1.0
    max_tokens: int = 32000
    max_tool_calls: int = 10
    max_retries: int = 0
    stream: bool = Field(
        default_factory=lambda: _default_stream_value(),
        description="Enable streaming output from the underlying DSPy program. Auto-disables in pytest.",
    )
    no_output: bool = Field(
        default=False,
        description="Disable output from the underlying DSPy program.",
    )
    stream_vertical_overflow: Literal["crop", "ellipsis", "crop_above", "visible"] = Field(
        default="crop_above",
        description=(
            "Rich Live vertical overflow strategy; select how tall output is handled; 'crop_above' keeps the most recent rows visible."
        ),
    )
    status_output_field: str = Field(
        default="_status_output",
        description="The field name for the status output.",
    )
    theme: str = Field(
        default="afterglow",
        description="Theme name for Rich output formatting.",
    )
    enable_cache: bool = Field(
        default=False,
        description="Enable caching of DSPy program results",
    )

    async def evaluate(self, agent, ctx, inputs: EvalInputs) -> EvalResult:  # type: ignore[override]
        if not inputs.artifacts:
            return EvalResult(artifacts=[], state=dict(inputs.state))

        model_name = self._resolve_model_name()
        dspy_mod = self._import_dspy()

        lm = dspy_mod.LM(
            model=model_name,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
            cache=self.enable_cache,
            num_retries=self.max_retries,
        )

        primary_artifact = self._select_primary_artifact(inputs.artifacts)
        input_model = self._resolve_input_model(primary_artifact)
        validated_input = self._validate_input_payload(input_model, primary_artifact.payload)
        output_model = self._resolve_output_model(agent)

        # Fetch conversation context from blackboard
        context_history = await self.fetch_conversation_context(ctx)
        has_context = bool(context_history) and self.should_use_context(inputs)

        # Prepare signature with optional context field
        signature = self._prepare_signature_with_context(
            dspy_mod,
            description=self.instructions or agent.description,
            input_schema=input_model,
            output_schema=output_model,
            has_context=has_context,
        )

        sys_desc = self._system_description(self.instructions or agent.description)

        # Pre-generate the artifact ID so it's available from the start
        from uuid import uuid4

        pre_generated_artifact_id = uuid4()

        # Build execution payload with context
        if has_context:
            execution_payload = {
                "input": validated_input,
                "context": context_history,
            }
        else:
            # Backwards compatible - direct input
            execution_payload = validated_input

        # Merge native tools with MCP tools
        native_tools = list(agent.tools or [])

        # Lazy-load MCP tools for this agent
        try:
            mcp_tools = await agent._get_mcp_tools(ctx)
            logger.debug(f"Loaded {len(mcp_tools)} MCP tools for agent {agent.name}")
        except Exception as e:
            # Architecture Decision: AD007 - Graceful Degradation
            # If MCP loading fails, continue with native tools only
            logger.error(f"Failed to load MCP tools in engine: {e}", exc_info=True)
            mcp_tools = []

        # Combine both lists
        # Architecture Decision: AD003 - MCP tools are namespaced, so no conflicts
        combined_tools = native_tools + mcp_tools
        logger.debug(
            f"Total tools for agent {agent.name}: {len(combined_tools)} (native: {len(native_tools)}, mcp: {len(mcp_tools)})"
        )

        with dspy_mod.context(lm=lm):
            program = self._choose_program(dspy_mod, signature, combined_tools)

            # Detect if there's already an active Rich Live context
            should_stream = self.stream
            orchestrator = getattr(ctx, "orchestrator", None)
            if orchestrator:
                is_dashboard = getattr(orchestrator, "is_dashboard", False) if ctx else False
                # if dashboard we always stream, streamin queue only for CLI output
                if should_stream and ctx and not is_dashboard:
                    if not hasattr(orchestrator, "_active_streams"):
                        orchestrator._active_streams = 0

                    if orchestrator._active_streams > 0:
                        should_stream = False
                    else:
                        orchestrator._active_streams += 1

            try:
                if should_stream:
                    (
                        raw_result,
                        _stream_final_display_data,
                    ) = await self._execute_streaming(
                        dspy_mod,
                        program,
                        signature,
                        description=sys_desc,
                        payload=execution_payload,
                        agent=agent,
                        ctx=ctx,
                        pre_generated_artifact_id=pre_generated_artifact_id,
                    )
                    if not self.no_output and ctx:
                        ctx.state["_flock_stream_live_active"] = True
                else:
                    orchestrator = getattr(ctx, "orchestrator", None) if ctx else None

                    raw_result = await self._execute_standard(
                        dspy_mod,
                        program,
                        description=sys_desc,
                        payload=execution_payload,
                    )
                    if ctx and orchestrator and getattr(orchestrator, "_active_streams", 0) > 0:
                        ctx.state["_flock_output_queued"] = True
            finally:
                if should_stream and ctx:
                    if orchestrator is None:
                        orchestrator = getattr(ctx, "orchestrator", None)
                    if orchestrator and hasattr(orchestrator, "_active_streams"):
                        orchestrator._active_streams = max(0, orchestrator._active_streams - 1)

        normalized_output = self._normalize_output_payload(getattr(raw_result, "output", None))
        artifacts, errors = self._materialize_artifacts(
            normalized_output,
            agent.outputs,
            agent.name,
            pre_generated_id=pre_generated_artifact_id,
        )

        state = dict(inputs.state)
        state.setdefault("dspy", {})
        state["dspy"].update({"model": model_name, "raw": normalized_output})

        logs: list[str] = []
        if normalized_output is not None:
            try:
                logs.append(f"dspy.output={json.dumps(normalized_output)}")
            except TypeError:
                logs.append(f"dspy.output={normalized_output!r}")
        logs.extend(f"dspy.error={message}" for message in errors)

        result_artifacts = artifacts if artifacts else list(inputs.artifacts)
        return EvalResult(artifacts=result_artifacts, state=state, logs=logs)

    # ------------------------------------------------------------------
    # Helpers mirroring the design engine

    def _resolve_model_name(self) -> str:
        model = self.model or os.getenv("TRELLIS_MODEL") or os.getenv("OPENAI_MODEL")
        if not model:
            raise NotImplementedError(
                "DSPyEngine requires a configured model (set TRELLIS_MODEL, OPENAI_MODEL, or pass model=...)."
            )
        return model

    def _import_dspy(self):  # pragma: no cover - import guarded by optional dependency
        try:
            import dspy
        except Exception as exc:
            raise NotImplementedError("DSPy is not installed or failed to import.") from exc
        return dspy

    def _select_primary_artifact(self, artifacts: Sequence[Artifact]) -> Artifact:
        return artifacts[-1]

    def _resolve_input_model(self, artifact: Artifact) -> type[BaseModel] | None:
        try:
            return type_registry.resolve(artifact.type)
        except KeyError:
            return None

    def _resolve_output_model(self, agent) -> type[BaseModel] | None:
        if not getattr(agent, "outputs", None):
            return None
        return agent.outputs[0].spec.model

    def _validate_input_payload(
        self,
        schema: type[BaseModel] | None,
        payload: Mapping[str, Any] | None,
    ) -> dict[str, Any]:
        data = dict(payload or {})
        if schema is None:
            return data
        try:
            return schema(**data).model_dump()
        except Exception:
            return data

    def _prepare_signature_with_context(
        self,
        dspy_mod,
        *,
        description: str | None,
        input_schema: type[BaseModel] | None,
        output_schema: type[BaseModel] | None,
        has_context: bool = False,
    ) -> Any:
        """Prepare DSPy signature, optionally including context field."""
        fields = {
            "description": (str, dspy_mod.InputField()),
        }

        # Add context field if we have conversation history
        if has_context:
            fields["context"] = (
                list,
                dspy_mod.InputField(
                    desc="Previous conversation artifacts providing context for this request"
                ),
            )

        fields["input"] = (input_schema or dict, dspy_mod.InputField())
        fields["output"] = (output_schema or dict, dspy_mod.OutputField())

        signature = dspy_mod.Signature(fields)

        instruction = description or "Produce a valid output that matches the 'output' schema."
        if has_context:
            instruction += " Consider the conversation context provided to inform your response."
        instruction += " Return only JSON."

        return signature.with_instructions(instruction)

    def _choose_program(self, dspy_mod, signature, tools: Iterable[Any]):
        tools_list = list(tools or [])
        try:
            if tools_list:
                return dspy_mod.ReAct(signature, tools=tools_list, max_iters=self.max_tool_calls)
            return dspy_mod.Predict(signature)
        except Exception:
            return dspy_mod.Predict(signature)

    def _system_description(self, description: str | None) -> str:
        if description:
            return description
        return "Produce a valid output that matches the 'output' schema. Return only JSON."

    def _normalize_output_payload(self, raw: Any) -> dict[str, Any]:
        if isinstance(raw, BaseModel):
            return raw.model_dump()
        if isinstance(raw, str):
            try:
                return json.loads(raw)
            except json.JSONDecodeError:
                return {"text": raw}
        if isinstance(raw, Mapping):
            return dict(raw)
        return {"value": raw}

    def _materialize_artifacts(
        self,
        payload: dict[str, Any],
        outputs: Iterable[Any],
        produced_by: str,
        pre_generated_id: Any = None,
    ):
        artifacts: list[Artifact] = []
        errors: list[str] = []
        for output in outputs or []:
            model_cls = output.spec.model
            data = self._select_output_payload(payload, model_cls, output.spec.type_name)
            try:
                instance = model_cls(**data)
            except Exception as exc:  # noqa: BLE001 - collect validation errors for logs
                errors.append(str(exc))
                continue

            # Use the pre-generated ID if provided (for streaming), otherwise let Artifact auto-generate
            artifact_kwargs = {
                "type": output.spec.type_name,
                "payload": instance.model_dump(),
                "produced_by": produced_by,
            }
            if pre_generated_id is not None:
                artifact_kwargs["id"] = pre_generated_id

            artifacts.append(Artifact(**artifact_kwargs))
        return artifacts, errors

    def _select_output_payload(
        self,
        payload: Mapping[str, Any],
        model_cls: type[BaseModel],
        type_name: str,
    ) -> dict[str, Any]:
        candidates = [
            payload.get(type_name),
            payload.get(model_cls.__name__),
            payload.get(model_cls.__name__.lower()),
        ]
        for candidate in candidates:
            if isinstance(candidate, Mapping):
                return dict(candidate)
        if isinstance(payload, Mapping):
            return dict(payload)
        return {}

    async def _execute_standard(
        self, dspy_mod, program, *, description: str, payload: dict[str, Any]
    ) -> Any:
        """Execute DSPy program in standard mode (no streaming)."""
        # Handle new format: {"input": ..., "context": ...}
        if isinstance(payload, dict) and "input" in payload:
            return program(
                description=description,
                input=payload["input"],
                context=payload.get("context", []),
            )

        # Handle old format: direct payload (backwards compatible)
        return program(description=description, input=payload, context=[])

    async def _execute_streaming(
        self,
        dspy_mod,
        program,
        signature,
        *,
        description: str,
        payload: dict[str, Any],
        agent: Any,
        ctx: Any = None,
        pre_generated_artifact_id: Any = None,
    ) -> Any:
        """Execute DSPy program in streaming mode with Rich table updates."""
        from rich.console import Console
        from rich.live import Live

        console = Console()

        # Get WebSocketManager for frontend streaming
        ws_manager = None
        if ctx:
            orchestrator = getattr(ctx, "orchestrator", None)
            if orchestrator:
                collector = getattr(orchestrator, "_dashboard_collector", None)
                if collector:
                    ws_manager = getattr(collector, "_websocket_manager", None)

        # Prepare stream listeners for output field
        listeners = []
        try:
            streaming_mod = getattr(dspy_mod, "streaming", None)
            if streaming_mod and hasattr(streaming_mod, "StreamListener"):
                for name, field in signature.output_fields.items():
                    if field.annotation is str:
                        listeners.append(streaming_mod.StreamListener(signature_field_name=name))
        except Exception:
            listeners = []

        streaming_task = dspy_mod.streamify(
            program,
            is_async_program=True,
            stream_listeners=listeners if listeners else None,
        )

        # Handle new format vs old format
        if isinstance(payload, dict) and "input" in payload:
            stream_generator = streaming_task(
                description=description,
                input=payload["input"],
                context=payload.get("context", []),
            )
        else:
            # Old format - backwards compatible
            stream_generator = streaming_task(description=description, input=payload, context=[])

        signature_order = []
        status_field = self.status_output_field
        try:
            signature_order = list(signature.output_fields.keys())
        except Exception:
            signature_order = []

        # Initialize display data in full artifact format (matching OutputUtilityComponent display)
        display_data: OrderedDict[str, Any] = OrderedDict()

        # Use the pre-generated artifact ID that was created before execution started
        display_data["id"] = str(pre_generated_artifact_id)

        # Get the output type from agent configuration
        output_type = "output"
        if hasattr(agent, "outputs") and agent.outputs:
            output_type = agent.outputs[0].spec.type_name

        display_data["type"] = output_type
        display_data["payload"] = OrderedDict()

        # Add output fields to payload section
        for field_name in signature_order:
            if field_name != "description":  # Skip description field
                display_data["payload"][field_name] = ""

        display_data["produced_by"] = agent.name
        display_data["correlation_id"] = (
            str(ctx.correlation_id) if ctx and ctx.correlation_id else None
        )
        display_data["partition_key"] = None
        display_data["tags"] = "set()"
        display_data["visibility"] = OrderedDict([("kind", "Public")])
        display_data["created_at"] = "streaming..."
        display_data["version"] = 1
        display_data["status"] = status_field

        stream_buffers: defaultdict[str, list[str]] = defaultdict(list)
        stream_buffers[status_field] = []
        stream_sequence = 0  # Monotonic sequence for ordering

        formatter = theme_dict = styles = agent_label = None
        live_cm = nullcontext()
        overflow_mode = self.stream_vertical_overflow

        if not self.no_output:
            _ensure_live_crop_above()
            (
                formatter,
                theme_dict,
                styles,
                agent_label,
            ) = self._prepare_stream_formatter(agent)
            initial_panel = formatter.format_result(display_data, agent_label, theme_dict, styles)
            live_cm = Live(
                initial_panel,
                console=console,
                refresh_per_second=4,
                transient=False,
                vertical_overflow=overflow_mode,
            )

        final_result: Any = None

        with live_cm as live:

            def _refresh_panel() -> None:
                if formatter is None or live is None:
                    return
                live.update(formatter.format_result(display_data, agent_label, theme_dict, styles))

            async for value in stream_generator:
                try:
                    from dspy.streaming import StatusMessage, StreamResponse
                    from litellm import ModelResponseStream
                except Exception:
                    StatusMessage = object  # type: ignore
                    StreamResponse = object  # type: ignore
                    ModelResponseStream = object  # type: ignore

                if isinstance(value, StatusMessage):
                    token = getattr(value, "message", "")
                    if token:
                        stream_buffers[status_field].append(str(token) + "\n")
                        display_data["status"] = "".join(stream_buffers[status_field])

                        # Emit to WebSocket
                        if ws_manager and token:
                            try:
                                event = StreamingOutputEvent(
                                    correlation_id=str(ctx.correlation_id)
                                    if ctx and ctx.correlation_id
                                    else "",
                                    agent_name=agent.name,
                                    run_id=ctx.task_id if ctx else "",
                                    output_type="llm_token",
                                    content=str(token + "\n"),
                                    sequence=stream_sequence,
                                    is_final=False,
                                )
                                await ws_manager.broadcast(event)
                                stream_sequence += 1
                            except Exception as e:
                                logger.warning(f"Failed to emit streaming event: {e}")

                        if formatter is not None:
                            _refresh_panel()
                    continue

                if isinstance(value, StreamResponse):
                    token = getattr(value, "chunk", None)
                    signature_field = getattr(value, "signature_field_name", None)
                    if signature_field and signature_field != "description":
                        # Update payload section - accumulate in "output" buffer
                        buffer_key = f"_stream_{signature_field}"
                        if token:
                            stream_buffers[buffer_key].append(str(token))
                            # Show streaming text in payload
                            display_data["payload"]["_streaming"] = "".join(
                                stream_buffers[buffer_key]
                            )

                            # Emit to WebSocket
                            if ws_manager:
                                logger.info(
                                    f"[STREAMING] Emitting StreamResponse token='{token}', sequence={stream_sequence}"
                                )
                                try:
                                    event = StreamingOutputEvent(
                                        correlation_id=str(ctx.correlation_id)
                                        if ctx and ctx.correlation_id
                                        else "",
                                        agent_name=agent.name,
                                        run_id=ctx.task_id if ctx else "",
                                        output_type="llm_token",
                                        content=str(token),
                                        sequence=stream_sequence,
                                        is_final=False,
                                    )
                                    await ws_manager.broadcast(event)
                                    stream_sequence += 1
                                except Exception as e:
                                    logger.warning(f"Failed to emit streaming event: {e}")

                        if formatter is not None:
                            _refresh_panel()
                    continue

                if isinstance(value, ModelResponseStream):
                    chunk = value
                    token = chunk.choices[0].delta.content or ""
                    signature_field = getattr(value, "signature_field_name", None)

                    # Determine output type based on signature field
                    output_type = "llm_token"  # if signature_field and signature_field != "description" else "log"

                    if signature_field and signature_field != "description":
                        # Update payload section - accumulate in buffer
                        buffer_key = f"_stream_{signature_field}"
                        if token:
                            stream_buffers[buffer_key].append(str(token))
                            # Show streaming text in payload
                            display_data["payload"]["_streaming"] = "".join(
                                stream_buffers[buffer_key]
                            )
                    elif token:
                        stream_buffers[status_field].append(str(token))
                        display_data["status"] = "".join(stream_buffers[status_field])

                    # Emit to WebSocket
                    if ws_manager and token:
                        try:
                            event = StreamingOutputEvent(
                                correlation_id=str(ctx.correlation_id)
                                if ctx and ctx.correlation_id
                                else "",
                                agent_name=agent.name,
                                run_id=ctx.task_id if ctx else "",
                                output_type=output_type,
                                content=str(token),
                                sequence=stream_sequence,
                                is_final=False,
                            )
                            await ws_manager.broadcast(event)
                            stream_sequence += 1
                        except Exception as e:
                            logger.warning(f"Failed to emit streaming event: {e}")

                    if formatter is not None:
                        _refresh_panel()
                    continue

                if isinstance(value, dspy_mod.Prediction):
                    final_result = value

                    # Emit final streaming event
                    if ws_manager:
                        try:
                            event = StreamingOutputEvent(
                                correlation_id=str(ctx.correlation_id)
                                if ctx and ctx.correlation_id
                                else "",
                                agent_name=agent.name,
                                run_id=ctx.task_id if ctx else "",
                                output_type="log",
                                content="\nAmount of output tokens: " + str(stream_sequence),
                                sequence=stream_sequence,
                                is_final=True,  # Mark as final
                            )
                            await ws_manager.broadcast(event)
                            event = StreamingOutputEvent(
                                correlation_id=str(ctx.correlation_id)
                                if ctx and ctx.correlation_id
                                else "",
                                agent_name=agent.name,
                                run_id=ctx.task_id if ctx else "",
                                output_type="log",
                                content="--- End of output ---",
                                sequence=stream_sequence,
                                is_final=True,  # Mark as final
                            )
                            await ws_manager.broadcast(event)
                        except Exception as e:
                            logger.warning(f"Failed to emit final streaming event: {e}")

                    if formatter is not None:
                        # Update payload section with final values
                        payload_data = OrderedDict()
                        for field_name in signature_order:
                            if field_name != "description" and hasattr(final_result, field_name):
                                field_value = getattr(final_result, field_name)
                                # If the field is a BaseModel, unwrap it to dict
                                if isinstance(field_value, BaseModel):
                                    payload_data.update(field_value.model_dump())
                                else:
                                    payload_data[field_name] = field_value

                        # Update all fields with actual values
                        display_data["payload"].clear()
                        display_data["payload"].update(payload_data)

                        # Update timestamp
                        from datetime import datetime, timezone

                        display_data["created_at"] = datetime.now(timezone.utc).isoformat()

                        # Remove status field from display
                        display_data.pop("status", None)
                        _refresh_panel()

        if final_result is None:
            raise RuntimeError("Streaming did not yield a final prediction.")

        # Return both the result and the display data for final ID update
        return final_result, (formatter, display_data, theme_dict, styles, agent_label)

    def _prepare_stream_formatter(
        self, agent: Any
    ) -> tuple[Any, dict[str, Any], dict[str, Any], str]:
        """Build formatter + theme metadata for streaming tables."""
        import pathlib

        from flock.logging.formatters.themed_formatter import (
            ThemedAgentResultFormatter,
            create_pygments_syntax_theme,
            get_default_styles,
            load_syntax_theme_from_file,
            load_theme_from_file,
        )

        themes_dir = pathlib.Path(__file__).resolve().parents[1] / "themes"
        theme_filename = self.theme
        if not theme_filename.endswith(".toml"):
            theme_filename = f"{theme_filename}.toml"
        theme_path = themes_dir / theme_filename

        try:
            theme_dict = load_theme_from_file(theme_path)
        except Exception:
            fallback_path = themes_dir / "afterglow.toml"
            theme_dict = load_theme_from_file(fallback_path)
            theme_path = fallback_path

        from flock.logging.formatters.themes import OutputTheme

        formatter = ThemedAgentResultFormatter(theme=OutputTheme.afterglow)
        styles = get_default_styles(theme_dict)
        formatter.styles = styles

        try:
            syntax_theme = load_syntax_theme_from_file(theme_path)
            formatter.syntax_style = create_pygments_syntax_theme(syntax_theme)
        except Exception:
            formatter.syntax_style = None

        model_label = self.model or ""
        agent_label = agent.name if not model_label else f"{agent.name} - {model_label}"

        return formatter, theme_dict, styles, agent_label

    def _print_final_stream_display(
        self,
        stream_display_data: tuple[Any, OrderedDict, dict, dict, str],
        artifact_id: str,
        artifact: Artifact,
    ) -> None:
        """Print the final streaming display with the real artifact ID."""
        from rich.console import Console

        formatter, display_data, theme_dict, styles, agent_label = stream_display_data

        # Update display_data with the real artifact information
        display_data["id"] = artifact_id
        display_data["created_at"] = artifact.created_at.isoformat()

        # Update all artifact metadata
        display_data["correlation_id"] = (
            str(artifact.correlation_id) if artifact.correlation_id else None
        )
        display_data["partition_key"] = artifact.partition_key
        display_data["tags"] = "set()" if not artifact.tags else f"set({list(artifact.tags)})"

        # Print the final panel
        console = Console()
        final_panel = formatter.format_result(display_data, agent_label, theme_dict, styles)
        console.print(final_panel)


__all__ = ["DSPyEngine"]


# Apply the Rich Live patch when this module is imported
_apply_live_patch_on_import()
