from dsrag_sdk.main import get_base_url, make_api_call


def create_document_via_text(knowledge_base_id: str, content: str, title: str = None, description: str = None, supp_id: str = None, auto_context_config: dict = None) -> dict:
    """
    Create a document in a knowledge base with raw text. For example, if you are adding Tweets to a knowledge base, you can just pass the text of the Tweet to this method.
    
    Args:
        knowledge_base_id (str): The ID of the knowledge base.
        content (str): The content of the document.
        title (str, optional): The title of the document. Defaults to None.
        link_to_source (str, optional): A link to the source of the document. Defaults to None.
        description (str, optional): The description of the document. Defaults to None.
        supp_id (str, optional): The ID of the document in your system. Defaults to None.
        chunk_header (str, optional): This is used to prepend context to every chunk created from this document.  Must be <= 500 characters. Defaults to None.
        auto_context (bool, optional): Whether to automatically prepend document-level context to every chunk. Cannot be True when `chunk_header` is set. ``NOTE:`` Setting this to `True` will incur an extra charge. Please see our pricing for more details. Defaults to None.
    
    Returns:
        dict: The document object.

    References:
        ``POST /knowledge_bases/{knowledge_base_id}/documents/raw_text``
    """
    data = {
        'text': content,
    }
    if title:
        data['title'] = title
    if description:
        data['description'] = description
    if supp_id:
        data['supp_id'] = supp_id
    if auto_context_config is not None:
        data['auto_context_config'] = auto_context_config
    args = {
        'method': 'POST',
        'url': f'{get_base_url()}/knowledge_bases/{knowledge_base_id}/documents/raw_text',
        'json': data,
    }
    return make_api_call(args)



def list_documents(knowledge_base_id: str, title_begins_with: str = None, link_to_source: str = None, supp_id: str = None, vectorization_status: str = None) -> list:
    """
    List the documents in a knowledge base.

    Args:
        knowledge_base_id (str): The ID of the knowledge base.
        title_begins_with (str, optional): A title prefix filter. Defaults to None.
        link_to_source (str, optional): Filter documents by their URL. Defaults to None.
        supp_id (str, optional): Filter documents by ``supp_id``. Defaults to None.

    Returns:
        list: A list of document objects.

    References:
        ``GET /knowledge_bases/{knowledge_base_id}/documents``
    """
    params = {}
    if title_begins_with:
        params['title_begins_with'] = title_begins_with
    if supp_id:
        params['supp_id'] = supp_id
    if vectorization_status:
        params['status'] = vectorization_status
    if link_to_source:
        params['link_to_source'] = link_to_source

    args = {
        'method': 'GET',
        'url': f'{get_base_url()}/knowledge_bases/{knowledge_base_id}/documents',
        'params': params,
    }
    resp = make_api_call(args)
    #documents = resp.get('documents', [])
    """while resp.get('next_page_token'):
        args['params']['next_page_token'] = resp['next_page_token']
        resp = superpowered.make_api_call(args)
        documents.extend(resp.get('documents', []))"""

    return resp


def get_document(knowledge_base_id: str, document_id: str, include_content: bool = False) -> dict:
    """
    Get an individual document from a knowledge base.

    Args:
        knowledge_base_id (str): The ID of the knowledge base.
        document_id (str): The ID of the document.
        include_content (bool, optional): Whether to include the document content in the response. If you don't plan on displaying the content, setting this parameter to ``False`` could improve latency for large documents. Defaults to True.

    Returns:
        dict: A document object.

    References:
        ``GET /knowledge_bases/{knowledge_base_id}/documents/{document_id}``
    """
    params = {
        'include_content': include_content,
    }
    args = {
        'method': 'GET',
        'url': f'{get_base_url()}/knowledge_bases/{knowledge_base_id}/documents/{document_id}',
        'params': params,
    }
    return make_api_call(args)

def delete_document(knowledge_base_id: str, document_id: str) -> bool:
    """
    Delete a document from a knowledge base.

    Args:
        knowledge_base_id (str): The ID of the knowledge base.
        document_id (str): The ID of the document.

    Returns:
        bool: Whether the document was successfully deleted.

    References:
        DELETE /knowledge_bases/{knowledge_base_id}/documents/{document_id}
    """
    args = {
        'method': 'DELETE',
        'url': f'{get_base_url()}/knowledge_bases/{knowledge_base_id}/documents/{document_id}',
    }
    return make_api_call(args)
