# Valispace Script for Custom Actions
# Purpose: Generate and update types for specified requirements based on their text
# Developed by [Paul Grey]
# Date: [Date]
#IMPORTANT: Provide your valispace url in the Configuration Section below

from typing import Any, Dict
from valispace import API

# ---------------------------------------------
# Configuration Section

VALISPACE = {
    "domain": "https://YOUR_DEPLOYMENT_NAME.valispace.com",  # Your Valispace instance URL
    "warn_https": False,  # Set to True to enable HTTPS warnings
}

# ---------------------------------------------

def initialize_api(temporary_access_token) -> API:
    """
    Initialize the Valispace API.

    Returns:
        API: Initialized Valispace API object.
    """
    return API(
        url=VALISPACE["domain"],
        session_token=temporary_access_token,
        warn_https=VALISPACE.get("warn_https", False)
    )

def get_all_types(api: API, project_id: int) -> list:
    """
    Get all requirement types for a given project ID.

    Args:
        api (API): Initialized Valispace API object.
        project_id (int): The project ID.

    Returns:
        list: List of all requirement types for the project.
    """
    types = api.request('GET', f'requirements/types/?project={project_id}')
    return [t['name'] for t in types]

def get_type_id(api: API, project_id: int, type_name: str) -> int:
    """
    Get the type ID corresponding to the given type name and project ID.

    Args:
        api (API): Initialized Valispace API object.
        project_id (int): The project ID.
        type_name (str): The name of the type to search for.

    Returns:
        int: The type ID corresponding to the given type name and project ID, or None if not found.
    """
    types = api.request('GET', f'requirements/types/?project={project_id}')
    for t in types:
        if t['name'].lower() == type_name.lower() and t['project'] == project_id:
            return t['id']
    return None

def main(**kwargs) -> Dict[str, Any]:
    """
    Main function to execute the script.

    Args:
        **kwargs: Additional arguments passed to the script.

    Returns:
        Dict[str, Any]: Result data to be sent back to Valispace.
    """
    print(kwargs)
    # Authenticate with the API using the provided temporary access token from kwargs
    api = initialize_api(kwargs['temporary_access_token'])

    # Retrieve object IDs from the kwargs
    objects_ids = kwargs.get("objects_ids", [])
    print(objects_ids)

    # Define the project ID
    project_id = api.request('GET', f'requirements/{objects_ids[0]}/')['project']
    print(project_id)
    
    # Fetch all requirement types
    requirement_types = get_all_types(api, project_id)
    types_string = ", ".join(requirement_types)

    # Construct the prompt with the fetched requirement types
    custom_prompt = f"Respond only with the type. Take into consideration the following types: {types_string}. Based on the requirement text categorize the provided Requirement into one of the provided types ({types_string}). Return only the type in the form of 'type'."
    print(custom_prompt)
    
    # Send prompt to assistant and get response
    assistant_response = api.general_prompt(
        custom_prompt=custom_prompt,
        model = 'requirement',
        field = 'text',
        objects_ids = objects_ids,
        parallel = True
    )
    print(assistant_response)

    # Process the response and update the requirements
    for req_id, category in assistant_response.get('results', {}).items():
        type_id = get_type_id(api, project_id, category)  # Get the type ID for the category
        if type_id is not None:
            api.request('PATCH', f'requirements/{req_id}/', {'type': type_id})
            print(f"Updated Requirement ID {req_id} to type '{category}'")

    # Return a success message
    return {
        "result": "Requirements updated based on categorization by the assistant."
    }

# This script is designed to be run through custom actions in Valispace.