"""
pipeline_executor.py — Executes AI pipeline node graphs.

Each pipeline is a React Flow graph (nodes + edges) stored as JSON.
Execution walks the graph in BFS order, passing a shared context dict
between nodes. LLM calls use the user's stored API keys.

Node types:
  p_input     → injects query into context
  p_prompt    → renders system + user template with {query}/{context} vars
  p_llm       → calls OpenAI or Anthropic with the rendered prompt
  p_tool_http → HTTP call to an external API
  p_memory    → read/write to ChromaDB vector store (long-term memory)
  p_agent     → calls another pipeline as a sub-agent (multi-agent)
  p_guard     → scores the accumulated steps with LabelFreeScorer
                and fires alert workflows via run_automations()
  p_route     → branches: "safe" edge vs "review" edge based on risk score
  p_output    → collects final result

Context dict passed between nodes:
{
    "query":       str,        # original user input
    "prompt":      str,        # rendered prompt after p_prompt
    "llm_output":  str,        # LLM response text
    "tool_output": str,        # HTTP tool response
    "context":     str,        # general carry-forward string
    "risk_score":  float|None, # set by p_guard
    "steps":       list[dict], # accumulated reasoning steps for scoring
    "trace":       list[dict], # [{node_id, node_type, label, output, duration_ms, error}]
}

Streaming:
    Use stream_pipeline() as an async generator to get per-node events.
    Use run_pipeline() to collect the final result dict synchronously.
"""

from __future__ import annotations

import asyncio
import logging
import os
import re
import time
from typing import TYPE_CHECKING, AsyncIterator

if TYPE_CHECKING:
    from .store import ChainStore

logger = logging.getLogger("pipeline_executor")


# ── Retry helper ───────────────────────────────────────────────────────────────

async def _with_retry(coro_factory, max_retries: int = 2):
    """Retry with exponential backoff. Gives up immediately on auth errors."""
    last_exc: Exception = RuntimeError("no attempts")
    for attempt in range(max_retries + 1):
        try:
            return await coro_factory()
        except Exception as e:
            last_exc = e
            msg = str(e).lower()
            if any(kw in msg for kw in ("api key", "authentication", "invalid_api_key", "unauthorized")):
                raise  # don't retry auth failures
            if attempt < max_retries:
                await asyncio.sleep(2 ** attempt)  # 1s, 2s
    raise last_exc


# ── Template renderer ──────────────────────────────────────────────────────────

def _render(template: str, ctx: dict) -> str:
    """Replace {query}, {context}, {llm_output}, {tool_output} in template."""
    try:
        return template.format(
            query=ctx.get("query", ""),
            context=ctx.get("context", ""),
            llm_output=ctx.get("llm_output", ""),
            tool_output=ctx.get("tool_output", ""),
        )
    except (KeyError, ValueError):
        return template


# ── LLM helpers ────────────────────────────────────────────────────────────────

async def _call_openai(prompt: str, system: str, config: dict, api_key: str) -> str:
    """Call OpenAI chat completion with timeout + retry."""
    from openai import AsyncOpenAI
    client  = AsyncOpenAI(api_key=api_key, timeout=30.0)
    model   = config.get("model") or "gpt-4o-mini"
    temp    = float(config.get("temperature", 0.7))
    max_tok = int(config.get("max_tokens", 1024))
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})
    resp = await _with_retry(
        lambda: client.chat.completions.create(
            model=model, messages=messages, temperature=temp, max_tokens=max_tok
        )
    )
    return resp.choices[0].message.content or ""


async def _call_anthropic(prompt: str, system: str, config: dict, api_key: str) -> str:
    """Call Anthropic messages API with timeout + retry."""
    import anthropic
    client  = anthropic.AsyncAnthropic(api_key=api_key, timeout=30.0)
    model   = config.get("model") or "claude-haiku-4-5-20251001"
    max_tok = int(config.get("max_tokens", 1024))
    kwargs: dict = dict(
        model=model,
        max_tokens=max_tok,
        messages=[{"role": "user", "content": prompt}],
    )
    if system:
        kwargs["system"] = system
    resp = await _with_retry(lambda: client.messages.create(**kwargs))
    return resp.content[0].text if resp.content else ""


