"""
Chain Orchestrator — Chains multiple LLM calls with automatic context passing.

Executes a sequence of LLM steps, feeding each step's output as context into
subsequent steps via variable substitution.  Supports per-step API keys (BYOK),
graceful partial-failure handling, and detailed timing/token tracking.

Usage:
    from brynq.chain_orchestrator import ChainOrchestrator, ChainStep

    chain = ChainOrchestrator()
    result = await chain.execute([
        ChainStep(
            backend="claude",
            model="claude-sonnet-4-20250514",
            prompt="Analyze this data: {original_input}",
            api_key="user_claude_key",
        ),
        ChainStep(
            backend="gemini",
            model="gemini-2.5-pro",
            prompt="Review this analysis:\\n{step_1_output}",
            api_key="user_gemini_key",
        ),
        ChainStep(
            backend="ollama",
            model="llama3",
            prompt="Summarize both:\\n{step_1_output}\\n{step_2_output}",
        ),
    ], original_input="Q4 revenue was $2.3M, up 15% YoY...")

All stdlib + brynq.llm_router — no external packages.
"""

from __future__ import annotations

import asyncio
import logging
import os
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from . import llm_router
from .context_optimizer import estimate_tokens

log = logging.getLogger("brynq.chain_orchestrator")

# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

# Maps backend names to the env var that holds their API key.
_BACKEND_KEY_ENV: Dict[str, str] = {
    "claude": "ANTHROPIC_API_KEY",
    "gemini": "GEMINI_API_KEY",
}


@dataclass
class ChainStep:
    """Definition of a single step in a chain.

    Attributes:
        backend:  LLM backend name (``"ollama"``, ``"claude"``, ``"gemini"``,
                  ``"lmstudio"``, or ``"auto"``).
        model:    Model name (e.g. ``"llama3"``, ``"claude-sonnet-4-20250514"``).
        prompt:   Prompt template.  May contain ``{original_input}``,
                  ``{step_N_output}`` (1-indexed), or any key passed via
                  *extra_vars*.
        api_key:  Optional per-step API key (BYOK).  Temporarily injected into
                  the environment for the duration of the LLM call.
        system_prompt: Optional system-level instruction prepended to messages.
        temperature: Optional temperature override for this step.
        max_tokens: Optional max_tokens override for this step.
        name:     Human-readable label for logs/results.  Defaults to
                  ``"step_N"`` at execution time.
    """

    backend: str = "auto"
    model: str = "auto"
    prompt: str = ""
    api_key: Optional[str] = None
    system_prompt: Optional[str] = None
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    name: Optional[str] = None


@dataclass
class StepResult:
    """Result of executing a single chain step.

    Attributes:
        step_index: 0-based index within the chain.
        step_name:  Name of the step (``"step_1"``, ``"step_2"``, etc.).
        backend:    Backend that was used.
        model:      Model that was used.
        output:     The LLM's text response (empty string on failure).
        prompt_rendered: The prompt after variable substitution.
        elapsed_seconds: Wall-clock time for this step.
        token_estimate:  Rough token count of the output (via ``estimate_tokens``).
        error:      Error message if the step failed, else ``None``.
        success:    Whether the step completed without error.
    """

    step_index: int = 0
    step_name: str = ""
    backend: str = ""
    model: str = ""
    output: str = ""
    prompt_rendered: str = ""
    elapsed_seconds: float = 0.0
    token_estimate: int = 0
    error: Optional[str] = None
    success: bool = True


@dataclass
class ChainResult:
    """Aggregate result of executing an entire chain.

    Attributes:
        steps:            List of per-step results.
        final_output:     The output of the last *successful* step, or ``""``
                          if every step failed.
        total_elapsed:    Wall-clock time for the full chain.
        total_tokens:     Sum of estimated tokens across all steps.
        success:          ``True`` if every step succeeded.
        error:            Summary error message if any step failed.
        original_input:   The original input passed to ``execute()``.
    """

    steps: List[StepResult] = field(default_factory=list)
    final_output: str = ""
    total_elapsed: float = 0.0
    total_tokens: int = 0
    success: bool = True
    error: Optional[str] = None
    original_input: str = ""


# ---------------------------------------------------------------------------
# Variable substitution
# ---------------------------------------------------------------------------


def _substitute_variables(
    template: str,
    original_input: str,
    step_outputs: Dict[str, str],
    extra_vars: Optional[Dict[str, str]] = None,
) -> str:
    """Replace ``{original_input}``, ``{step_N_output}``, and extra vars in *template*.

    Unknown placeholders are left as-is so downstream prompts remain readable
    rather than raising KeyError.
    """
    result = template.replace("{original_input}", original_input)

    for key, value in step_outputs.items():
        result = result.replace("{" + key + "}", value)

    if extra_vars:
        for key, value in extra_vars.items():
            result = result.replace("{" + key + "}", value)

    return result


# ---------------------------------------------------------------------------
# Chain Orchestrator
# ---------------------------------------------------------------------------


