"""
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 asyncio
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")

STEP_TIMEOUT_SECONDS = 120

# ---------------------------------------------------------------------------
# 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).
        critical: If True (default), failure aborts the plan. If False, failure
            is recorded in the step result but execution continues.
    """

    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 = ""
    critical: bool = True  # If False, failure won't abort the plan


@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
    error_category: Optional[str] = None  # "timeout", "auth_error", "model_error", "network_error"
    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


# ---------------------------------------------------------------------------
# Error categorization
# ---------------------------------------------------------------------------


def _categorize_error(exc: BaseException) -> str:
    """Classify an exception into a standard error category.

    Categories:
        timeout      — The step exceeded its time limit.
        auth_error   — Authentication or authorization failure (401/403).
        model_error  — The model/backend returned an application-level error.
        network_error — Connection or transport failure.

    Returns one of: "timeout", "auth_error", "model_error", "network_error".
    """
    import http.client
    import socket
    import urllib.error

    # Timeout variants
    if isinstance(exc, (TimeoutError, asyncio.TimeoutError, socket.timeout)):
        return "timeout"

    # HTTP errors — status code drives the category
    if isinstance(exc, urllib.error.HTTPError):
        if exc.code in (401, 403):
            return "auth_error"
        if exc.code in (400, 404, 422, 429, 500, 502, 503):
            return "model_error"
        return "model_error"

    # URL / connection errors
    if isinstance(exc, (urllib.error.URLError, ConnectionError, socket.error,
                        http.client.HTTPException, OSError)):
        return "network_error"

    # Fallback — treat unknown exceptions as model errors
    return "model_error"


