"""Run orchestration: our DB rows + LFX execution + terminal-state mapping."""

from __future__ import annotations

import json
from collections.abc import AsyncIterator
from uuid import UUID

import httpx
from sqlmodel import Session

from app.db import get_engine
from app.models import FlowRun, FlowVersion

from . import input_schema as form_schema
from . import store
from .lfx_client import LfxClient, LfxNotFound, LfxRunFailed
from .store import InvalidError


def build_envelope(
    *,
    payload: dict,
    run_id: UUID,
    callback_base_url: str,
    continuation_flow_id: UUID | None,
    callback_token: str,
    fields: list[str] | None = None,
) -> str:
    """Run envelope for LFX. ``fields`` (asked form-field names, from the
    version's cached input schema) lets form components emit defaults only
    for asked fields; omitted (None) preserves the legacy shape."""
    body: dict = {
        "payload": payload,
        "run_id": str(run_id),
        "callback_url": callback_base_url.rstrip("/"),
        "continuation_flow_id": (str(continuation_flow_id) if continuation_flow_id else None),
        "callback_token": callback_token,
    }
    if fields is not None:
        body["fields"] = list(fields)
    return json.dumps(body)


def _task_exists(session: Session, run_id: UUID) -> bool:
    return store.has_open_task(session, run_id)


async def execute_run(
    session: Session,
    lfx: LfxClient,
    version: FlowVersion,
    *,
    payload: dict,
    tenant_id: str = "default",
    session_id: str | None = None,
    callback_base_url: str,
    continuation_flow_id: UUID | None = None,
    callback_token: str = "",
) -> FlowRun:
    """Run a deployed version to a terminal state (no interrupts).

    The flow always terminates: either it completes, or its approval branch
    POSTs back to /api/approvals (creating a task) and terminates with a
    pending marker. Status is derived from that side effect, not from
    result-text parsing.
    """
    if not version.lfx_flow_id:
        raise InvalidError(f"version {version.id} is not deployed to LFX")
    run = store.create_run(
        session, version, input=payload, tenant_id=tenant_id, session_id=session_id
    )
    cached = version.input_schema
    envelope = build_envelope(
        payload=payload,
        run_id=run.id,
        callback_base_url=callback_base_url,
        continuation_flow_id=continuation_flow_id,
        callback_token=callback_token,
        fields=form_schema.schema_field_names(cached),
    )
    try:
        result = await lfx.run_flow(
            version.lfx_flow_id,
            input_value=envelope,
            session_id=session_id or str(run.id),
        )
    except LfxRunFailed as exc:
        return store.set_run_result(session, run, "failed", {"error": str(exc)})
    output = {
        "result": result.result,
        "component": result.component,
        "type": result.type,
    }
    if _task_exists(session, run.id):
        return store.set_run_result(session, run, "pending_approval", output)
    return store.set_run_result(session, run, "completed", output)


async def stream_run(
    session: Session,
    lfx: LfxClient,
    version: FlowVersion,
    *,
    payload: dict,
    tenant_id: str,
    session_id: str | None,
    callback_base_url: str,
    continuation_flow_id: UUID | None,
    callback_token: str,
) -> tuple[FlowRun, str, str, str]:
    """Create a run + envelope for SSE streaming. Returns (run, envelope, lfx_id, sid).

    The actual SSE bytes are produced by the API layer; finalization uses a
    fresh session (the request session may be closed when the generator ends).
    """
    if not version.lfx_flow_id:
        raise InvalidError(f"version {version.id} is not deployed to LFX")
    run = store.create_run(
        session, version, input=payload, tenant_id=tenant_id, session_id=session_id
    )
    cached = version.input_schema
    envelope = build_envelope(
        payload=payload,
        run_id=run.id,
        callback_base_url=callback_base_url,
        continuation_flow_id=continuation_flow_id,
        callback_token=callback_token,
        fields=form_schema.schema_field_names(cached),
    )
    return run, envelope, version.lfx_flow_id, session_id or str(run.id)


async def stream_lfx_chunks(
    lfx: LfxClient, lfx_flow_id: str, envelope: str, sid: str
) -> AsyncIterator[str]:
    """Yield raw SSE chunks from LFX, mapping transport errors to error events."""
    try:
        async with (
            httpx.AsyncClient(
                base_url=lfx.base_url, timeout=lfx.timeout, headers=lfx._headers()
            ) as client,
            client.stream(
                "POST",
                f"/flows/{lfx_flow_id}/stream",
                json={"input_value": envelope, "session_id": sid},
            ) as response,
        ):
            if response.status_code == 404:
                raise LfxNotFound(f"flow {lfx_flow_id} unknown to LFX")
            response.raise_for_status()
            async for chunk in response.aiter_text():
                yield chunk
    except (LfxNotFound, httpx.HTTPError) as exc:
        yield f"event: error\ndata: {json.dumps({'error': str(exc)})}\n\n"


def finalize_streamed_run(run_id: UUID) -> str:
    """Mark a streamed run completed/pending_approval from a fresh session."""
    with Session(get_engine()) as fresh:
        try:
            fresh_run = store.get_run(fresh, run_id)
            status = "pending_approval" if store.has_open_task(fresh, run_id) else "completed"
            store.set_run_result(fresh, fresh_run, status, {"streamed": True})
        except store.NotFoundError:
            status = "failed"
    return status