level1.py

from typing import Optional, Tuple, Literal, IO, Union
from .common import request, Result, utils, constants
import os

DateTimeTuple = Tuple[str, str]

def find(
        instrument: Literal['MSC', 'IFS', 'MCI', 'HSTDM', 'CPIC'],
        dataset: str,
        batch_id: str,        
        obs_group: Optional[str] = None,
        obs_id: Optional[str] = None,
        level0_id: Optional[str] = None,
        detector: Optional[str] = None,
        data_model: Optional[str] = None,
        obs_type: Optional[str] = None,
        filter: Optional[str] = None,
        obs_date: Optional[DateTimeTuple] = None,
        create_time: Optional[DateTimeTuple] = None,
        qc_status: Optional[int] = None,
        prc_status: Optional[int] = None,
        file_name: Optional[str] = None,
        ra: Optional[int] = None,
        dec: Optional[int] = None,
        radius: Optional[float] = None,
        object: Optional[str] = None,
        rss_id: Optional[str] = None,
        cube_id: Optional[str] = None,
        build: Optional[int] = None,
        pmapname: Optional[str] = None,
        page: int = 1,
        limit: int = 0,
        **extra_kwargs) -> Result:
    """
    根据给定的参数搜索1级数据文件记录
    
    Args:
        obs_group (Optional[str], optional): 项目ID. Defaults to None.
        obs_id (Optional[str], optional): 观测ID. Defaults to None.
        instrument (Optional[str], optional): 模块ID，如'MSC', 'IFS'. Defaults to None.
        detector (Optional[str], optional): 探测器. Defaults to None.
        data_model (Optional[str], optional): 数据类型，如'csst-msc-l1-mbi'. Defaults to None.
        obs_type (Optional[str], optional): 观测类型，如'01'. Defaults to None.
        filter (Optional[str], optional): 滤光片. Defaults to None.
        obs_date (Optional[DateTimeTuple], optional): 观测时间范围. Defaults to None.
        create_time (Optional[DateTimeTuple], optional): 创建时间范围. Defaults to None.
        qc_status (Optional[int], optional): QC1状态. Defaults to None.
        prc_status (Optional[int], optional): 处理状态. Defaults to None.
        file_name (Optional[str], optional): 文件名. Defaults to None.
        ra (Optional[int], optional): 中心赤经. Defaults to None.
        dec (Optional[int], optional): 中心赤纬. Defaults to None.
        radius (Optional[float], optional): 搜索半径. Defaults to None.
        object (Optional[str], optional): 天体名称. Defaults to None.
        rss_id (Optional[str], optional): RSS ID (IFS) Defaults to None.
        cube_id (Optional[str], optional): Cube ID (IFS). Defaults to None.
        dataset (Optional[str], optional): 数据集名称. Defaults to None.
        batch_id (Optional[str], optional): 批次ID. Defaults to None.
        build (Optional[int], optional): 构建版本. Defaults to None.
        pmapname (Optional[str], optional): CCDS pmap名. Defaults to None.
        page (int, optional): 页码. Defaults to 1.
        limit (int, optional): 每页数量. Defaults to 0.
    
    Returns:
        Result: 搜索结果对象.
    
    """

    params = {
        'obs_group': obs_group,
        'level0_id': level0_id,
        'obs_id': obs_id,
        'instrument': instrument,
        'detector': detector,
        'data_model': data_model,
        'obs_type': obs_type,
        'filter': filter,
        'qc_status': qc_status,
        'prc_status': prc_status,
        'file_name': file_name,
        'ra_cen': ra,
        'dec_cen': dec,
        'radius': radius,
        'object': object,
        'obs_date_start': None,
        'obs_date_end': None,
        'create_time_start': None,
        'create_time_end': None,
        'rss_id': rss_id,
        'cube_id': cube_id,
        'dataset': dataset,
        'batch_id': batch_id,
        'build': build,
        'pmapname': pmapname,
        'page': page,
        'limit': limit,
    }
    params.update(extra_kwargs)
    
    if obs_date is not None:
        params['obs_date_start'], params['obs_date_end'] = obs_date
        utils.is_valid_datetime_format(params['obs_date_start']) or not utils.is_valid_datetime_format(params['obs_date_end'])
    if create_time is not None:
        params['create_time_start'], params['create_time_end'] = create_time
        utils.is_valid_datetime_format(params['create_time_start']) or utils.is_valid_datetime_format(params['create_time_end'])
    
    return request.post("/api/level1", params)

def find_by_level1_id(level1_id: str) -> Result:
    """
    通过 level1 的 ID 查询1级数据
    
    Args:
        level1_id (str): 1级数据的ID
    
    Returns:
        Result: 查询结果
    
    """
    return request.get(f"/api/level1/{level1_id}")

def find_by_brick_id(brick_id: int) -> Result:
    """
    通过 brick 的 ID 查询1级数据
    
    Args:
        brick_id (int): 天区ID
    
    Returns:
        Result: 查询结果
    
    """
    return request.get(f"/api/level1/brick/{brick_id}")

def sls_find_by_qc_status(qc_status: int, batch_id: str, limit: int = 1) -> Result:
    return request.post(f"/api/level1/sls/qc_status/{qc_status}", {'limit': limit, 'batch_id': batch_id})

def update_qc_status_by_ids(ids: list[str], qc_status: int) -> Result:
    """
    根据内部_id，批量更新1级数据的QC状态
    
    Args:
        ids (List[str]): 内部_id列表
        qc_status (int): QC状态
    
    Returns:
        Result: 更新结果
    """
    return request.put("/api/level1/qc_status/batch/update", {'qc_status': qc_status, 'ids': ids})

