TJWaterServerBinary/app/services/scheme_management.py

from datetime import date, datetime
from typing import Any
from uuid import UUID, uuid4

from psycopg.types.json import Jsonb

from app.native.wndb.core.connection import project_connection
from app.services.time_api import parse_utc_time


def scheme_name_exists(name: str, scheme_name: str) -> bool:
    with project_connection(name) as conn, conn.cursor() as cur:
        cur.execute(
            "select exists(select 1 from analysis.runs where name = %s)",
            (scheme_name,),
        )
        row = cur.fetchone()
    return bool(row and row[0])


def store_scheme_info(
    name: str,
    scheme_name: str,
    scheme_type: str,
    username: str,
    scheme_start_time: datetime | str,
    scheme_detail: dict,
) -> UUID:
    """Create one completed, immutable analysis run."""
    return create_analysis_run(
        name=name,
        scheme_name=scheme_name,
        scheme_type=scheme_type,
        username=username,
        scheme_start_time=scheme_start_time,
        scheme_detail=scheme_detail,
        status="completed",
    )


def create_analysis_run(
    name: str,
    scheme_name: str,
    scheme_type: str,
    username: str,
    scheme_start_time: datetime | str,
    scheme_detail: dict,
    *,
    status: str = "running",
) -> UUID:
    """Create a distinct execution record; names are labels, not identities."""
    started_at = parse_utc_time(scheme_start_time, field_name="scheme_start_time")
    run_id = uuid4()
    with project_connection(name) as conn, conn.cursor() as cur:
        cur.execute(
            """
            insert into analysis.runs
                (run_id, name, run_type, created_by, created_at, started_at, status, parameters)
            values (%s, %s, %s, %s, now(), %s, %s, %s)
            """,
            (
                run_id,
                scheme_name,
                scheme_type,
                username,
                started_at,
                status,
                Jsonb(scheme_detail),
            ),
        )
    return run_id


def update_analysis_run(
    name: str,
    run_id: UUID,
    *,
    status: str,
    username: str,
    scheme_detail: dict,
) -> None:
    """Update lifecycle state and metadata for one execution identity."""
    with project_connection(name) as conn, conn.cursor() as cur:
        cur.execute(
            """
            update analysis.runs
            set created_by = %s, status = %s, parameters = %s
            where run_id = %s
            """,
            (username, status, Jsonb(scheme_detail), run_id),
        )
        if cur.rowcount != 1:
            raise LookupError(f"analysis run {run_id} does not exist")


def _run_row(row: dict[str, Any]) -> dict[str, Any]:
    parameters = row.get("parameters") if isinstance(row.get("parameters"), dict) else {}
    return {
        "run_id": row["run_id"],
        "name": row["name"],
        "run_type": row["run_type"],
        "created_by": row["created_by"],
        "created_at": row["created_at"],
        "started_at": row["started_at"],
        "status": row["status"],
        "parameters": parameters,
    }


def _list_runs(
    name: str,
    run_type: str | None = None,
    query_date: date | None = None,
) -> list[dict[str, Any]]:
    clauses: list[str] = []
    params: list[Any] = []
    if run_type:
        clauses.append("run_type = %s")
        params.append(run_type)
    if query_date is not None:
        clauses.append("created_at::date = %s")
        params.append(query_date)
    where = f"where {' and '.join(clauses)}" if clauses else ""
    with project_connection(name) as conn, conn.cursor() as cur:
        cur.execute(
            f"select run_id, name, run_type, created_by, created_at, started_at, status, parameters from analysis.runs {where} order by created_at desc",
            params,
        )
        return [_run_row(row) for row in cur.fetchall()]


def query_scheme_list(
    name: str,
    scheme_type: str | None = None,
    query_date: date | None = None,
) -> list[dict[str, Any]]:
    return _list_runs(name, scheme_type, query_date)


