"""
Plan Executor — The dumb-but-reliable execution engine.

Receives signed execution plans from Brynq Cloud and executes them step by
step. Routes each step to the correct bridge (local/cloud/agent), captures
output and timing, and stores results for downstream steps.

Design principles:
  - Deliberately simple: a for-loop over steps.
  - Rock solid: every error is caught and surfaced, never swallowed.
  - Privacy-first: NEVER sends user content to Brynq servers — only metadata
    (timing, token counts, success/fail).
  - Tamper-evident: validates plan signature before executing.
  - TTL-enforced: rejects expired plans.

All stdlib — no external packages.
"""

from __future__ import annotations

import hashlib
import hmac
import json
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from .agent_bridge import AgentBridge
from .cloud_bridge import CloudBridge
from .context import resolve_context
from .local_bridge import BridgeResponse, LocalBridge, _estimate_tokens

log = logging.getLogger("brynq.runtime.plan_executor")

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


@dataclass
class ExecutionStep:
    """A single step in an execution plan.

    Attributes:
        step_id: Unique identifier for this step (e.g. "step_1", "step_2").
        type: The bridge type — "local", "cloud", "local_agent", "remote_agent".
        backend: Backend/provider name (e.g. "ollama", "claude", "gemini", "gpt").
        model: Model or agent identifier.
        prompt: Prompt template with {context_ref} placeholders.
        context_refs: List of context references this step depends on
            (e.g. ["original_input", "step_1"]).
        config: Optional per-step config (temperature, max_tokens, etc.).
        endpoint: For agent steps — the agent's HTTP endpoint.
        auth_token: For remote agent steps — marketplace auth token.
        api_key: For cloud steps — the user's API key (from local vault).
    """

    step_id: str = ""
    type: str = "local"  # "local" | "cloud" | "local_agent" | "remote_agent"
    backend: str = ""
    model: str = ""
    prompt: str = ""
    context_refs: List[str] = field(default_factory=list)
    config: Dict[str, Any] = field(default_factory=dict)
    endpoint: str = ""
    auth_token: str = ""
    api_key: str = ""