def update_prc_status_by_ids(ids: list[str], prc_status: int) -> Result:
    """
    根据内部_id，批量更新1级数据的处理状态
    
    Args:
        ids (List[str]): 内部_id列表
        prc_status (int): 处理状态
    
    Returns:
        Result: 更新结果
    """
    return request.put("/api/level1/prc_status/batch/update", {'prc_status': prc_status, 'ids': ids})

def update_qc_status(level1_id: str, data_model: str, qc_status: int, batch_id: str) -> Result:
    """
    更新1级数据的QC状态
    
    Args:
        level1_id (str): 1级数据的ID
        data_model (str): 数据类型
        qc_status (int): QC状态
        batch_id (str): 批次ID
    
    Returns:
        Result: 更新结果
    """
    return request.put(f"/api/level1/qc_status/{level1_id}", {'data_model': data_model, 'qc_status': qc_status, 'batch_id': batch_id})

def update_prc_status(level1_id: str, data_model: str, prc_status: int, batch_id: str) -> Result:
    """
    更新1级数据的处理状态
    
    Args:
        level1_id (str): 1级数据的ID
        data_model (str): 数据类型
        prc_status (int): 处理状态
        batch_id (str): 批次ID
    
    Returns:
        Result: 操作结果
    """
    return request.put(f"/api/level1/prc_status/{level1_id}", {'data_model': data_model, 'prc_status': prc_status, 'batch_id': batch_id})

def write(local_file: Union[IO, str]) -> Result:
    '''
    将本地的1级数据文件写入到DFS中其他参数如rss_id, cube_id等，可通过extra_kwargs传入
    
    Args:
        local_file (Union[IO, str]): 文件路径或文件对象
    
    Returns:
        Result: 操作的结果对象，包含操作是否成功以及相关的错误信息，成功返回data为1级数据对象
    '''
    params = {}
    if local_file is None:
        raise ValueError("local_file is required")
    if isinstance(local_file, str):
        if not os.path.exists(local_file):
            raise FileNotFoundError(local_file)        
        return request.post_file("/api/level1/file", local_file, params)
    return request.post_bytesio("/api/level1/file", local_file, params)

def find_process(dag: Optional[str] = None,  
                dag_run: Optional[str] = None, 
                batch_id: Optional[str] = None, 
                level1_id: Optional[str] = None,
                dataset: Optional[str] = None,
                prc_module: Optional[str] = None,
                prc_status: Optional[int] = None,
                prc_date: Optional[DateTimeTuple] = None,
                page: int = 1,
                limit: int = 0) -> Result:
    """
    查询0级数据处理过程
    
    Args:
        dag (str): DAG标识
        dag_run (str): DAG运行标识
        batch_id (str): 批次ID
        level1_id (str): 1级数据的ID
        dataset (str): 数据集
        prc_module (str): 处理模块
        prc_status (int): 处理状态
        prc_date (DateTimeTuple): 处理时间范围
        page (int): 页码，默认为1
        limit (int): 每页数量 0: 不限制
    
    Returns:
        Result: 成功后，Result.data为数据列表，失败message为失败原因
    
    """
    params = {
        'dag': dag,
        'dag_run': dag_run,
        'batch_id': batch_id,
        'level1_id': level1_id,
        'dataset': dataset,
        'prc_module': prc_module,
        'prc_status': prc_status,
        'prc_date_start': None,
        'prc_date_end': None,
        'page': page,
        'limit': limit
    }
    if prc_date is not None:
        params['prc_date_start'], params['prc_date_end'] = prc_date
        if params['prc_date_start'] and utils.is_valid_datetime_format(params['prc_date_start']):
            pass
        if params['prc_date_end'] and utils.is_valid_datetime_format(params['prc_date_end']):
            pass 
    return request.post("/api/level1/process", params)

def add_process(level1_id: str, 
                dag: str, 
                dag_run: str, 
                dataset: str = constants.DEFAULT_DATASET,
                batch_id: str = constants.DEFAULT_BATCH_ID,
                prc_date: str = utils.get_current_time(),
                prc_status: int = -1024,
                prc_module: str = "",
                message: str = "") -> Result:
    """
    添加1级数据处理记录
    
    Args:
        level1_id (str): 1级数据的ID
        dag (str): DAG标识
        dag_run (str): DAG运行标识
        dataset (str): 数据集
        batch_id (str): 批次ID
        prc_date (str): 处理时间,格式为"YYYY-MM-DD HH:MM:SS"
        prc_status (int): 处理状态
        prc_module (str): 处理模块
        message (str): 处理消息
    
    Returns:
        Result: 成功后，Result.data为写入记录，失败message为失败原因
    
    """
    params = {
        'level1_id': level1_id,
        'dag': dag,
        'dag_run': dag_run,
        'dataset': dataset,
        'batch_id': batch_id,
        'prc_date': prc_date,
        'prc_status': prc_status,
        'prc_module': prc_module,
        'message': message,
    }
    utils.is_valid_datetime_format(prc_date)
    return request.post("/api/level1/prc", params)

def delete(dataset: str, batch_id: str, data_model: str) -> Result:
    """
    删除1级数据处理记录
    
    Args:
        dataset (str): 数据集
        batch_id (str): 批次ID
        data_model (str): 数据模型
    
    Returns:
        Result: 操作的结果对象，包含操作是否成功以及相关的错误信息，成功返回数据对象
    
    """
    return request.delete("/api/level1/delete", {"dataset": dataset, "batch_id": batch_id, "data_model": data_model})

def count_by_dataset(dataset: Optional[str] = None) -> Result:
    """
    查询0级数据按数据集分组的数量

    Args:
        dataset (str): 数据集

    Returns:
        Result: 成功后，Result.data为数据列表，失败message为失败原因

    """
    params = {"dataset": dataset}
    return request.post("/api/level1/count_by_dataset", params)