# services/trial_manager.py
"""
Clinical Trial Data Management Service

This module is responsible for retrieving, processing, and managing clinical trial data.
It handles fetching trial data from ClinicalTrials.gov, identifying and structuring criteria,
and persisting the processed data.

The trial management process follows these steps:
1. Fetch raw trial data from ClinicalTrials.gov API
2. Extract and normalize inclusion, exclusion, and miscellaneous criteria
3. Process trial data through identification and logical structuring pipelines
4. Store the processed data as JSON files

Functions:
    get_extra_criteria: Extract additional criteria from the eligibility module.
    convert_std_ages_to_numerical_ages: Convert standardized age groups to numerical ranges.
    get_trial_data: Retrieve raw trial data from ClinicalTrials.gov API.
    remove_pesky_slash: Clean backslashes from text content.
    process_trial: Process a trial through the entire pipeline.
"""

import logging
import os
import re

import rich

from src.models.identified_criteria import IdentifiedTrial, RawTrialData
from src.models.logical_criteria import LogicalTrial
from src.repositories.trial_repository import export_pydantic_to_json
from src.services.identifier import identify_criterions_from_rawTrial
from src.services.logical_structurizer import logically_structurize_trial
from src.utils.config import DEFAULT_OUTPUT_DIR
from src.utils.helpers import curl_with_status_check

# Configure logging
logger = logging.getLogger(__name__)


def remove_pesky_slash(text: str) -> str:
    """
    Remove backslashes from text to normalize the content.

    Args:
        text (str): The text to clean.

    Returns:
        str: The cleaned text without backslashes.
    """
    return re.sub(r"\\", "", text)


def get_extra_criteria(eligibility_module: dict) -> list[str]:
    """
    Process the eligibility module to extract additional criteria.

    Args:
        eligibility_module (dict): The eligibility module dictionary from ClinicalTrials.gov.

    Returns:
        list[str]: A list of additional criteria extracted from the eligibility module.
    """
    criteria: list[str] = []

    # Extract key eligibility fields
    healthy_volunteers = eligibility_module.get("healthyVolunteers")
    sex = eligibility_module.get("sex")
    minimum_age = eligibility_module.get("minimumAge")
    maximum_age = eligibility_module.get("maximumAge")
    std_ages = eligibility_module.get("stdAges")

    # Process healthy volunteers information
    if healthy_volunteers is not None:
        if healthy_volunteers == "false":
            criteria.append("No healthy volunteers allowed")
        else:
            criteria.append("Healthy volunteers allowed")

    # Process sex requirements
    if sex is not None and sex != "ALL":
        criteria.append(f"Must be {sex}")

    # Process age requirements
    if minimum_age is not None:
        criteria.append(f"Must have minimum age of {minimum_age}")

    if maximum_age is not None:
        criteria.append(f"Must have maximum age of {maximum_age}")

    # Handle standardized age groups if specific age limits aren't provided
    if std_ages and not minimum_age and not maximum_age:
        convert_std_ages_to_numerical_ages(std_ages, criteria)

    return criteria


def convert_std_ages_to_numerical_ages(std_ages: list, criteria: list[str]) -> None:
    """
    Convert standardized age groups to numerical age ranges and add to criteria.

    Args:
        std_ages (list): List of standardized age groups from ClinicalTrials.gov.
        criteria (list[str]): List to append the converted age criteria to.
    """
    min_age = 100
    max_age = 0

    # Define age group mappings
    age_groups = {"CHILD": (0, 17), "ADULT": (18, 64), "OLDER_ADULT": (65, 100)}

    # Determine the minimum and maximum ages based on std_ages
    for age_group in std_ages:
        if age_group in age_groups:
            min_age = min(min_age, age_groups[age_group][0])
            max_age = max(max_age, age_groups[age_group][1])

    # Add age criteria if they are restrictive
    if min_age != 0:
        criteria.append(f"Must be {min_age} or older")
    if max_age != 100:
        criteria.append(f"Must be {max_age} or younger")


