import json
import logging
from datetime import UTC, datetime
from typing import Any

from gobby.storage.database import DatabaseProtocol

from .definitions import WorkflowState

logger = logging.getLogger(__name__)


class WorkflowStateManager:
    """
    Manages persistence of WorkflowState and Handoffs.
    """

    def __init__(self, db: DatabaseProtocol):
        self.db = db

    def get_state(self, session_id: str) -> WorkflowState | None:
        row = self.db.fetchone("SELECT * FROM workflow_states WHERE session_id = ?", (session_id,))
        if not row:
            return None

        try:
            return WorkflowState(
                session_id=row["session_id"],
                workflow_name=row["workflow_name"],
                step=row["step"],
                step_entered_at=(
                    datetime.fromisoformat(row["step_entered_at"])
                    if row["step_entered_at"]
                    else datetime.now(UTC)
                ),
                step_action_count=row["step_action_count"],
                total_action_count=row["total_action_count"],
                artifacts=json.loads(row["artifacts"]) if row["artifacts"] else {},
                observations=json.loads(row["observations"]) if row["observations"] else [],
                reflection_pending=bool(row["reflection_pending"]),
                context_injected=bool(row["context_injected"]),
                variables=json.loads(row["variables"]) if row["variables"] else {},
                task_list=json.loads(row["task_list"]) if row["task_list"] else None,
                current_task_index=row["current_task_index"],
                files_modified_this_task=row["files_modified_this_task"],
                updated_at=(
                    datetime.fromisoformat(row["updated_at"])
                    if row["updated_at"]
                    else datetime.now(UTC)
                ),
            )
        except Exception as e:
            logger.error(
                f"Failed to parse workflow state for session {session_id}: {e}", exc_info=True
            )
            return None

    def save_state(self, state: WorkflowState) -> None:
        """Upsert workflow state."""
        self.db.execute(
            """
            INSERT INTO workflow_states (
                session_id, workflow_name, step, step_entered_at,
                step_action_count, total_action_count, artifacts,
                observations, reflection_pending, context_injected, variables,
                task_list, current_task_index, files_modified_this_task,
                updated_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(session_id) DO UPDATE SET
                workflow_name = excluded.workflow_name,
                step = excluded.step,
                step_entered_at = excluded.step_entered_at,
                step_action_count = excluded.step_action_count,
                total_action_count = excluded.total_action_count,
                artifacts = excluded.artifacts,
                observations = excluded.observations,
                reflection_pending = excluded.reflection_pending,
                context_injected = excluded.context_injected,
                variables = excluded.variables,
                task_list = excluded.task_list,
                current_task_index = excluded.current_task_index,
                files_modified_this_task = excluded.files_modified_this_task,
                updated_at = excluded.updated_at
            """,
            (
                state.session_id,
                state.workflow_name,
                state.step,
                state.step_entered_at.isoformat(),
                state.step_action_count,
                state.total_action_count,
                json.dumps(state.artifacts),
                json.dumps(state.observations),
                1 if state.reflection_pending else 0,
                1 if state.context_injected else 0,
                json.dumps(state.variables),
                json.dumps(state.task_list) if state.task_list else None,
                state.current_task_index,
                state.files_modified_this_task,
                datetime.now(UTC).isoformat(),
            ),
        )

    def merge_variables(self, session_id: str, updates: dict[str, Any]) -> bool:
        """Atomically merge variable updates into existing state.

        Uses BEGIN IMMEDIATE to serialize the read-modify-write,
        preventing concurrent evaluations from clobbering each other.

        Returns:
            True if the merge succeeded, False if the session was not found.
        """
        if not updates:
            return True
        with self.db.transaction_immediate() as conn:
            row = conn.execute(
                "SELECT variables FROM workflow_states WHERE session_id = ?",
                (session_id,),
            ).fetchone()
            if not row:
                logger.warning(
                    "merge_variables: no workflow state found for session %s", session_id
                )
                return False
            current = json.loads(row["variables"]) if row["variables"] else {}
            current.update(updates)
            conn.execute(
                "UPDATE workflow_states SET variables = ?, updated_at = ? WHERE session_id = ?",
                (json.dumps(current), datetime.now(UTC).isoformat(), session_id),
            )
            return True

    def delete_state(self, session_id: str) -> None:
        """
        Clear step workflow state while preserving lifecycle variables.

        This clears workflow_name, step, and related step-tracking fields,
        but keeps the `variables` JSON which contains lifecycle workflow
        state like unlocked_tools, task_claimed, etc.

        Uses '__ended__' placeholder instead of NULL to satisfy NOT NULL constraints.
        """
        self.db.execute(
            """
            UPDATE workflow_states SET
                workflow_name = '__ended__',
                step = '__ended__',
                step_entered_at = NULL,
                step_action_count = 0,
                total_action_count = 0,
                artifacts = '{}',
                observations = '[]',
                reflection_pending = 0,
                context_injected = 0,
                task_list = NULL,
                current_task_index = 0,
                files_modified_this_task = 0,
                updated_at = ?
            WHERE session_id = ?
            """,
            (datetime.now(UTC).isoformat(), session_id),
        )
