"""Phase C: Orchestrator Engine.

Multi-agent workflow orchestrator with shared context.
Coordinates multiple agent calls within a single chain, manages
context propagation between agents, and handles fan-out/fan-in
patterns for parallel agent execution.

Orchestration model:
  - A workflow is a chain of agent_call steps with shared context
  - Context accumulates across steps (agent A's output → agent B's input)
  - Fan-out: one step routes to multiple agents in parallel
  - Fan-in: parallel results merge back into a single context key
"""

import json
import logging
import threading
import time
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

import requests

from app.models_agents import (
    get_agent_with_api_key,
    log_agent_call,
    update_agent_call,
)
from app.services.chain_engine import (
    STEP_HANDLERS,
    _build_step_context,
    _interpolate_config,
    execute_step,
)

logger = logging.getLogger("agentforms.orchestrator")


# ─── Shared Context ───────────────────────────────────────────────────────────


class SharedContext:
    """Thread-safe context store shared across agent calls in a workflow."""

    def __init__(self, site: dict, submission: dict):
        self._data: dict[str, Any] = {}
        self._site = site
        self._submission = submission
        self._lock = threading.Lock()

    def get(self, key: str, default: Any = None) -> Any:
        """Read a value from shared context."""
        with self._lock:
            return self._data.get(key, default)

    def set(self, key: str, value: Any):
        """Write a value to shared context."""
        with self._lock:
            self._data[key] = value

    def merge(self, updates: dict[str, Any], prefix: str | None = None):
        """Merge multiple values into context, optionally with a prefix."""
        with self._lock:
            for k, v in updates.items():
                self._data[f"{prefix}.{k}" if prefix else k] = v

    def to_dict(self) -> dict[str, Any]:
        """Snapshot of current context (for template interpolation)."""
        with self._lock:
            return dict(self._data)

    def build_interpolation_context(self) -> dict[str, Any]:
        """Full interpolation context including site, submission, and shared data."""
        base = _build_step_context(self._site, self._submission, {}, {})
        with self._lock:
            base["context"] = dict(self._data)
        return base


# ─── Single Agent Execution ───────────────────────────────────────────────────


def execute_agent_call(
    agent_id: int,
    user_id: int,
    context: SharedContext,
    prompt: str = "",
    chain_action_id: int | None = None,
    timeout: int | None = None,
) -> dict[str, Any]:
    """Execute a single agent call with shared context.

    Returns:
        {
            "success": bool,
            "agent_id": int,
            "output": dict,
            "error": str|None,
            "call_id": int|None,
            "duration_ms": int,
        }
    """
    agent = get_agent_with_api_key(agent_id, user_id)
    if not agent:
        return {
            "success": False,
            "agent_id": agent_id,
            "output": {},
            "error": f"Agent {agent_id} not found or disabled",
            "call_id": None,
            "duration_ms": 0,
        }

    # Build payload with current shared context
    payload = {
        "type": "agent_call",
        "site_id": context._site.get("id"),
        "submission_id": context._submission.get("id"),
        "submission_data": context._submission,
        "chain_context": context.to_dict(),
        "prompt": prompt,
    }

    # Log the call
    call_log_id = None
    try:
        call_log_id = log_agent_call(
            user_id=user_id,
            agent_id=agent_id,
            chain_action_id=chain_action_id,
            site_id=context._site.get("id"),
            submission_id=context._submission.get("id"),
            request_payload=payload,
        )
    except Exception:
        pass

    start_time = time.time()
    effective_timeout = timeout or agent.get("max_timeout_seconds", 60)

    headers = {"Content-Type": "application/json"}
    if agent.get("api_key"):
        headers["Authorization"] = f"Bearer {agent['api_key']}"

    try:
        resp = requests.post(
            agent["endpoint_url"],
            json=payload,
            headers=headers,
            timeout=effective_timeout,
        )
        duration_ms = int((time.time() - start_time) * 1000)

        if resp.status_code == 200:
            result = resp.json()
            agent_output = result.get("output", result) if isinstance(result, dict) else {"response": result}

            try:
                if call_log_id:
                    update_agent_call(call_log_id, "completed", response_payload=agent_output, duration_ms=duration_ms)
            except Exception:
                pass

            return {
                "success": True,
                "agent_id": agent_id,
                "output": agent_output,
                "error": None,
                "call_id": call_log_id,
                "duration_ms": duration_ms,
            }

        else:
            error_msg = f"Agent returned HTTP {resp.status_code}"
            try:
                if call_log_id:
                    update_agent_call(call_log_id, "failed", error_message=error_msg, duration_ms=duration_ms)
            except Exception:
                pass
            return {
                "success": False,
                "agent_id": agent_id,
                "output": {},
                "error": error_msg,
                "call_id": call_log_id,
                "duration_ms": duration_ms,
            }

    except requests.exceptions.Timeout:
        duration_ms = int((time.time() - start_time) * 1000)
        error_msg = f"Agent call timed out after {effective_timeout}s"
        try:
            if call_log_id:
                update_agent_call(call_log_id, "timeout", error_message=error_msg, duration_ms=duration_ms)
        except Exception:
            pass
        return {
            "success": False,
            "agent_id": agent_id,
            "output": {},
            "error": error_msg,
            "call_id": call_log_id,
            "duration_ms": duration_ms,
        }
    except Exception as e:
        duration_ms = int((time.time() - start_time) * 1000)
        error_msg = f"Agent call failed: {str(e)}"
        try:
            if call_log_id:
                update_agent_call(call_log_id, "failed", error_message=str(e), duration_ms=duration_ms)
        except Exception:
            pass
        return {
            "success": False,
            "agent_id": agent_id,
            "output": {},
            "error": error_msg,
            "call_id": call_log_id,
            "duration_ms": duration_ms,
        }