class ChainOrchestrator:
    """Execute a sequence of LLM calls, threading outputs between steps.

    The orchestrator is stateless — all state lives inside the ``ChainResult``
    returned by ``execute()``.  This makes it safe to reuse a single instance
    for many chains, even concurrently.
    """

    def __init__(self, default_max_tokens: int = 2048) -> None:
        self.default_max_tokens = default_max_tokens

    # ------------------------------------------------------------------
    # Internal: run a single step (sync, called from thread)
    # ------------------------------------------------------------------

    def _run_step(
        self,
        step: ChainStep,
        step_index: int,
        original_input: str,
        step_outputs: Dict[str, str],
        extra_vars: Optional[Dict[str, str]] = None,
    ) -> StepResult:
        """Execute one step synchronously and return its result."""
        step_name = step.name or f"step_{step_index + 1}"
        rendered = _substitute_variables(
            step.prompt, original_input, step_outputs, extra_vars
        )

        result = StepResult(
            step_index=step_index,
            step_name=step_name,
            backend=step.backend,
            model=step.model,
            prompt_rendered=rendered,
        )

        t0 = time.monotonic()

        # Build kwargs for llm_router.chat()
        kwargs: Dict[str, Any] = {}
        if step.temperature is not None:
            kwargs["temperature"] = step.temperature
        if step.max_tokens is not None:
            kwargs["max_tokens"] = step.max_tokens
        elif self.default_max_tokens:
            kwargs["max_tokens"] = self.default_max_tokens

        # Build messages
        messages: List[Dict[str, str]] = []
        if step.system_prompt:
            messages.append({"role": "system", "content": step.system_prompt})
        messages.append({"role": "user", "content": rendered})

        # BYOK: temporarily inject API key into env if provided
        env_var = _BACKEND_KEY_ENV.get(step.backend)
        old_key: Optional[str] = None
        key_was_set = False

        try:
            if step.api_key and env_var:
                old_key = os.environ.get(env_var)
                os.environ[env_var] = step.api_key
                key_was_set = True

            output = llm_router.chat(
                backend=step.backend,
                model=step.model,
                messages=messages,
                **kwargs,
            )

            result.output = output
            result.token_estimate = estimate_tokens(output)
            result.success = True

        except Exception as exc:
            log.warning(
                "Chain step %s (%s/%s) failed: %s",
                step_name,
                step.backend,
                step.model,
                exc,
            )
            result.error = str(exc)
            result.success = False

        finally:
            # Restore the original env var
            if key_was_set and env_var:
                if old_key is not None:
                    os.environ[env_var] = old_key
                else:
                    os.environ.pop(env_var, None)

            result.elapsed_seconds = time.monotonic() - t0

        return result

    # ------------------------------------------------------------------
    # Public: async execute
    # ------------------------------------------------------------------

    async def execute(
        self,
        steps: List[ChainStep],
        original_input: str = "",
        extra_vars: Optional[Dict[str, str]] = None,
    ) -> ChainResult:
        """Execute a chain of LLM steps sequentially.

        Each step's output is available to subsequent steps via
        ``{step_N_output}`` (1-indexed) in their prompt templates.

        Args:
            steps:          Ordered list of chain steps.
            original_input: The initial user input, available as
                            ``{original_input}`` in every prompt template.
            extra_vars:     Additional variables available for substitution
                            in all prompt templates (e.g. ``{"domain": "finance"}``).

        Returns:
            A ``ChainResult`` with per-step details and the final output.
        """
        if not steps:
            return ChainResult(
                success=True,
                original_input=original_input,
            )

        chain_result = ChainResult(original_input=original_input)
        step_outputs: Dict[str, str] = {}
        t_start = time.monotonic()

        loop = asyncio.get_running_loop()

        for idx, step in enumerate(steps):
            # Run the (blocking) LLM call in a thread to avoid blocking the
            # event loop.
            step_result = await loop.run_in_executor(
                None,
                self._run_step,
                step,
                idx,
                original_input,
                step_outputs,
                extra_vars,
            )

            chain_result.steps.append(step_result)
            chain_result.total_tokens += step_result.token_estimate

            if step_result.success:
                # Register output for downstream steps (1-indexed)
                step_outputs[f"step_{idx + 1}_output"] = step_result.output
                chain_result.final_output = step_result.output
            else:
                # Record partial failure but continue — downstream steps will
                # get an empty string for this step's output.
                chain_result.success = False
                chain_result.error = (
                    f"Step {idx + 1} ({step_result.step_name}) failed: "
                    f"{step_result.error}"
                )
                step_outputs[f"step_{idx + 1}_output"] = ""
                # Stop the chain on failure — downstream steps depend on this output
                break

        chain_result.total_elapsed = time.monotonic() - t_start
        return chain_result

    # ------------------------------------------------------------------
    # Convenience: synchronous execute
    # ------------------------------------------------------------------

    def execute_sync(
        self,
        steps: List[ChainStep],
        original_input: str = "",
        extra_vars: Optional[Dict[str, str]] = None,
    ) -> ChainResult:
        """Synchronous wrapper around ``execute()`` for non-async callers.

        Creates a new event loop if none is running.
        """
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None

        if loop and loop.is_running():
            # Already inside an async context — use nest_asyncio pattern
            # or create a new thread.  For simplicity, run in a new thread.
            import concurrent.futures

            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                future = pool.submit(
                    asyncio.run,
                    self.execute(steps, original_input, extra_vars),
                )
                return future.result()
        else:
            return asyncio.run(self.execute(steps, original_input, extra_vars))