def _categorize_error_message(error_msg: str) -> str:
    """Infer error category from an error message string.

    Used when the bridge caught the exception internally and returned it
    as an error string in BridgeResponse rather than raising.
    """
    if not error_msg:
        return "model_error"

    lower = error_msg.lower()

    # Timeout patterns
    if any(kw in lower for kw in ("timeout", "timed out", "time out", "deadline exceeded")):
        return "timeout"

    # Auth patterns
    if any(kw in lower for kw in ("401", "403", "unauthorized", "forbidden",
                                   "auth", "authentication", "permission denied",
                                   "invalid api key", "invalid key", "expired token")):
        return "auth_error"

    # Network patterns
    if any(kw in lower for kw in ("connection refused", "connection reset",
                                   "network", "dns", "unreachable", "econnrefused",
                                   "econnreset", "enotfound", "not running",
                                   "could not connect", "urlopen error")):
        return "network_error"

    # Default
    return "model_error"


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


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

    .. deprecated::
        The production code path in ``runtime/app.py`` uses its own
        ``_execute_plan_steps()`` which has OAuth-first credential
        resolution and retry logic. This class is retained for tests
        and backward compatibility. New features should go in app.py.

    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,
        workspace_root: str = ".",
    ) -> 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.
            workspace_root: Path to the repository root for codebase indexing.
        """
        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
        self.workspace_root = workspace_root

    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. Inject Codebase Context (if relevant)
        3. Route to correct bridge (local/cloud/agent)
        4. Execute and capture output
        5. Store result for downstream steps
        6. 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()

        # Initialize codebase indexer for contextual injection
        import sys
        import os
        from .codebase_indexer import CodebaseIndexer
        import json
        indexer = CodebaseIndexer(self.workspace_root)
        
        # Pull global project context once per plan
        project_context = indexer.get_project_context()

        # --- 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("[plan:%s] Signature validation failed", 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("[plan:%s] TTL expired", plan.plan_id)
                return result

        # --- Validate plan has steps ---
        if not plan.steps:
            result.success = False
            result.error = "Plan has no steps to execute"
            result.total_elapsed = time.monotonic() - t_start
            log.error("[plan:%s] Empty steps list", plan.plan_id)
            return result

        # --- Validate unique step_ids ---
        seen_ids: Dict[str, int] = {}
        for step in plan.steps:
            seen_ids[step.step_id] = seen_ids.get(step.step_id, 0) + 1
        duplicates = {sid for sid, count in seen_ids.items() if count > 1}
        if duplicates:
            log.warning(
                "[plan:%s] Duplicate step_id values found: %s",
                plan.plan_id, ", ".join(sorted(duplicates)),
            )

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

        total_steps = len(plan.steps)
        for step_index, step in enumerate(plan.steps, 1):
            log.info(
                "[plan:%s] Step %d/%d: %s via %s",
                plan.plan_id, step_index, total_steps, step.type, step.model,
            )
            log.info(
                "[plan:%s] Executing %s: type=%s backend=%s model=%s",
                plan.plan_id, step.step_id, step.type, step.backend, step.model,
            )

            # 0. Validate prompt is non-empty before routing
            if not step.prompt or not step.prompt.strip():
                log.warning(
                    "[plan:%s] Step %s has empty prompt — skipping",
                    plan.plan_id, step.step_id,
                )
                step_result = StepResult(
                    step_id=step.step_id,
                    output="",
                    success=False,
                    error="Step has empty prompt",
                    elapsed_seconds=0.0,
                    token_estimate=0,
                    backend=step.backend,
                    model=step.model,
                )
                result.steps.append(step_result)
                previous_results[step.step_id] = step_result
                continue

            # 1. Resolve context references
            resolved_prompt = resolve_context(step, plan, previous_results)
            
            # 1.5 Inject Codebase Context automatically if the step seems to deal with code
            # In a real scenario, this would be triggered by a specific flag in `step.config`,
            # but for this MVP we heuristically append it if not already present.
            if "project_context" not in resolved_prompt and "code" in plan.original_input.lower():
                context_header = f"\n\n--- AUTO-INJECTED PROJECT CONTEXT ---\n{project_context}\n-----------------------------------\n"
                resolved_prompt += context_header

            # 2. Route to correct bridge and execute (with per-step timeout)
            step_t0 = time.monotonic()
            try:
                bridge_response = await asyncio.wait_for(
                    asyncio.to_thread(self._route_step, step, resolved_prompt),
                    timeout=STEP_TIMEOUT_SECONDS,
                )
            except TimeoutError:
                elapsed = time.monotonic() - step_t0
                log.warning(
                    "[plan:%s] Step %s timed out after %.2fs (limit %ds)",
                    plan.plan_id, step.step_id, elapsed, STEP_TIMEOUT_SECONDS,
                )
                step_result = StepResult(
                    step_id=step.step_id,
                    output="",
                    success=False,
                    error=f"Step timed out after {elapsed:.2f}s (limit {STEP_TIMEOUT_SECONDS}s)",
                    error_category="timeout",
                    elapsed_seconds=elapsed,
                    token_estimate=0,
                    backend=step.backend,
                    model=step.model,
                )
                result.steps.append(step_result)
                previous_results[step.step_id] = step_result
                if step.critical:
                    result.success = False
                    result.error = f"Step {step.step_id} failed: {step_result.error}"
                    log.error("[plan:%s] Critical step %s timed out — aborting", plan.plan_id, step.step_id)
                    break
                continue
            except Exception as exc:
                elapsed = time.monotonic() - step_t0
                category = _categorize_error(exc)
                log.warning(
                    "[plan:%s] Step %s raised %s (category=%s): %s",
                    plan.plan_id, step.step_id, type(exc).__name__, category, exc,
                )
                step_result = StepResult(
                    step_id=step.step_id,
                    output="",
                    success=False,
                    error=f"{type(exc).__name__}: {exc}",
                    error_category=category,
                    elapsed_seconds=elapsed,
                    token_estimate=0,
                    backend=step.backend,
                    model=step.model,
                )
                result.steps.append(step_result)
                previous_results[step.step_id] = step_result
                if step.critical:
                    result.success = False
                    result.error = f"Step {step.step_id} failed: {step_result.error}"
                    log.error("[plan:%s] Critical step %s raised %s — aborting", plan.plan_id, step.step_id, type(exc).__name__)
                    break
                continue

            # 3. Capture result
            error_cat = None
            if not bridge_response.success and bridge_response.error:
                error_cat = _categorize_error_message(bridge_response.error)

            step_result = StepResult(
                step_id=step.step_id,
                output=bridge_response.output,
                success=bridge_response.success,
                error=bridge_response.error,
                error_category=error_cat,
                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)

            # 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(
                    "[plan:%s] Step %s completed: %d tokens in %.2fs",
                    plan.plan_id, step.step_id, step_result.token_estimate, step_result.elapsed_seconds,
                )
            elif step.critical:
                # Critical step failed — abort the plan
                result.success = False
                result.error = (
                    f"Step {step.step_id} failed: {step_result.error}"
                )
                log.error(
                    "[plan:%s] Step %s failed: %s — aborting",
                    plan.plan_id, step.step_id, step_result.error,
                )
                break
            else:
                # Non-critical step failed — continue execution
                log.warning(
                    "[plan:%s] Non-critical step %s failed: %s — continuing",
                    plan.plan_id, step.step_id, step_result.error,
                )

        result.total_elapsed = time.monotonic() - t_start
        result.total_tokens = sum(len(s.output) // 4 for s in result.steps)

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

        passed = sum(1 for s in result.steps if s.success)
        failed = len(result.steps) - passed
        log.info(
            "[plan:%s] Done: %d steps (%d passed, %d failed) in %.2fs, ~%d tokens total",
            plan.plan_id,
            len(result.steps),
            passed,
            failed,
            result.total_elapsed,
            result.total_tokens,
        )

        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":
            # Check if agent is installed locally using get_installed_agent_endpoint
            from .agent_bridge import get_installed_agent_endpoint
            manifest = get_installed_agent_endpoint(step.model)
            if manifest:
                return self.agent_bridge.call_installed_agent(
                    agent_id=step.model,
                    prompt=resolved_prompt,
                    config=step.config,
                )
            return self.agent_bridge.call_local(
                agent_id=step.model,
                endpoint=step.endpoint,
                prompt=resolved_prompt,
                config=step.config,
            )

        elif step_type == "remote_agent":
            from .agent_bridge import get_installed_agent_endpoint
            manifest = get_installed_agent_endpoint(step.model)
            if manifest:
                return self.agent_bridge.call_installed_agent(
                    agent_id=step.model,
                    prompt=resolved_prompt,
                    config=step.config,
                )
            return self.agent_bridge.call_remote(
                agent_id=step.model,
                endpoint=step.endpoint,
                auth_token=step.auth_token,
                prompt=resolved_prompt,
                config=step.config,
            )

        elif step_type == "agent":
            from .agent_bridge import get_installed_agent_endpoint
            manifest = get_installed_agent_endpoint(step.model)
            if manifest:
                return self.agent_bridge.call_installed_agent(
                    agent_id=step.model,
                    prompt=resolved_prompt,
                    config=step.config,
                )
            # Default to remote if not specified
            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:
            entry: Dict[str, Any] = {
                "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)
            }
            if step.error_category:
                entry["error_category"] = step.error_category
            step_metadata.append(entry)

        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,
        }