# ─── Parallel Agent Fan-Out ───────────────────────────────────────────────────


def execute_agents_parallel(
    agent_configs: list[dict[str, Any]],
    user_id: int,
    context: SharedContext,
    chain_action_id: int | None = None,
) -> list[dict[str, Any]]:
    """Execute multiple agents in parallel.

    Each agent_config: {agent_id, prompt, output_key, timeout}.
    output_key is the context key where the result gets stored.

    Returns list of result dicts from execute_agent_call.
    """
    results: list[dict[str, Any]] = []
    errors: list[str] = []

    def _run(config: dict[str, Any]):
        result = execute_agent_call(
            agent_id=config["agent_id"],
            user_id=user_id,
            context=context,
            prompt=config.get("prompt", ""),
            chain_action_id=chain_action_id,
            timeout=config.get("timeout"),
        )
        results.append(result)
        if not result["success"]:
            errors.append(result["error"] or "Unknown error")

    threads = []
    for config in agent_configs:
        t = threading.Thread(target=_run, args=(config,), daemon=True)
        threads.append(t)
        t.start()

    for t in threads:
        t.join(timeout=120)

    # Store results in shared context under their output keys
    for result in results:
        # Find the matching config for the output key
        for config in agent_configs:
            if config["agent_id"] == result["agent_id"]:
                output_key = config.get("output_key", f"agent_{result['agent_id']}")
                if result["success"]:
                    context.set(output_key, result["output"])
                else:
                    context.set(output_key, {"error": result["error"]})
                break

    return results


# ─── Workflow Orchestrator ────────────────────────────────────────────────────


def execute_workflow(
    workflow: dict[str, Any],
    site: dict,
    submission: dict,
) -> dict[str, Any]:
    """Execute a multi-agent workflow.

    Workflow format:
        {
            "phases": [
                {
                    "id": "phase_1",
                    "agents": [
                        {"agent_id": 1, "prompt": "...", "output_key": "analysis"},
                    ],
                    "parallel": false,   # true = fan-out, false = sequential
                    "on_success": "phase_2",
                    "on_failure": "fallback_phase",
                },
            ],
            "max_phases": 20,
            "continue_on_error": false,
        }

    Returns:
        {
            "success": bool,
            "phases_executed": int,
            "context": dict,
            "results": list,
            "error": str|None,
        }
    """
    user_id = site.get("user_id", 0)
    shared = SharedContext(site, submission)
    phases = workflow.get("phases", [])
    max_phases = workflow.get("max_phases", 20)
    continue_on_error = workflow.get("continue_on_error", False)

    if not phases:
        return {
            "success": True,
            "phases_executed": 0,
            "context": {},
            "results": [],
            "error": None,
        }

    phase_map = {p["id"]: p for p in phases}
    current_phase_id = phases[0]["id"]
    phases_executed = 0
    all_results: list[dict[str, Any]] = []
    workflow_error = None

    while current_phase_id and phases_executed < max_phases:
        phase = phase_map.get(current_phase_id)
        if not phase:
            workflow_error = f"Unknown phase: {current_phase_id}"
            break

        logger.info("Executing phase %s (%d agents)", current_phase_id, len(phase.get("agents", [])))
        phases_executed += 1

        # Execute phase agents
        if phase.get("parallel", False):
            phase_results = execute_agents_parallel(
                phase.get("agents", []),
                user_id=user_id,
                context=shared,
            )
        else:
            phase_results = []
            for agent_config in phase.get("agents", []):
                result = execute_agent_call(
                    agent_id=agent_config["agent_id"],
                    user_id=user_id,
                    context=shared,
                    prompt=agent_config.get("prompt", ""),
                )
                phase_results.append(result)

                # Store result in context
                output_key = agent_config.get("output_key", f"agent_{result['agent_id']}")
                if result["success"]:
                    shared.set(output_key, result["output"])
                else:
                    shared.set(output_key, {"error": result["error"]})

        all_results.extend(phase_results)

        # Check phase success
        phase_success = any(r["success"] for r in phase_results)

        if not phase_success:
            workflow_error = "Phase failed: " + " | ".join(r["error"] for r in phase_results if not r["success"])
            if not continue_on_error:
                break
            # Continue to next phase even on error
            current_phase_id = phase.get("on_failure") or phase.get("on_success")
        else:
            current_phase_id = phase.get("on_success")

    return {
        "success": workflow_error is None,
        "phases_executed": phases_executed,
        "context": shared.to_dict(),
        "results": all_results,
        "error": workflow_error,
    }