# ── Node executors ─────────────────────────────────────────────────────────────

async def _exec_input(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    ctx["context"] = ctx.get("query", "")
    _trace(ctx, node, "Input received", ctx["query"][:200])


async def _exec_prompt(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    config     = node.get("data", {}).get("config", {})
    system_tpl = config.get("system_prompt", "You are a helpful assistant.")
    user_tpl   = config.get("user_template", "{query}")
    ctx["_system_prompt"] = _render(system_tpl, ctx)
    ctx["prompt"]         = _render(user_tpl, ctx)
    _trace(ctx, node, "Prompt rendered", ctx["prompt"][:300])


async def _exec_llm(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    config   = node.get("data", {}).get("config", {})
    provider = config.get("provider", "openai").lower()
    prompt   = ctx.get("prompt") or ctx.get("query", "")
    system   = ctx.get("_system_prompt", "You are a helpful assistant.")

    if provider == "anthropic":
        api_key = user_configs.get("anthropic_key", "")
        if not api_key:
            raise ValueError("Anthropic API key not configured. Go to Settings → LLM API Keys.")
        output = await _call_anthropic(prompt, system, config, api_key)
    else:
        api_key = user_configs.get("openai_key", "")
        if not api_key:
            raise ValueError("OpenAI API key not configured. Go to Settings → LLM API Keys.")
        output = await _call_openai(prompt, system, config, api_key)

    ctx["llm_output"] = output
    ctx["context"]    = output
    ctx.setdefault("steps", []).append({
        "thought":     f"LLM ({provider}) generated response",
        "action":      "llm_call",
        "observation": output[:500],
    })
    _trace(ctx, node, f"LLM ({provider})", output[:400])


async def _exec_tool_http(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    import httpx
    config   = node.get("data", {}).get("config", {})
    url      = config.get("url", "")
    method   = config.get("method", "POST").upper()
    body_tpl = config.get("body_template", '{"query": "{query}"}')
    if not url:
        _trace(ctx, node, "HTTP Tool", "No URL configured")
        return
    body_str = _render(body_tpl, ctx)
    try:
        import json as _json
        body = _json.loads(body_str)
    except Exception:
        body = {"query": ctx.get("query", "")}

    async with httpx.AsyncClient(timeout=15.0) as client:
        if method == "GET":
            resp = await client.get(url, params=body)
        else:
            resp = await client.post(url, json=body)
        output = resp.text[:1000]

    ctx["tool_output"] = output
    ctx["context"]     = output
    ctx.setdefault("steps", []).append({
        "thought":     f"HTTP {method} to {url}",
        "action":      "tool_call",
        "observation": output[:500],
    })
    _trace(ctx, node, f"HTTP {method}", output[:300])


async def _exec_memory(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    """ChromaDB vector memory — store embeddings or retrieve similar docs."""
    config          = node.get("data", {}).get("config", {})
    action          = config.get("action", "retrieve")   # "store" or "retrieve"
    collection_name = config.get("collection", f"user_{user_id}_memory")
    top_k           = int(config.get("top_k", 3))

    try:
        import chromadb
        chroma_path = os.path.expanduser(os.environ.get("QPPG_CHROMA_PATH", "~/.qppg/chroma"))
        chroma_client = chromadb.PersistentClient(path=chroma_path)
        collection = chroma_client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )

        if action == "store":
            text = ctx.get("llm_output") or ctx.get("context") or ctx.get("query", "")
            if text:
                collection.add(
                    documents=[text],
                    ids=[f"doc_{int(time.time() * 1000)}"],
                    metadatas=[{"query": ctx.get("query", "")[:200], "user_id": user_id}],
                )
            output = f"Stored {len(text)} chars to '{collection_name}' ({collection.count()} total)"

        else:  # retrieve
            query_text = ctx.get("query", "")
            count = collection.count()
            if count == 0:
                output = "Memory collection is empty."
            else:
                results = collection.query(
                    query_texts=[query_text],
                    n_results=min(top_k, count),
                )
                docs = results.get("documents", [[]])[0]
                output = "\n\n---\n\n".join(docs) if docs else "No relevant memories found."
            ctx["context"] = output
            ctx.setdefault("steps", []).append({
                "thought":     f"Retrieved {top_k} memories from '{collection_name}'",
                "action":      "memory_retrieve",
                "observation": output[:300],
            })

        _trace(ctx, node, f"Memory {action}", output[:300])

    except ImportError:
        _trace(ctx, node, "Memory", "chromadb not installed (pip install chromadb)")
    except Exception as e:
        logger.warning(f"Memory node error: {e}")
        _trace(ctx, node, "Memory Error", str(e)[:200])


async def _exec_agent(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    """Call another pipeline as a sub-agent (multi-agent orchestration)."""
    config      = node.get("data", {}).get("config", {})
    pipeline_id = config.get("pipeline_id")

    if not pipeline_id:
        _trace(ctx, node, "Agent", "No pipeline_id configured")
        return

    sub_pipeline = store.get_pipeline(int(pipeline_id), user_id)
    if not sub_pipeline:
        _trace(ctx, node, "Agent Error", f"Pipeline {pipeline_id} not found")
        return

    sub_query = ctx.get("llm_output") or ctx.get("context") or ctx.get("query", "")
    result    = await run_pipeline(sub_pipeline, sub_query, user_id, user_configs, store)

    output = result.get("result", "")
    ctx["context"]    = output
    ctx["llm_output"] = output
    ctx.setdefault("steps", []).append({
        "thought":     f"Sub-agent (pipeline {pipeline_id}) returned",
        "action":      "agent_call",
        "observation": output[:300],
    })
    # Merge sub-agent trace with prefix so it's visible in the parent trace
    for sub_entry in result.get("trace", []):
        ctx["trace"].append({
            **sub_entry,
            "node_id": f"sub[{pipeline_id}]/{sub_entry.get('node_id', '')}",
        })
    _trace(ctx, node, f"Agent [pipeline {pipeline_id}]", output[:300])


async def _exec_guard(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    from .label_free_scorer import LabelFreeScorer
    from .workflow_executor import run_automations

    config   = node.get("data", {}).get("config", {})
    domain   = config.get("domain") or "pipeline"
    question = ctx.get("query", "")
    steps    = ctx.get("steps") or []

    scorer = LabelFreeScorer()
    result = scorer.score(question, steps, ctx.get("llm_output", ""), finished=True)
    ctx["risk_score"] = result.risk_score

    chain_id = store.add_chain(
        f"user_{user_id}:{domain}",
        {"question": question, "steps": steps, "final_answer": ctx.get("llm_output", ""), "finished": True},
        risk_score=result.risk_score,
        alert=result.needs_review,
    )

    automation = await run_automations(
        user_id=user_id, domain=domain, risk_score=result.risk_score,
        question=question, chain_id=chain_id, store=store,
    )
    ctx["guardrail_action"]  = automation.get("guardrail_action")
    ctx["fallback_required"] = automation.get("fallback_required", False)

    badge = "🔴 HIGH" if result.risk_score >= 0.65 else "🟡 MED" if result.risk_score >= 0.4 else "🟢 LOW"
    _trace(ctx, node, "Guard Score", f"{badge} {result.risk_score:.0%} | action: {ctx['guardrail_action'] or 'none'}")


async def _exec_eval(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    """Run one or more pre-built evaluators (PII, toxicity, hallucination, off-topic)."""
    from .evaluators import run_evaluator, EVALUATOR_REGISTRY
    config     = node.get("data", {}).get("config", {})
    # evaluators: comma-sep string e.g. "pii_detector,toxicity_classifier"
    eval_names = [e.strip() for e in config.get("evaluators", "pii_detector").split(",") if e.strip()]
    text       = ctx.get("llm_output") or ctx.get("context") or ctx.get("query", "")
    reference  = ctx.get("query", "") + " " + ctx.get("tool_output", "")

    eval_results: dict = {}
    any_failed = False
    for name in eval_names:
        if name not in EVALUATOR_REGISTRY:
            continue
        r = run_evaluator(name, text, reference, config)
        eval_results[name] = r.to_dict()
        if not r.passed:
            any_failed = True

    ctx["eval_results"] = eval_results
    ctx.setdefault("steps", []).append({
        "thought":     f"Evaluators ran: {', '.join(eval_names)}",
        "action":      "evaluate",
        "observation": str({k: v["label"] for k, v in eval_results.items()}),
    })
    summary = "; ".join(f"{k}:{v['label']}" for k, v in eval_results.items())
    flag = "⚠️ FAIL" if any_failed else "✓ PASS"
    _trace(ctx, node, f"Eval {flag}", summary[:400])


async def _exec_react(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    """
    ReAct loop node: Thought → Action → Observation, iterates until stop condition.

    config fields:
      provider        : openai | anthropic
      model           : model id
      max_steps       : max loop iterations (default 5)
      stop_token      : string that ends the loop if found in LLM output (default "Final Answer:")
      tools_json      : JSON string of available tools [{name, description, url}]
      system_prompt   : system prompt for the loop
    """
    config     = node.get("data", {}).get("config", {})
    provider   = config.get("provider", "openai").lower()
    max_steps  = int(config.get("max_steps", 5))
    stop_token = config.get("stop_token", "Final Answer:")
    tools_raw  = config.get("tools_json", "[]")
    sys_prompt = config.get("system_prompt",
        "You are a ReAct agent. Think step by step, use tools, then give a Final Answer.")

    try:
        import json as _json
        tools: list[dict] = _json.loads(tools_raw)
    except Exception:
        tools = []

    # Build tool descriptions for the prompt
    tool_desc = ""
    if tools:
        tool_desc = "\nAvailable tools:\n" + "\n".join(
            f"  {t.get('name', 'tool')}: {t.get('description', '')}" for t in tools
        )

    # Initial prompt
    query = ctx.get("query", "")
    loop_history = f"Question: {query}\n{tool_desc}\n\nThought: "

    api_key_openai    = user_configs.get("openai_key", "")
    api_key_anthropic = user_configs.get("anthropic_key", "")

    steps_taken: list[dict] = []
    final_answer = ""

    for step in range(max_steps):
        # Call LLM
        try:
            if provider == "anthropic":
                if not api_key_anthropic:
                    raise ValueError("Anthropic API key not configured.")
                llm_out = await _call_anthropic(loop_history, sys_prompt, config, api_key_anthropic)
            else:
                if not api_key_openai:
                    raise ValueError("OpenAI API key not configured.")
                llm_out = await _call_openai(loop_history, sys_prompt, config, api_key_openai)
        except Exception as e:
            _trace(ctx, node, f"ReAct error (step {step+1})", str(e)[:200])
            break

        loop_history += llm_out + "\n"

        # Check stop condition
        if stop_token and stop_token.lower() in llm_out.lower():
            # Extract final answer after stop token
            idx = llm_out.lower().find(stop_token.lower())
            final_answer = llm_out[idx + len(stop_token):].strip()
            steps_taken.append({"step": step + 1, "thought": llm_out[:300], "observation": "[DONE]"})
            # Ensure the final step is included in ctx["steps"] for Guard scoring
            ctx.setdefault("steps", []).append({
                "thought":     llm_out[:300],
                "action":      "final_answer",
                "observation": final_answer[:300],
            })
            break

        # Parse Action from LLM output (look for "Action: tool_name(args)")
        action_match = re.search(r'Action:\s*(\w+)\s*\(([^)]*)\)', llm_out, re.I)
        observation = ""
        if action_match and tools:
            tool_name = action_match.group(1)
            tool_args = action_match.group(2).strip()
            # Find matching tool
            tool_def = next((t for t in tools if t.get("name", "") == tool_name), None)
            if tool_def and tool_def.get("url"):
                try:
                    import httpx
                    async with httpx.AsyncClient(timeout=10.0) as client:
                        resp = await client.post(tool_def["url"], json={"query": tool_args})
                        observation = resp.text[:500]
                except Exception as e:
                    observation = f"Tool error: {e}"
            else:
                observation = f"Tool '{tool_name}' not found."
        else:
            observation = "(no action taken)"

        loop_history += f"Observation: {observation}\nThought: "
        steps_taken.append({
            "step": step + 1,
            "thought": llm_out[:300],
            "observation": observation[:300],
        })
        ctx.setdefault("steps", []).append({
            "thought":     llm_out[:300],
            "action":      action_match.group(1) if action_match else "think",
            "observation": observation[:300],
        })

    if not final_answer:
        final_answer = ctx.get("llm_output", "") or loop_history[-500:]

    ctx["llm_output"]  = final_answer
    ctx["context"]     = final_answer
    ctx["react_steps"] = steps_taken
    _trace(ctx, node, f"ReAct ({len(steps_taken)} steps)", final_answer[:400])


async def _exec_route(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> str:
    """Returns edge label to follow: 'safe' or 'review'."""
    config    = node.get("data", {}).get("config", {})
    threshold = float(config.get("threshold", 0.65))
    risk      = ctx.get("risk_score", 0.0)
    edge      = "safe" if risk < threshold else "review"
    _trace(ctx, node, "Route", f"risk={risk:.0%} threshold={threshold:.0%} → {edge}")
    return edge


async def _exec_output(node: dict, ctx: dict, user_configs: dict, store: "ChainStore", user_id: int) -> None:
    output = ctx.get("llm_output") or ctx.get("tool_output") or ctx.get("context") or ctx.get("query", "")
    ctx["final_output"] = output
    _trace(ctx, node, "Output", output[:400])


# ── Trace helper ───────────────────────────────────────────────────────────────

def _trace(ctx: dict, node: dict, label: str, output: str) -> None:
    ctx.setdefault("trace", []).append({
        "node_id":   node.get("id", ""),
        "node_type": node.get("data", {}).get("nodeType", ""),
        "label":     label,
        "output":    output,
        "timestamp": time.time(),
    })


# ── Node executor dispatch ─────────────────────────────────────────────────────

_HANDLERS = {
    "p_input":     _exec_input,
    "p_prompt":    _exec_prompt,
    "p_llm":       _exec_llm,
    "p_tool_http": _exec_tool_http,
    "p_memory":    _exec_memory,
    "p_agent":     _exec_agent,
    "p_eval":      _exec_eval,
    "p_react":     _exec_react,
    "p_guard":     _exec_guard,
    "p_route":     _exec_route,
    "p_output":    _exec_output,
}


# ── Core streaming generator ───────────────────────────────────────────────────

async def stream_pipeline(
    pipeline: dict,
    query: str,
    user_id: int,
    user_configs: dict,
    store: "ChainStore",
) -> AsyncIterator[dict]:
    """
    Execute a pipeline and yield SSE-style event dicts as each node completes.

    Event shapes:
        {"type": "node_start",  "node_id": str, "node_type": str, "label": str}
        {"type": "node_done",   "node_id": str, "node_type": str, "label": str,
                                "output": str, "duration_ms": int}
        {"type": "node_error",  "node_id": str, "node_type": str, "error": str}
        {"type": "done",        "result": str, "trace": list, "risk_score": float|None,
                                "guardrail_action": str|None, "fallback_required": bool,
                                "status": "done"|"error", "error": str|None}
    """
    nodes: list[dict] = pipeline.get("nodes") or []
    edges: list[dict] = pipeline.get("edges") or []

    if not nodes:
        yield {"type": "done", "result": "", "trace": [], "risk_score": None,
               "guardrail_action": None, "fallback_required": False,
               "status": "error", "error": "Pipeline has no nodes"}
        return

    adj: dict[str, list[tuple[str, str]]] = {}
    incoming: set[str] = set()
    for edge in edges:
        src   = edge.get("source", "")
        tgt   = edge.get("target", "")
        label = edge.get("label") or edge.get("sourceHandle") or ""
        if src and tgt:
            adj.setdefault(src, []).append((tgt, label))
            incoming.add(tgt)

    node_by_id: dict[str, dict] = {n["id"]: n for n in nodes}
    roots = [n for n in nodes if n["id"] not in incoming] or nodes[:1]

    ctx: dict = {
        "query": query, "prompt": "", "llm_output": "", "tool_output": "",
        "context": "", "risk_score": None, "guardrail_action": None,
        "fallback_required": False, "steps": [], "trace": [], "final_output": "",
    }

    queue: list[dict] = list(roots)
    visited: set[str] = set()

    try:
        while queue:
            node  = queue.pop(0)
            nid   = node["id"]
            if nid in visited:
                continue
            visited.add(nid)

            ntype   = node.get("data", {}).get("nodeType", "")
            handler = _HANDLERS.get(ntype)
            label   = node.get("data", {}).get("label", ntype)

            yield {"type": "node_start", "node_id": nid, "node_type": ntype, "label": label}

            edge_to_follow: str | None = None
            t0 = time.time()
            error_msg: str | None = None

            try:
                if handler:
                    ret = await handler(node, ctx, user_configs, store, user_id)
                    if ntype == "p_route":
                        edge_to_follow = ret
                else:
                    _trace(ctx, node, f"Unknown node ({ntype})", "skipped")
            except Exception as e:
                error_msg = str(e)
                logger.error(f"Node {nid} ({ntype}) error: {e}", exc_info=True)
                ctx["trace"].append({
                    "node_id": nid, "node_type": ntype,
                    "label": "Error", "output": error_msg,
                    "timestamp": time.time(), "error": True,
                })
                yield {"type": "node_error", "node_id": nid, "node_type": ntype, "error": error_msg}
                if "API key" in error_msg:
                    yield {
                        "type": "done", "result": "", "trace": ctx["trace"],
                        "risk_score": ctx.get("risk_score"),
                        "guardrail_action": ctx.get("guardrail_action"),
                        "fallback_required": ctx.get("fallback_required", False),
                        "status": "error", "error": error_msg,
                    }
                    return

            duration_ms = int((time.time() - t0) * 1000)
            last_trace  = ctx["trace"][-1] if ctx["trace"] else {}
            yield {
                "type":        "node_done",
                "node_id":     nid,
                "node_type":   ntype,
                "label":       last_trace.get("label", label),
                "output":      last_trace.get("output", ""),
                "duration_ms": duration_ms,
            }

            for child_id, lbl in adj.get(nid, []):
                child = node_by_id.get(child_id)
                if not child:
                    continue
                if edge_to_follow is not None and lbl and lbl != edge_to_follow:
                    continue
                queue.append(child)

        yield {
            "type":              "done",
            "result":            ctx.get("final_output") or ctx.get("llm_output") or "",
            "trace":             ctx["trace"],
            "risk_score":        ctx.get("risk_score"),
            "guardrail_action":  ctx.get("guardrail_action"),
            "fallback_required": ctx.get("fallback_required", False),
            "status":            "done",
            "error":             None,
        }

    except Exception as e:
        logger.error(f"Pipeline execution failed: {e}", exc_info=True)
        yield {
            "type":              "done",
            "result":            "",
            "trace":             ctx.get("trace", []),
            "risk_score":        None,
            "guardrail_action":  None,
            "fallback_required": False,
            "status":            "error",
            "error":             str(e),
        }


# ── Synchronous wrapper ────────────────────────────────────────────────────────

async def run_pipeline(
    pipeline: dict,
    query: str,
    user_id: int,
    user_configs: dict,
    store: "ChainStore",
) -> dict:
    """
    Execute a pipeline graph and return the final result dict.
    Collects all stream events and returns the "done" event payload.
    """
    result: dict = {
        "result": "", "trace": [], "risk_score": None,
        "guardrail_action": None, "fallback_required": False,
        "status": "error", "error": "no events",
    }
    async for event in stream_pipeline(pipeline, query, user_id, user_configs, store):
        if event["type"] == "done":
            result = {k: v for k, v in event.items() if k != "type"}
    return result
