vent_ai_platform/app/db/repositories/diagnosis.py

from __future__ import annotations

from sqlalchemy import func, select
from sqlalchemy.orm import Session

from app.db.models.diagnosis import (
    DiagnosisEvidence,
    DiagnosisResult,
    DiagnosisRule,
    DiagnosisTask,
)
from app.db.repositories.base import create_model, update_model


def list_rules(
    db: Session,
    *,
    mine_id: str | None = None,
    scene: str | None = None,
) -> list[DiagnosisRule]:
    statement = select(DiagnosisRule).where(DiagnosisRule.enabled.is_(True))
    if mine_id:
        statement = statement.where(
            (DiagnosisRule.mine_id == mine_id) | (DiagnosisRule.mine_id.is_(None))
        )
    if scene and scene != "all":
        statement = statement.where(DiagnosisRule.scene == scene)
    return list(db.scalars(statement.order_by(DiagnosisRule.scene, DiagnosisRule.code)).all())


def get_rule(db: Session, rule_id: str) -> DiagnosisRule | None:
    return db.get(DiagnosisRule, rule_id)


def list_tasks(
    db: Session,
    mine_id: str,
    *,
    scene: str | None = None,
    status: str | None = None,
    limit: int = 100,
) -> list[DiagnosisTask]:
    statement = select(DiagnosisTask).where(DiagnosisTask.mine_id == mine_id)
    if scene and scene != "all":
        statement = statement.where(DiagnosisTask.scene == scene)
    if status:
        statement = statement.where(DiagnosisTask.status == status)
    return list(
        db.scalars(
            statement.order_by(DiagnosisTask.created_at.desc()).limit(limit)
        ).all()
    )


def get_task(db: Session, task_id: str) -> DiagnosisTask | None:
    return db.get(DiagnosisTask, task_id)


def list_results(db: Session, task_id: str) -> list[DiagnosisResult]:
    return list(
        db.scalars(
            select(DiagnosisResult)
            .where(DiagnosisResult.task_id == task_id)
            .order_by(DiagnosisResult.confidence.desc())
        ).all()
    )


def list_evidence_for_results(
    db: Session,
    result_ids: list[str],
) -> list[DiagnosisEvidence]:
    if not result_ids:
        return []
    return list(
        db.scalars(
            select(DiagnosisEvidence)
            .where(DiagnosisEvidence.result_id.in_(result_ids))
            .order_by(DiagnosisEvidence.created_at)
        ).all()
    )


def result_summary(db: Session, mine_id: str) -> dict:
    task_ids = select(DiagnosisTask.id).where(DiagnosisTask.mine_id == mine_id)
    rows = db.execute(
        select(DiagnosisResult.severity, func.count(DiagnosisResult.id))
        .where(DiagnosisResult.task_id.in_(task_ids))
        .group_by(DiagnosisResult.severity)
    ).all()
    counts = {"high": 0, "medium": 0, "low": 0}
    for severity, count in rows:
        counts[str(severity)] = int(count)
    return counts


def create_task(db: Session, data: dict) -> DiagnosisTask:
    return create_model(db, DiagnosisTask, data)


def update_task(db: Session, task: DiagnosisTask, data: dict) -> DiagnosisTask:
    return update_model(db, task, data)


def create_result(db: Session, data: dict) -> DiagnosisResult:
    return create_model(db, DiagnosisResult, data)


def create_evidence(db: Session, data: dict) -> DiagnosisEvidence:
    return create_model(db, DiagnosisEvidence, data)