def get_trial_data(nct_id: str) -> RawTrialData:
    """
    Retrieve trial data from ClinicalTrials.gov API and format it for processing.

    Args:
        nct_id (str): The NCT ID of the clinical trial.

    Returns:
        RawTrialData: Structured raw trial data.

    Raises:
        ValueError: If the trial data cannot be retrieved or processed.
    """
    logger.info("Fetching trial data for NCT ID: %s", nct_id)
    try:
        # Request data from ClinicalTrials.gov API
        url = f"https://clinicaltrials.gov/api/v2/studies/{nct_id}?fields=NCTId,OfficialTitle,EligibilityModule"
        data = curl_with_status_check(url)

        # Extract study data from the response
        study = data.get("studies", [{}])[0]
        if not study:
            study = data.get("protocolSection", None)
            if not study:
                logger.error("No data found for NCT ID: %s", nct_id)
                logger.debug("Response data: %s", data)
                raise ValueError(f"No data found for NCT ID: {nct_id}")

        # Extract key fields from the study data
        official_title = study.get("identificationModule", {}).get("officialTitle", "")
        eligibility_module = study.get("eligibilityModule", {})
        eligibility = remove_pesky_slash(
            eligibility_module.get("eligibilityCriteria", "")
        )
        extra_criteria = "\n".join(get_extra_criteria(eligibility_module))

        # Split the eligibility text into sections
        inclusion_pos = eligibility.find("Inclusion Criteria:")
        exclusion_pos = eligibility.find("Exclusion Criteria:")

        # Process the eligibility text based on the presence of section markers
        if inclusion_pos != -1 and exclusion_pos != -1:
            inclusion_text = eligibility[
                inclusion_pos + len("Inclusion Criteria:") : exclusion_pos
            ].strip()
            exclusion_text = eligibility[
                exclusion_pos + len("Exclusion Criteria:") :
            ].strip()
            miscellaneous_text = eligibility[:inclusion_pos].strip()

        elif inclusion_pos != -1:
            inclusion_text = eligibility[
                inclusion_pos + len("Inclusion Criteria:") :
            ].strip()
            exclusion_text = ""
            miscellaneous_text = eligibility[:inclusion_pos].strip()
        elif exclusion_pos != -1:
            inclusion_text = ""
            exclusion_text = eligibility[
                exclusion_pos + len("Exclusion Criteria:") :
            ].strip()
            miscellaneous_text = eligibility[:exclusion_pos].strip()
        else:
            inclusion_text = ""
            exclusion_text = ""
            miscellaneous_text = eligibility.strip()

        # Add extra criteria to inclusion criteria
        inclusion_text = (inclusion_text + "\n" + extra_criteria).strip()

        # Create the raw trial data object
        raw_data = RawTrialData(
            nct_id=nct_id,
            official_title=official_title,
            inclusion_criteria=inclusion_text,
            exclusion_criteria=exclusion_text,
            miscellaneous_criteria=miscellaneous_text,
        )

        logger.info("Successfully retrieved trial data.")
        logger.debug("Fully raw input: %s", data)
        logger.debug("Trial data: %s", raw_data)
        return raw_data

    except Exception as e:
        logger.error("Error fetching trial data: %s", e)
        raise ValueError(f"Error fetching trial data: {e}") from e


def process_trial(nct_id: str, folder: str = DEFAULT_OUTPUT_DIR) -> LogicalTrial:
    """
    Process a clinical trial through the complete identification and structuring pipeline.

    Args:
        nct_id (str): The NCT ID of the clinical trial.
        folder (str, optional): The output directory for storing results. Defaults to DEFAULT_OUTPUT_DIR.

    Returns:
        LogicalTrial: The trial data with criteria identified and logically structured.

    Raises:
        ValueError: If trial data processing fails at any stage.
    """
    logger.info("Starting processing for trial NCT ID: %s", nct_id)

    # Fetch raw trial data
    raw_data = get_trial_data(nct_id)
    if not raw_data:
        raise ValueError(f"Failed to fetch trial data for NCT ID: {nct_id}")

    # Identify atomic criteria from the raw trial data
    identified_trial: IdentifiedTrial = identify_criterions_from_rawTrial(raw_data)
    if logger.level <= logging.DEBUG:
        rich.print(
            identified_trial
        )  # Using rich.print for better readability in debug mode

    # Save the identified trial data
    export_pydantic_to_json(
        identified_trial,
        f"{nct_id}_identified.json",
        os.path.join(folder, "identified"),
    )

    # Structure the identified criteria into logical relationships
    logical_trial = logically_structurize_trial(identified_trial)
    if logger.level <= logging.DEBUG:
        rich.print(
            logical_trial
        )  # Using rich.print for better readability in debug mode

    # Save the logical trial data
    export_pydantic_to_json(
        logical_trial, f"{nct_id}_logical.json", os.path.join(folder, "logical")
    )

    logger.info("Trial processing complete for NCT ID: %s", nct_id)

    return logical_trial