"""Workflow execution engine — runs workflow steps sequentially with state persistence."""

from __future__ import annotations

import json
import logging
from datetime import datetime, timezone

import asyncpg
from jinja2 import Template
from redis.asyncio import Redis

from loom.bus.events import EventType
from loom.bus.publisher import publish_event
from loom.config import SkillsConfig
from loom.graph import store
from loom.skills.loader import load_skill
from loom.skills.runner import run_skill
from loom.workflows.loader import WorkflowDefinition, WorkflowStep, load_workflow

log = logging.getLogger(__name__)


def _render_inputs(inputs: dict, state: dict) -> dict:
    """Render Jinja2 templates in step inputs using accumulated workflow state."""
    rendered = {}
    for key, val in inputs.items():
        if isinstance(val, str) and "{{" in val:
            rendered[key] = Template(val).render(**state)
        else:
            rendered[key] = val
    return rendered


async def _execute_skill_step(
    step: WorkflowStep,
    state: dict,
    config: SkillsConfig,
    pool: asyncpg.Pool,
    project_id: str,
    project_dir: str | None,
) -> dict:
    """Execute a skill step: load skill, render inputs, run, return output."""
    rendered = _render_inputs(step.inputs, state)
    skill = load_skill(step.skill, project_dir)
    return await run_skill(skill, rendered, config, pool, project_id)


async def _execute_action_step(
    step: WorkflowStep,
    state: dict,
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
) -> dict:
    """Execute a built-in action step."""
    rendered = _render_inputs(step.inputs, state)
    action = step.action

    if action == "write_task_graph":
        # Accept task graph data and write tasks to DB via store
        tasks_data = rendered.get("tasks")
        if isinstance(tasks_data, str):
            tasks_data = json.loads(tasks_data)

        if isinstance(tasks_data, dict) and "tasks" in tasks_data:
            tasks_data = tasks_data["tasks"]

        if not isinstance(tasks_data, list):
            return {"written": 0, "error": "No valid task list provided"}

        from loom.graph.task import Priority, Task, TaskStatus
        from loom.ids import task_id as gen_task_id

        # Build Task objects from the raw dicts
        task_objects = []
        for td in tasks_data:
            tid = gen_task_id()
            status_str = td.get("status", "pending")
            try:
                status = TaskStatus(status_str)
            except ValueError:
                status = TaskStatus.PENDING

            priority_str = td.get("priority", "p1")
            try:
                priority = Priority(priority_str)
            except ValueError:
                priority = Priority.P1

            task_objects.append(Task(
                id=tid,
                project_id=project_id,
                title=td.get("title", "Untitled"),
                status=status,
                priority=priority,
                context=td.get("context", {}),
                done_when=td.get("done_when"),
                depends_on=td.get("depends_on", []),
            ))

        created = await store.create_tasks_batch(pool, task_objects)
        return {"written": len(created), "task_ids": [t.id for t in created]}

    elif action == "load_task":
        task_id = rendered.get("task_id", "")
        task = await store.get_task(pool, task_id)
        return task.model_dump(mode="json")

    else:
        return {"error": f"Unknown action: {action}"}


async def _execute_condition_step(
    step: WorkflowStep,
    state: dict,
) -> str | None:
    """Evaluate a condition and return the next step name to jump to."""
    if not step.condition:
        return None
    result = Template(step.condition).render(**state)
    # Evaluate as truthy: non-empty, non-"false", non-"0"
    is_true = result.strip().lower() not in ("", "false", "0", "none")
    if is_true:
        return step.on_true
    return step.on_false


async def run_workflow(
    workflow: WorkflowDefinition,
    inputs: dict,
    config: SkillsConfig,
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
    project_dir: str | None = None,
) -> dict:
    """Execute a workflow step by step, persisting state to Postgres.

    Returns the workflow run dict (from store) on completion or pause.
    """
    # Create the workflow run record
    run = await store.create_workflow_run(pool, project_id, workflow.name, inputs)
    run_id = run["id"]

    # Initialize state: workflow inputs + outputs accumulator
    state: dict = {**inputs, "outputs": {}}

    # Publish start event
    await publish_event(
        redis, project_id, EventType.WORKFLOW_STARTED,
        payload={"run_id": run_id, "workflow": workflow.name},
    )

    # Update status to running
    await store.update_workflow_run(pool, run_id, status="running", state=state)

    try:
        step_index = 0
        while step_index < len(workflow.steps):
            step = workflow.steps[step_index]

            # Record step start
            await store.record_workflow_step(
                pool, run_id, step.name, "running",
                inputs=step.inputs,
            )

            if step.type == "confirm":
                # Pause the workflow
                await store.update_workflow_run(
                    pool, run_id,
                    status="paused",
                    current_step=step.name,
                    state=state,
                )
                # Record step as paused
                await store.record_workflow_step(
                    pool, run_id, step.name, "paused",
                )
                return await store.get_workflow_run(pool, run_id)

            elif step.type == "skill":
                output = await _execute_skill_step(
                    step, state, config, pool, project_id, project_dir,
                )
                if step.outputs_key:
                    state["outputs"][step.outputs_key] = output

            elif step.type == "action":
                output = await _execute_action_step(
                    step, state, pool, redis, project_id,
                )
                if step.outputs_key:
                    state["outputs"][step.outputs_key] = output

            elif step.type == "condition":
                jump_to = await _execute_condition_step(step, state)
                # Record step completion
                await store.record_workflow_step(
                    pool, run_id, step.name, "completed",
                    outputs={"jump_to": jump_to},
                )
                await publish_event(
                    redis, project_id, EventType.WORKFLOW_STEP_COMPLETED,
                    payload={"run_id": run_id, "step": step.name},
                )
                # Jump to named step if specified
                if jump_to:
                    found = False
                    for i, s in enumerate(workflow.steps):
                        if s.name == jump_to:
                            step_index = i
                            found = True
                            break
                    if not found:
                        step_index += 1
                else:
                    step_index += 1
                # Update state and continue (skip the normal increment)
                await store.update_workflow_run(
                    pool, run_id,
                    current_step=step.name,
                    state=state,
                )
                continue

            # Record step completion (for non-condition types)
            await store.record_workflow_step(
                pool, run_id, step.name, "completed",
                outputs=output if isinstance(output, dict) else {},
            )

            # Publish step completed event
            await publish_event(
                redis, project_id, EventType.WORKFLOW_STEP_COMPLETED,
                payload={"run_id": run_id, "step": step.name},
            )

            # Persist state after each step
            await store.update_workflow_run(
                pool, run_id,
                current_step=step.name,
                state=state,
            )

            step_index += 1

        # All steps completed
        now = datetime.now(timezone.utc)
        await store.update_workflow_run(
            pool, run_id,
            status="completed",
            outputs=state.get("outputs", {}),
            completed_at=now,
        )
        await publish_event(
            redis, project_id, EventType.WORKFLOW_COMPLETED,
            payload={"run_id": run_id, "workflow": workflow.name},
        )
        return await store.get_workflow_run(pool, run_id)

    except Exception as exc:
        # Mark workflow as failed
        await store.update_workflow_run(
            pool, run_id,
            status="failed",
            error=str(exc),
        )
        await publish_event(
            redis, project_id, EventType.WORKFLOW_FAILED,
            payload={"run_id": run_id, "error": str(exc)},
        )
        raise


