from __future__ import annotations

import json
import os
from typing import Any

import psycopg
from temporalio import activity


@activity.defn
async def run_gumbo_langgraph_job(params: dict[str, Any]) -> dict[str, Any]:
    from gumbo.graph import run_gumbo

    conn = os.environ["LANGGRAPH_CHECKPOINT_URI"]
    state = await run_gumbo(params["object_key"], params["thread_id"], conn)
    return {"summary": state["summary"], "object_key": params["object_key"], "thread_id": params["thread_id"]}


@activity.defn
async def persist_gumbo_summary(payload: dict[str, Any]) -> None:
    dsn = os.environ["GUMBO_RESULTS_DSN"]
    conn = await psycopg.AsyncConnection.connect(dsn)
    try:
        async with conn.cursor() as cur:

            await cur.execute(
                """
                INSERT INTO gumbo_summaries (workflow_id, object_key, summary)
                VALUES (%s, %s, %s);
                """,
                (payload["workflow_id"], payload["object_key"], payload["summary"]),
            )
        await conn.commit()
    finally:
        await conn.close()


@activity.defn
async def parse_gumbo_job_logs(job_log_bytes: bytes) -> dict[str, Any]:
    return json.loads(job_log_bytes.decode("utf-8"))
