"""Persistence operations for flow definitions, versions, runs, and tasks."""

from __future__ import annotations

from datetime import datetime, timezone
from uuid import UUID

from sqlalchemy.exc import IntegrityError
from sqlmodel import Session, func, select

from app.models import ApprovalTask, FlowDefinition, FlowRun, FlowVersion

from .schemas import ApprovalCreate, FlowDefinitionCreate, FlowVersionCreate


class NotFoundError(Exception):
    pass


class ConflictError(Exception):
    pass


class InvalidError(Exception):
    pass


def create_definition(session: Session, data: FlowDefinitionCreate) -> FlowDefinition:
    if data.continuation_flow_id is not None and session.get(
        FlowDefinition, data.continuation_flow_id
    ) is None:
        raise InvalidError("continuation_flow_id unknown")
    definition = FlowDefinition(
        name=data.name,
        description=data.description,
        continuation_flow_id=data.continuation_flow_id,
    )
    session.add(definition)
    try:
        session.commit()
    except IntegrityError as exc:
        session.rollback()
        raise ConflictError(f"flow {data.name!r} already exists") from exc
    session.refresh(definition)
    return definition


def get_definition(session: Session, flow_id: UUID) -> FlowDefinition:
    definition = session.get(FlowDefinition, flow_id)
    if definition is None:
        raise NotFoundError(f"flow {flow_id} not found")
    return definition


def list_definitions(session: Session) -> list[FlowDefinition]:
    return list(session.exec(select(FlowDefinition).order_by(FlowDefinition.name)).all())


def _validate_flow_json(flow_json: dict) -> None:
    if not isinstance(flow_json, dict):
        raise InvalidError("langflow_json must be an object")
    data = flow_json.get("data")
    if not isinstance(data, dict):
        raise InvalidError("langflow_json must contain a 'data' object")
    if not isinstance(data.get("nodes"), list) or not isinstance(data.get("edges"), list):
        raise InvalidError("langflow_json.data must contain 'nodes' and 'edges' lists")


def create_version(
    session: Session, flow_id: UUID, data: FlowVersionCreate
) -> FlowVersion:
    get_definition(session, flow_id)
    _validate_flow_json(data.langflow_json)
    max_no = session.exec(
        select(func.max(FlowVersion.version_no)).where(FlowVersion.flow_id == flow_id)
    ).one()
    version = FlowVersion(
        flow_id=flow_id,
        version_no=(max_no or 0) + 1,
        langflow_json=data.langflow_json,
        author=data.author,
    )
    session.add(version)
    session.commit()
    session.refresh(version)
    return version


def get_version(session: Session, version_id: UUID) -> FlowVersion:
    version = session.get(FlowVersion, version_id)
    if version is None:
        raise NotFoundError(f"flow version {version_id} not found")
    return version


def list_versions(session: Session, flow_id: UUID) -> list[FlowVersion]:
    get_definition(session, flow_id)
    return list(
        session.exec(
            select(FlowVersion)
            .where(FlowVersion.flow_id == flow_id)
            .order_by(FlowVersion.version_no)
        ).all()
    )


def get_version_by_no(session: Session, flow_id: UUID, version_no: int) -> FlowVersion:
    version = session.exec(
        select(FlowVersion).where(
            FlowVersion.flow_id == flow_id, FlowVersion.version_no == version_no
        )
    ).first()
    if version is None:
        raise NotFoundError(f"flow {flow_id} has no version {version_no}")
    return version


def get_latest_deployed_version(session: Session, flow_id: UUID) -> FlowVersion:
    version = session.exec(
        select(FlowVersion)
        .where(FlowVersion.flow_id == flow_id, FlowVersion.deployed_at.is_not(None))
        .order_by(FlowVersion.version_no.desc())
    ).first()
    if version is None:
        raise InvalidError(f"flow {flow_id} has no deployed version")
    return version


def mark_deployed(session: Session, version: FlowVersion, lfx_flow_id: str) -> FlowVersion:
    version.lfx_flow_id = lfx_flow_id
    version.deployed_at = datetime.now(timezone.utc)
    session.add(version)
    session.commit()
    session.refresh(version)
    return version


def create_run(
    session: Session,
    version: FlowVersion,
    *,
    input: dict | None,
    tenant_id: str = "default",
    session_id: str | None = None,
) -> FlowRun:
    run = FlowRun(
        version_id=version.id,
        tenant_id=tenant_id,
        input=input,
        session_id=session_id or str(version.id),
    )
    session.add(run)
    session.commit()
    session.refresh(run)
    return run


def set_run_result(
    session: Session, run: FlowRun, status: str, output: dict | None
) -> FlowRun:
    run.status = status
    run.output = output
    run.finished_at = datetime.now(timezone.utc)
    session.add(run)
    session.commit()
    session.refresh(run)
    return run


def get_run(session: Session, run_id: UUID) -> FlowRun:
    run = session.get(FlowRun, run_id)
    if run is None:
        raise NotFoundError(f"run {run_id} not found")
    return run


def list_runs(session: Session, flow_id: UUID) -> list[FlowRun]:
    get_definition(session, flow_id)
    return list(
        session.exec(
            select(FlowRun)
            .join(FlowVersion, FlowVersion.id == FlowRun.version_id)
            .where(FlowVersion.flow_id == flow_id)
            .order_by(FlowRun.created_at.desc())
        ).all()
    )


def create_approval_task(session: Session, data: ApprovalCreate) -> ApprovalTask:
    run = session.get(FlowRun, data.run_id)
    if run is None:
        raise NotFoundError(f"run {data.run_id} not found")
    if data.continuation_flow_id is not None and session.get(
        FlowDefinition, data.continuation_flow_id
    ) is None:
        raise InvalidError("continuation_flow_id unknown")
    existing = session.exec(
        select(ApprovalTask).where(
            ApprovalTask.run_id == data.run_id, ApprovalTask.status == "open"
        )
    ).first()
    if existing is not None:
        raise ConflictError(f"run {data.run_id} already has an open approval task")
    task = ApprovalTask(
        run_id=run.id,
        version_id=run.version_id,
        tenant_id=run.tenant_id,
        title=data.title,
        payload=data.payload,
        continuation_flow_id=data.continuation_flow_id,
    )
    session.add(task)
    run.status = "pending_approval"
    run.finished_at = datetime.now(timezone.utc)
    session.add(run)
    session.commit()
    session.refresh(task)
    return task


def list_approval_tasks(session: Session, tenant_id: str | None = None) -> list[ApprovalTask]:
    query = select(ApprovalTask).order_by(ApprovalTask.created_at.desc())
    if tenant_id is not None:
        query = query.where(ApprovalTask.tenant_id == tenant_id)
    return list(session.exec(query).all())