@dataclass
class ExecutionPlan:
    """A signed execution plan from Brynq Cloud.

    Attributes:
        plan_id: Unique plan identifier.
        steps: Ordered list of execution steps.
        original_input: The user's original input text.
        signature: HMAC signature for tamper detection.
        created_at: Unix timestamp when the plan was created.
        ttl_seconds: How long the plan is valid (default 300s = 5 min).
        metadata: Extra plan metadata from the cloud.
    """

    plan_id: str = ""
    steps: List[ExecutionStep] = field(default_factory=list)
    original_input: str = ""
    signature: str = ""
    created_at: float = 0.0
    ttl_seconds: int = 300  # 5 minutes default
    metadata: Dict[str, Any] = field(default_factory=dict)


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

    Attributes:
        step_id: Matches the step's step_id.
        output: The text output from the bridge.
        success: Whether the step completed without error.
        error: Error message if the step failed, else None.
        elapsed_seconds: Wall-clock time for this step.
        token_estimate: Rough token count of the output.
        backend: The backend that was actually used.
        model: The model that was actually used.
    """

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


@dataclass
class ExecutionResult:
    """Aggregate result of executing an entire plan.

    Attributes:
        plan_id: The plan that was executed.
        steps: Per-step results.
        final_output: Output from the last successful step.
        total_elapsed: Wall-clock time for the full plan.
        total_tokens: Sum of token estimates across all steps.
        success: True if every step succeeded.
        error: Summary error message if any step failed.
        metadata_for_cloud: Metadata safe to report to Brynq Cloud
            (timing, tokens, success — NO content).
    """

    plan_id: str = ""
    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
    metadata_for_cloud: Dict[str, Any] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Plan validation
# ---------------------------------------------------------------------------


def validate_signature(plan: ExecutionPlan, signing_key: str) -> bool:
    """Validate that a plan's HMAC signature is authentic.

    The signature covers the plan_id, step count, step types/models,
    original_input hash, and created_at timestamp. This ensures the plan
    was produced by Brynq Cloud and has not been tampered with.

    Args:
        plan: The execution plan to validate.
        signing_key: The shared secret used to sign plans. Derived from
            the user's session with Brynq Cloud.

    Returns:
        True if the signature matches, False if the plan is tampered.
    """
    if not plan.signature:
        return False
    if not signing_key:
        return False

    payload = _build_signature_payload(plan)
    expected = hmac.new(
        signing_key.encode("utf-8"),
        payload.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()

    return hmac.compare_digest(plan.signature, expected)


def sign_plan(plan: ExecutionPlan, signing_key: str) -> str:
    """Generate an HMAC signature for a plan.

    Used by Brynq Cloud (or tests) to sign plans before sending them
    to the runtime.

    Args:
        plan: The execution plan to sign.
        signing_key: The shared secret.

    Returns:
        Hex-encoded HMAC-SHA256 signature.
    """
    payload = _build_signature_payload(plan)
    return hmac.new(
        signing_key.encode("utf-8"),
        payload.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()


def _build_signature_payload(plan: ExecutionPlan) -> str:
    """Build the canonical string that gets signed.

    Covers all structural fields — any modification to the plan
    (adding/removing/reordering steps, changing models, etc.)
    will invalidate the signature.
    """
    # Hash the original input rather than including it directly,
    # so the signature payload is always a fixed, predictable size.
    input_hash = hashlib.sha256(plan.original_input.encode("utf-8")).hexdigest()

    step_summaries = []
    for step in plan.steps:
        step_summaries.append(f"{step.step_id}:{step.type}:{step.backend}:{step.model}")

    return "|".join([
        plan.plan_id,
        str(len(plan.steps)),
        ",".join(step_summaries),
        input_hash,
        str(plan.created_at),
    ])


def validate_ttl(plan: ExecutionPlan) -> bool:
    """Check that a plan has not expired.

    Args:
        plan: The execution plan to check.

    Returns:
        True if the plan is still within its TTL, False if expired.
    """
    if plan.ttl_seconds <= 0:
        return True  # TTL of 0 means no expiry
    now = time.time()
    return (now - plan.created_at) < plan.ttl_seconds


# ---------------------------------------------------------------------------
# Plan Executor
# ---------------------------------------------------------------------------


class PlanExecutor:
    """Execute signed plans from Brynq Cloud, step by step.

    The executor is stateless — all state lives in the ExecutionResult.
    Safe to reuse a single instance for many plans, even concurrently.
    """

    def __init__(
        self,
        signing_key: str = "",
        local_bridge: Optional[LocalBridge] = None,
        cloud_bridge: Optional[CloudBridge] = None,
        agent_bridge: Optional[AgentBridge] = None,
        skip_signature_check: bool = False,
        skip_ttl_check: bool = False,
    ) -> None:
        """Initialize the plan executor.

        Args:
            signing_key: Shared secret for signature validation. Required
                unless skip_signature_check is True.
            local_bridge: Bridge for local model calls. Auto-created if None.
            cloud_bridge: Bridge for cloud API calls. Auto-created if None.
            agent_bridge: Bridge for agent calls. Auto-created if None.
            skip_signature_check: If True, skip plan signature validation.
                Only for testing — NEVER in production.
            skip_ttl_check: If True, skip plan TTL validation.
                Only for testing — NEVER in production.
        """
        self.signing_key = signing_key
        self.local_bridge = local_bridge or LocalBridge()
        self.cloud_bridge = cloud_bridge or CloudBridge()
        self.agent_bridge = agent_bridge or AgentBridge()
        self.skip_signature_check = skip_signature_check
        self.skip_ttl_check = skip_ttl_check

    async def execute(self, plan: ExecutionPlan) -> ExecutionResult:
        """Execute a plan step by step. Dumb but reliable.

        For each step:
        1. Resolve context refs (get previous step outputs)
        2. Route to correct bridge (local/cloud/agent)
        3. Execute and capture output
        4. Store result for downstream steps
        5. Report metadata to cloud (no content!)

        On error, execution stops and returns partial results.

        Args:
            plan: A signed execution plan from Brynq Cloud.

        Returns:
            ExecutionResult with per-step details, final output, and
            metadata safe for cloud reporting.
        """
        result = ExecutionResult(plan_id=plan.plan_id)
        t_start = time.monotonic()

        # --- Validate signature ---
        if not self.skip_signature_check:
            if not validate_signature(plan, self.signing_key):
                result.success = False
                result.error = "Plan signature validation failed — plan may be tampered"
                result.total_elapsed = time.monotonic() - t_start
                log.error("Signature validation failed for plan %s", plan.plan_id)
                return result

        # --- Validate TTL ---
        if not self.skip_ttl_check:
            if not validate_ttl(plan):
                result.success = False
                result.error = "Plan has expired (TTL exceeded)"
                result.total_elapsed = time.monotonic() - t_start
                log.error("TTL expired for plan %s", plan.plan_id)
                return result

        # --- Validate plan has steps ---
        if not plan.steps:
            result.success = True
            result.total_elapsed = time.monotonic() - t_start
            return result

        # --- Execute step by step ---
        previous_results: Dict[str, StepResult] = {}

        for step in plan.steps:
            log.info(
                "Executing %s: type=%s backend=%s model=%s",
                step.step_id, step.type, step.backend, step.model,
            )

            # 1. Resolve context references
            resolved_prompt = resolve_context(step, plan, previous_results)

            # 2. Route to correct bridge and execute
            bridge_response = self._route_step(step, resolved_prompt)

            # 3. Capture result
            step_result = StepResult(
                step_id=step.step_id,
                output=bridge_response.output,
                success=bridge_response.success,
                error=bridge_response.error,
                elapsed_seconds=bridge_response.elapsed_seconds,
                token_estimate=bridge_response.token_estimate,
                backend=bridge_response.backend or step.backend,
                model=bridge_response.model or step.model,
            )

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

            # 4. Store result for downstream steps
            previous_results[step.step_id] = step_result

            if step_result.success:
                result.final_output = step_result.output
                log.info(
                    "Step %s completed: %d tokens in %.2fs",
                    step.step_id, step_result.token_estimate, step_result.elapsed_seconds,
                )
            else:
                # 5. Stop on error — downstream steps depend on this output
                result.success = False
                result.error = (
                    f"Step {step.step_id} failed: {step_result.error}"
                )
                log.error(
                    "Step %s failed: %s — aborting plan %s",
                    step.step_id, step_result.error, plan.plan_id,
                )
                break

        result.total_elapsed = time.monotonic() - t_start

        # Build metadata safe for cloud reporting (NO content)
        result.metadata_for_cloud = self._build_cloud_metadata(result)

        log.info(
            "Plan %s %s: %d steps, %d tokens, %.2fs",
            plan.plan_id,
            "completed" if result.success else "failed",
            len(result.steps),
            result.total_tokens,
            result.total_elapsed,
        )

        return result

    def _route_step(self, step: ExecutionStep, resolved_prompt: str) -> BridgeResponse:
        """Route a step to the correct bridge based on step.type.

        Args:
            step: The execution step.
            resolved_prompt: The prompt with all context refs resolved.

        Returns:
            BridgeResponse from the selected bridge.
        """
        step_type = step.type.lower()

        if step_type == "local":
            return self.local_bridge.call(
                model=step.model,
                prompt=resolved_prompt,
                config=step.config,
            )

        elif step_type == "cloud":
            return self.cloud_bridge.call(
                backend=step.backend,
                model=step.model,
                prompt=resolved_prompt,
                api_key=step.api_key,
                config=step.config,
            )

        elif step_type == "local_agent":
            return self.agent_bridge.call_local(
                agent_id=step.model,
                endpoint=step.endpoint,
                prompt=resolved_prompt,
                config=step.config,
            )

        elif step_type == "remote_agent":
            return self.agent_bridge.call_remote(
                agent_id=step.model,
                endpoint=step.endpoint,
                auth_token=step.auth_token,
                prompt=resolved_prompt,
                config=step.config,
            )

        else:
            return BridgeResponse(
                success=False,
                error=f"Unknown step type '{step.type}'. Expected: local, cloud, local_agent, remote_agent",
                backend=step.backend,
                model=step.model,
            )

    def _build_cloud_metadata(self, result: ExecutionResult) -> Dict[str, Any]:
        """Build metadata safe to send to Brynq Cloud.

        This includes ONLY execution metadata — timing, token counts,
        success/fail. NEVER any user content, prompts, or model outputs.
        """
        step_metadata = []
        for step in result.steps:
            step_metadata.append({
                "step_id": step.step_id,
                "backend": step.backend,
                "model": step.model,
                "success": step.success,
                "elapsed_seconds": round(step.elapsed_seconds, 3),
                "token_estimate": step.token_estimate,
                # Intentionally excluded: output, error details (may contain user content)
            })

        return {
            "plan_id": result.plan_id,
            "success": result.success,
            "total_elapsed": round(result.total_elapsed, 3),
            "total_tokens": result.total_tokens,
            "steps_completed": len(result.steps),
            "steps": step_metadata,
        }