def _get_run_by_name(
    name: str,
    run_name: str,
    run_type: str | None = None,
) -> dict[str, Any]:
    params: list[Any] = [run_name]
    type_clause = ""
    if run_type:
        type_clause = "and run_type = %s"
        params.append(run_type)
    with project_connection(name) as conn, conn.cursor() as cur:
        cur.execute(
            f"""
            select run_id, name, run_type, created_by, created_at, started_at,
                   status, parameters
            from analysis.runs
            where name = %s {type_clause}
            order by created_at desc
            limit 1
            """,
            params,
        )
        row = cur.fetchone()
    return _run_row(row) if row else {}


def get_analysis_run(name: str, run_id: UUID) -> dict[str, Any]:
    with project_connection(name) as conn, conn.cursor() as cur:
        cur.execute(
            """
            select run_id, name, run_type, created_by, created_at, started_at,
                   status, parameters
            from analysis.runs
            where run_id = %s
            """,
            (run_id,),
        )
        row = cur.fetchone()
    return _run_row(row) if row else {}


def query_scheme_detail(
    name: str,
    scheme_name: str,
    scheme_type: str | None = None,
) -> dict[str, Any]:
    return _get_run_by_name(name, scheme_name, scheme_type)


def store_leakage_identify_result(
    name: str,
    scheme_name: str,
    network: str,
    sensor_nodes: list[str],
    result_rows: list[dict],
    node_area_map: dict[str, str],
    areas: list[dict],
    drawing_payload: dict | None = None,
    run_status: str = "completed",
    error_message: str | None = None,
) -> None:
    run = _get_run_by_name(name, scheme_name, "dma_leak_identification")
    if not run:
        raise LookupError(f"analysis run {scheme_name!r} does not exist")
    payload = {
        "network": network,
        "run_status": run_status,
        "error_message": error_message,
        "sensor_nodes": sensor_nodes,
        "rows": result_rows,
        "node_area_map": node_area_map,
        "areas": areas,
        "drawing_payload": drawing_payload or {},
    }
    with project_connection(name) as conn, conn.cursor() as cur:
        cur.execute(
            "insert into analysis.results (run_id, result_type, payload) values (%s, 'leakage_identification', %s)",
            (run["run_id"], Jsonb(payload)),
        )


def _list_typed_runs(
    name: str,
    network: str,
    run_type: str,
    query_date: date | None,
) -> list[dict[str, Any]]:
    rows = _list_runs(name, run_type, query_date)
    return [
        row
        for row in rows
        if not network or row["parameters"].get("network") in (None, network)
    ]


def _typed_run_detail(name: str, run_name: str, run_type: str) -> dict[str, Any]:
    run = _get_run_by_name(name, run_name, run_type)
    if not run:
        return {}
    with project_connection(name) as conn, conn.cursor() as cur:
        cur.execute(
            "select result_type, payload, created_at from analysis.results where run_id = %s order by created_at, result_id",
            (run["run_id"],),
        )
        results = [dict(row) for row in cur.fetchall()]
    return run | {"results": results}


def query_leakage_identify_schemes(
    name: str,
    network: str,
    scheme_type: str = "dma_leak_identification",
    query_date: date | None = None,
) -> list[dict[str, Any]]:
    return _list_typed_runs(name, network, scheme_type, query_date)


def query_leakage_identify_scheme_detail(name: str, scheme_name: str) -> dict[str, Any]:
    return _typed_run_detail(name, scheme_name, "dma_leak_identification")


def query_burst_location_schemes(
    name: str,
    network: str,
    scheme_type: str = "burst_location",
    query_date: date | None = None,
) -> list[dict[str, Any]]:
    return _list_typed_runs(name, network, scheme_type, query_date)


def query_burst_location_scheme_detail(name: str, scheme_name: str) -> dict[str, Any]:
    return _typed_run_detail(name, scheme_name, "burst_location")


def query_burst_detection_schemes(
    name: str,
    network: str,
    scheme_type: str = "burst_detection",
    query_date: date | None = None,
) -> list[dict[str, Any]]:
    return _list_typed_runs(name, network, scheme_type, query_date)


def query_burst_detection_scheme_detail(name: str, scheme_name: str) -> dict[str, Any]:
    return _typed_run_detail(name, scheme_name, "burst_detection")