async def resume_workflow(
    run_id: str,
    pool: asyncpg.Pool,
    redis: Redis,
    config: SkillsConfig,
    project_id: str,
    project_dir: str | None = None,
    workflow: WorkflowDefinition | None = None,
) -> dict:
    """Resume a paused workflow from the step after its current_step.

    If ``workflow`` is provided it is used directly; otherwise the definition
    is loaded by name from the discovered workflow files.
    """
    # Load workflow run from DB
    run = await store.get_workflow_run(pool, run_id)

    if run["status"] != "paused":
        raise ValueError(f"Workflow run {run_id} is not paused (status: {run['status']})")

    # Load the workflow definition (use provided or discover)
    if workflow is None:
        workflow = load_workflow(run["workflow_name"], project_dir)

    # Restore state
    state = run.get("state") or {}
    if "outputs" not in state:
        state["outputs"] = {}

    # Find the step index after current_step
    current_step_name = run.get("current_step")
    start_index = 0
    if current_step_name:
        for i, s in enumerate(workflow.steps):
            if s.name == current_step_name:
                start_index = i + 1
                break

    # Update status to running
    await store.update_workflow_run(pool, run_id, status="running")

    try:
        step_index = start_index
        while step_index < len(workflow.steps):
            step = workflow.steps[step_index]

            # Record step start
            await store.record_workflow_step(
                pool, run_id, step.name, "running",
                inputs=step.inputs,
            )

            if step.type == "confirm":
                await store.update_workflow_run(
                    pool, run_id,
                    status="paused",
                    current_step=step.name,
                    state=state,
                )
                await store.record_workflow_step(
                    pool, run_id, step.name, "paused",
                )
                return await store.get_workflow_run(pool, run_id)

            elif step.type == "skill":
                output = await _execute_skill_step(
                    step, state, config, pool, project_id, project_dir,
                )
                if step.outputs_key:
                    state["outputs"][step.outputs_key] = output

            elif step.type == "action":
                output = await _execute_action_step(
                    step, state, pool, redis, project_id,
                )
                if step.outputs_key:
                    state["outputs"][step.outputs_key] = output

            elif step.type == "condition":
                jump_to = await _execute_condition_step(step, state)
                await store.record_workflow_step(
                    pool, run_id, step.name, "completed",
                    outputs={"jump_to": jump_to},
                )
                await publish_event(
                    redis, project_id, EventType.WORKFLOW_STEP_COMPLETED,
                    payload={"run_id": run_id, "step": step.name},
                )
                if jump_to:
                    found = False
                    for i, s in enumerate(workflow.steps):
                        if s.name == jump_to:
                            step_index = i
                            found = True
                            break
                    if not found:
                        step_index += 1
                else:
                    step_index += 1
                await store.update_workflow_run(
                    pool, run_id,
                    current_step=step.name,
                    state=state,
                )
                continue

            # Record step completion
            await store.record_workflow_step(
                pool, run_id, step.name, "completed",
                outputs=output if isinstance(output, dict) else {},
            )

            await publish_event(
                redis, project_id, EventType.WORKFLOW_STEP_COMPLETED,
                payload={"run_id": run_id, "step": step.name},
            )

            await store.update_workflow_run(
                pool, run_id,
                current_step=step.name,
                state=state,
            )

            step_index += 1

        # All steps completed
        now = datetime.now(timezone.utc)
        await store.update_workflow_run(
            pool, run_id,
            status="completed",
            outputs=state.get("outputs", {}),
            completed_at=now,
        )
        await publish_event(
            redis, project_id, EventType.WORKFLOW_COMPLETED,
            payload={"run_id": run_id, "workflow": workflow.name},
        )
        return await store.get_workflow_run(pool, run_id)

    except Exception as exc:
        await store.update_workflow_run(
            pool, run_id,
            status="failed",
            error=str(exc),
        )
        await publish_event(
            redis, project_id, EventType.WORKFLOW_FAILED,
            payload={"run_id": run_id, "error": str(exc)},
        )
        raise
