apps/cli/main.py

"""Command-line interface for running assessments and exporting reports."""

import json
from datetime import datetime
from pathlib import Path

import hydra
from hydra.utils import to_absolute_path
from omegaconf import DictConfig

from sentinel.config import AppConfig, ModelConfig, ResourcePaths
from sentinel.factory import SentinelFactory
from sentinel.models import (
    ConversationResponse,
    InitialAssessment,
)
from sentinel.reporting import generate_excel_report, generate_pdf_report
from sentinel.risk_models import RISK_MODELS
from sentinel.user_input import (
    Demographics,
    FamilyMemberCancer,
    FemaleSpecific,
    Lifestyle,
    PersonalMedicalHistory,
    UserInput,
)
from sentinel.utils import load_user_file


# Color codes for terminal output
class Colors:
    """ANSI color codes for terminal output formatting."""

    HEADER = "\033[95m"
    OKBLUE = "\033[94m"
    OKCYAN = "\033[96m"
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"


def _get_input(prompt: str, optional: bool = False) -> str:
    """Get a line of input from the user.

    Args:
        prompt: Message to display to the user.
        optional: If True, allow empty input to be returned as an empty string.

    Returns:
        The raw string entered by the user (may be empty if optional).
    """
    suffix = " (optional, press Enter to skip)" if optional else ""
    return input(f"{Colors.OKCYAN}{prompt}{suffix}:{Colors.ENDC} ")


def _get_int_input(prompt: str, optional: bool = False) -> int | None:
    """Get an integer from the user.

    Args:
        prompt: Message to display to the user.
        optional: If True, allow empty input and return None.

    Returns:
        The parsed integer value, or None if optional and left empty.
    """
    while True:
        val = _get_input(prompt, optional)
        if not val and optional:
            return None
        try:
            return int(val)
        except (ValueError, TypeError):
            print(f"{Colors.WARNING}Please enter a valid number.{Colors.ENDC}")


def collect_user_input() -> UserInput:
    """Collect user profile data interactively.

    Returns:
        UserInput: Structured demographics, lifestyle, and clinical data
        assembled from CLI prompts.
    """
    print(
        f"\n{Colors.HEADER}{Colors.BOLD}=== User Information Collection ==={Colors.ENDC}"
    )
    print("Please provide the following details for your assessment.")

    # --- DEMOGRAPHICS ---
    print(f"\n{Colors.OKBLUE}{Colors.BOLD}--- Demographics ---{Colors.ENDC}")
    age = _get_int_input("Age")
    sex = _get_input("Biological Sex (e.g., Male, Female)")
    ethnicity = _get_input("Ethnicity", optional=True)
    demographics = Demographics(age=age, sex=sex, ethnicity=ethnicity)

    # --- LIFESTYLE ---
    print(f"\n{Colors.OKBLUE}{Colors.BOLD}--- Lifestyle ---{Colors.ENDC}")
    smoking_status = _get_input("Smoking Status (e.g., never, former, current)")
    smoking_pack_years = (
        _get_int_input("Smoking Pack-Years", optional=True)
        if smoking_status in ["former", "current"]
        else None
    )
    alcohol_consumption = _get_input(
        "Alcohol Consumption (e.g., none, light, moderate, heavy)"
    )
    dietary_habits = _get_input("Dietary Habits", optional=True)
    physical_activity_level = _get_input("Physical Activity Level", optional=True)
    lifestyle = Lifestyle(
        smoking_status=smoking_status,
        smoking_pack_years=smoking_pack_years,
        alcohol_consumption=alcohol_consumption,
        dietary_habits=dietary_habits,
        physical_activity_level=physical_activity_level,
    )

    # --- PERSONAL MEDICAL HISTORY ---
    print(
        f"\n{Colors.OKBLUE}{Colors.BOLD}--- Personal Medical History ---{Colors.ENDC}"
    )
    mutations = _get_input("Known genetic mutations (comma-separated)", optional=True)
    cancers = _get_input("Previous cancers (comma-separated)", optional=True)
    illnesses = _get_input(
        "Chronic illnesses (e.g., IBD, comma-separated)", optional=True
    )
    personal_medical_history = PersonalMedicalHistory(
        known_genetic_mutations=[m.strip() for m in mutations.split(",")]
        if mutations
        else [],
        previous_cancers=[c.strip() for c in cancers.split(",")] if cancers else [],
        chronic_illnesses=[i.strip() for i in illnesses.split(",")]
        if illnesses
        else [],
    )

    # --- CLINICAL OBSERVATIONS ---
    print(
        f"\n{Colors.OKBLUE}{Colors.BOLD}--- Clinical Observations / Test Results (Optional) ---{Colors.ENDC}"
    )
    clinical_observations = []
    while True:
        add_test = _get_input(
            "Add a clinical observation or test result? (y/N)"
        ).lower()
        if add_test not in ["y", "yes"]:
            break
        test_name = _get_input("Test/Observation Name")
        value = _get_input("Value")
        unit = _get_input("Unit (e.g., ng/mL, or N/A)")
        reference_range = _get_input("Reference Range", optional=True)
        date = _get_input("Date of Test (YYYY-MM-DD)", optional=True)
        clinical_observations.append(
            {
                "test_name": test_name,
                "value": value,
                "unit": unit,
                "reference_range": reference_range or None,
                "date": date or None,
            }
        )

    # --- FAMILY HISTORY ---
    print(
        f"\n{Colors.OKBLUE}{Colors.BOLD}--- Family History of Cancer ---{Colors.ENDC}"
    )
    family_history = []
    while True:
        add_relative = _get_input("Add a family member with cancer? (y/N)").lower()
        if add_relative not in ["y", "yes"]:
            break
        relative = _get_input("Relative (e.g., mother, sister)")
        cancer_type = _get_input("Cancer Type")
        age_at_diagnosis = _get_int_input("Age at Diagnosis", optional=True)
        family_history.append(
            FamilyMemberCancer(
                relative=relative,
                cancer_type=cancer_type,
                age_at_diagnosis=age_at_diagnosis,
            )
        )

    # --- FEMALE-SPECIFIC ---
    female_specific = None
    if sex.lower() == "female":
        print(
            f"\n{Colors.OKBLUE}{Colors.BOLD}--- Female-Specific Information ---{Colors.ENDC}"
        )
        age_at_first_period = _get_int_input("Age at first period", optional=True)
        age_at_menopause = _get_int_input("Age at menopause", optional=True)
        num_live_births = _get_int_input("Number of live births", optional=True)
        age_at_first_live_birth = _get_int_input(
            "Age at first live birth", optional=True
        )
        hormone_therapy_use = _get_input("Hormone therapy use", optional=True)
        female_specific = FemaleSpecific(
            age_at_first_period=age_at_first_period,
            age_at_menopause=age_at_menopause,
            num_live_births=num_live_births,
            age_at_first_live_birth=age_at_first_live_birth,
            hormone_therapy_use=hormone_therapy_use,
        )

    # --- CURRENT CONCERNS ---
    print(f"\n{Colors.OKBLUE}{Colors.BOLD}--- Current Concerns ---{Colors.ENDC}")
    current_concerns_or_symptoms = _get_input(
        "Current symptoms or health concerns", optional=True
    )

    return UserInput(
        demographics=demographics,
        lifestyle=lifestyle,
        family_history=family_history,
        personal_medical_history=personal_medical_history,
        female_specific=female_specific,
        current_concerns_or_symptoms=current_concerns_or_symptoms,
        clinical_observations=clinical_observations,
    )


def format_risk_assessment(response: InitialAssessment, dev_mode: bool = False) -> None:
    """Pretty-print an initial risk assessment payload.

    Args:
        response (InitialAssessment): Parsed result returned by the assessment
            chain.
        dev_mode (bool): Flag enabling verbose debugging output.
    """
    # In dev mode, show everything
    if dev_mode:
        print(
            f"\n{Colors.WARNING}{Colors.BOLD}--- DEV MODE: RAW MODEL OUTPUT ---{Colors.ENDC}"
        )
        # Use model_dump instead of model_dump_json for direct printing
        print(json.dumps(response.model_dump(), indent=2))
        print(
            f"\n{Colors.WARNING}{Colors.BOLD}--- DEV MODE: PARSED & VALIDATED PYDANTIC OBJECT ---{Colors.ENDC}"
        )
        if response.thinking:
            print(
                f"{Colors.OKCYAN}{Colors.BOLD}🤔 Chain of Thought (`<think>` block):{Colors.ENDC}"
            )
            print(response.thinking)
            print(f"{Colors.WARNING}{Colors.BOLD}{'-' * 30}{Colors.ENDC}")
        if response.reasoning:
            print(
                f"{Colors.OKCYAN}{Colors.BOLD}🧠 Reasoning (`<reasoning>` block):{Colors.ENDC}"
            )
            print(response.reasoning)
            print(f"{Colors.WARNING}{Colors.BOLD}{'-' * 30}{Colors.ENDC}")
        print(f"{Colors.OKCYAN}{Colors.BOLD}Full Pydantic Object:{Colors.ENDC}")

        # return
        print(
            f"\n{Colors.WARNING}{Colors.BOLD}--- DEV MODE: FORMATTED MODEL OUTPUT ---{Colors.ENDC}"
        )

    # User-friendly formatting
    print(f"\n{Colors.HEADER}{Colors.BOLD}{'=' * 60}")
    print("🏥 CANCER RISK ASSESSMENT REPORT")
    print(f"{'=' * 60}{Colors.ENDC}")

    # Display the primary user-facing response first
    if response.response:
        print(f"\n{Colors.OKCYAN}{Colors.BOLD}🤖 BiOS:{Colors.ENDC}")
        print(response.response)

    # Then display the structured summary and details
    print(f"\n{Colors.OKBLUE}{Colors.BOLD}📋 OVERALL SUMMARY{Colors.ENDC}")
    if response.overall_risk_score is not None:
        print(
            f"{Colors.OKCYAN}Overall Risk Score: {Colors.BOLD}{response.overall_risk_score}/100{Colors.ENDC}"
        )
    if response.overall_summary:
        print(f"{Colors.OKCYAN}{response.overall_summary}{Colors.ENDC}")

    # Risk assessments
    risk_assessments = response.risk_assessments
    if risk_assessments:
        print(
            f"\n{Colors.OKBLUE}{Colors.BOLD}🎯 DETAILED RISK ASSESSMENTS{Colors.ENDC}"
        )
        print(f"{Colors.OKBLUE}{'─' * 40}{Colors.ENDC}")

        for i, assessment in enumerate(risk_assessments, 1):
            cancer_type = assessment.cancer_type
            risk_level = assessment.risk_level
            explanation = assessment.explanation

            # Color code risk levels
            if risk_level is None:
                risk_color = Colors.ENDC
            elif risk_level <= 2:
                risk_color = Colors.OKGREEN
            elif risk_level == 3:
                risk_color = Colors.WARNING
            else:  # 4-5
                risk_color = Colors.FAIL

            print(f"\n{Colors.BOLD}{i}. {cancer_type.upper()}{Colors.ENDC}")
            print(
                f"   🎚️  Risk Level: {risk_color}{Colors.BOLD}{risk_level or 'N/A'}{Colors.ENDC}"
            )
            print(f"   💭 Explanation: {explanation}")

            # Optional fields
            if assessment.recommended_steps:
                print("   📝 Recommended Steps:")
                if isinstance(assessment.recommended_steps, list):
                    for step in assessment.recommended_steps:
                        print(f"      • {step}")
                else:
                    print(f"      • {assessment.recommended_steps}")

            if assessment.lifestyle_advice:
                print(f"   🌟 Lifestyle Advice: {assessment.lifestyle_advice}")

            if i < len(risk_assessments):
                print(f"   {Colors.OKBLUE}{'─' * 40}{Colors.ENDC}")

    # Diagnostic recommendations
    dx_recommendations = response.dx_recommendations
    if dx_recommendations:
        print(
            f"\n{Colors.OKBLUE}{Colors.BOLD}🔬 DIAGNOSTIC RECOMMENDATIONS{Colors.ENDC}"
        )
        print(f"{Colors.OKBLUE}{'─' * 40}{Colors.ENDC}")

        for i, dx_rec in enumerate(dx_recommendations, 1):
            test_name = dx_rec.test_name
            frequency = dx_rec.frequency
            rationale = dx_rec.rationale
            recommendation_level = dx_rec.recommendation_level

            level_text = ""
            if recommendation_level is not None:
                level_map = {
                    1: "Unsuitable",
                    2: "Unnecessary",
                    3: "Optional",
                    4: "Recommended",
                    5: "Critical - Do not skip",
                }
                level_text = f" ({level_map.get(recommendation_level, 'Unknown')})"

            print(f"\n{Colors.BOLD}{i}. {test_name.upper()}{Colors.ENDC}")
            if recommendation_level is not None:
                print(
                    f"   ⭐ Recommendation Level: {Colors.BOLD}{recommendation_level}/5{level_text}{Colors.ENDC}"
                )
            print(f"   📅 Frequency: {Colors.OKGREEN}{frequency}{Colors.ENDC}")
            print(f"   💭 Rationale: {rationale}")

            if dx_rec.applicable_guideline:
                print(f"   📜 Applicable Guideline: {dx_rec.applicable_guideline}")

            if i < len(dx_recommendations):
                print(f"   {Colors.OKBLUE}{'─' * 40}{Colors.ENDC}")

    print(
        f"\n{Colors.WARNING}⚠️  IMPORTANT: This assessment does not replace professional medical advice.{Colors.ENDC}"
    )
    print(f"{Colors.HEADER}{'=' * 60}{Colors.ENDC}")


def format_followup_response(
    response: ConversationResponse, dev_mode: bool = False
) -> None:
    """Display follow-up conversation output.

    Args:
        response (ConversationResponse): Conversation exchange returned by the
            LLM chain.
        dev_mode (bool): Flag enabling verbose debugging output.
    """
    if dev_mode:
        print(
            f"\n{Colors.WARNING}{Colors.BOLD}--- DEV MODE: RAW MODEL OUTPUT ---{Colors.ENDC}"
        )
        # Use model_dump instead of model_dump_json for direct printing
        print(json.dumps(response.model_dump(), indent=2))
        print(
            f"\n{Colors.WARNING}{Colors.BOLD}--- DEV MODE: PARSED RESPONSE ---{Colors.ENDC}"
        )
        if response.thinking:
            print(f"\n{Colors.OKCYAN}{Colors.BOLD}🤔 Chain of Thought:{Colors.ENDC}")
            print(f"{Colors.OKCYAN}{response.thinking}{Colors.ENDC}")

    print(f"\n{Colors.OKCYAN}{Colors.BOLD}🤖 BiOS:{Colors.ENDC}")
    print(f"{response.response}")


@hydra.main(config_path="../../configs", config_name="config", version_base=None)
def main(cfg: DictConfig) -> None:
    """Entry point for the CLI tool invoked via Hydra.

    Args:
        cfg (DictConfig): Hydra configuration containing model, knowledge base,
            and runtime settings.
    """
    print(
        f"{Colors.HEADER}{Colors.BOLD}Welcome to the Cancer Risk Assessment Tool{Colors.ENDC}"
    )
    print(
        f"{Colors.OKBLUE}This tool provides preliminary cancer risk assessments based on your input.{Colors.ENDC}\n"
    )

    dev_mode = cfg.dev_mode

    if dev_mode:
        print(
            f"{Colors.WARNING}🔧 Running in developer mode - raw JSON output enabled{Colors.ENDC}"
        )
    else:
        print(
            f"{Colors.OKGREEN}👤 Running in user mode - formatted output enabled{Colors.ENDC}"
        )

    model = cfg.model.model_name
    provider = cfg.model.provider
    print(f"{Colors.OKBLUE}🤖 Using model: {model} from {provider}{Colors.ENDC}")

    # Create ResourcePaths with resolved absolute paths
    knowledge_base_paths = ResourcePaths(
        persona=Path(to_absolute_path("prompts/persona/default.md")),
        instruction_assessment=Path(
            to_absolute_path("prompts/instruction/assessment.md")
        ),
        instruction_conversation=Path(
            to_absolute_path("prompts/instruction/conversation.md")
        ),
        output_format_assessment=Path(
            to_absolute_path("configs/output_format/assessment.yaml")
        ),
        output_format_conversation=Path(
            to_absolute_path("configs/output_format/conversation.yaml")
        ),
        cancer_modules_dir=Path(
            to_absolute_path("configs/knowledge_base/cancer_modules")
        ),
        dx_protocols_dir=Path(to_absolute_path("configs/knowledge_base/dx_protocols")),
    )

    # Create AppConfig from Hydra config
    app_config = AppConfig(
        model=ModelConfig(provider=cfg.model.provider, model_name=cfg.model.model_name),
        knowledge_base_paths=knowledge_base_paths,
        selected_cancer_modules=list(cfg.knowledge_base.cancer_modules),
        selected_dx_protocols=list(cfg.knowledge_base.dx_protocols),
    )

    # Create factory and conversation manager
    factory = SentinelFactory(app_config)
    conversation = factory.create_conversation_manager()

    if cfg.user_file:
        print(f"{Colors.OKBLUE}📂 Loading user data from: {cfg.user_file}{Colors.ENDC}")
        user = load_user_file(cfg.user_file)
    else:
        user = collect_user_input()

    print(f"\n{Colors.OKCYAN}🔄 Running risk scoring tools...{Colors.ENDC}")
    risks_scores = []
    for model in RISK_MODELS:
        try:
            risk_score = model().run(user)
            # Handle models that return multiple scores (e.g., QCancer)
            if isinstance(risk_score, list):
                risks_scores.extend(risk_score)
            else:
                risks_scores.append(risk_score)
        except ValueError as e:
            # Skip models that aren't applicable or have validation errors
            print(f"{Colors.WARNING}⚠️  Skipping {model().name}: {e!s}{Colors.ENDC}")
            continue

    for risk_score in risks_scores:
        # Format output based on whether cancer type is specified
        if risk_score.cancer_type and risk_score.cancer_type not in [
            "multiple",
            "Multiple Cancer Sites",
        ]:
            display = (
                f"{risk_score.name} ({risk_score.cancer_type}): {risk_score.score}"
            )
        else:
            display = f"{risk_score.name}: {risk_score.score}"
        print(f"{Colors.OKCYAN}🔄 {display}{Colors.ENDC}")

    print(f"\n{Colors.OKGREEN}🔄 Analyzing your information...{Colors.ENDC}")
    response = None
    try:
        response = conversation.initial_assessment(user, risk_scores=risks_scores)
        format_risk_assessment(response, dev_mode)
    except Exception as e:
        print(f"{Colors.FAIL}❌ Error generating assessment: {e}{Colors.ENDC}")
        return

    if response:
        export_choice = input(
            f"\n{Colors.OKCYAN}Export full report to a file? (pdf/excel/both/N):{Colors.ENDC} "
        ).lower()
        if export_choice in ["pdf", "excel", "both"]:
            output_dir = Path("outputs")
            output_dir.mkdir(exist_ok=True)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            base_filename = f"Cancer_Risk_Report_{timestamp}"

            if export_choice in ["pdf", "both"]:
                pdf_filename = output_dir / f"{base_filename}.pdf"
                try:
                    print(f"{Colors.OKCYAN}Generating PDF report...{Colors.ENDC}")
                    generate_pdf_report(response, user, str(pdf_filename))
                    print(
                        f"{Colors.OKGREEN}✅ Successfully generated {pdf_filename}{Colors.ENDC}"
                    )
                except Exception as e:
                    print(
                        f"{Colors.FAIL}❌ Error generating PDF report: {e}{Colors.ENDC}"
                    )

            if export_choice in ["excel", "both"]:
                excel_filename = output_dir / f"{base_filename}.xlsx"
                try:
                    print(f"{Colors.OKCYAN}Generating Excel report...{Colors.ENDC}")
                    generate_excel_report(response, user, str(excel_filename))
                    print(
                        f"{Colors.OKGREEN}✅ Successfully generated {excel_filename}{Colors.ENDC}"
                    )
                except Exception as e:
                    print(
                        f"{Colors.FAIL}❌ Error generating Excel report: {e}{Colors.ENDC}"
                    )

    # Follow-up conversation loop
    print(
        f"\n{Colors.OKBLUE}{Colors.BOLD}💬 You can now ask follow-up questions. Type 'quit' to exit.{Colors.ENDC}"
    )
    while True:
        q = input(f"\n{Colors.BOLD}You: {Colors.ENDC}")
        if q.lower() in {"quit", "exit", "q"}:
            print(
                f"{Colors.OKGREEN}👋 Thank you for using the Cancer Risk Assessment Tool!{Colors.ENDC}"
            )
            break

        if not q.strip():
            continue

        try:
            text = conversation.follow_up(q)
            format_followup_response(text, dev_mode)
        except Exception as e:
            print(f"{Colors.FAIL}❌ Error: {e}{Colors.ENDC}")


if __name__ == "__main__":
    main()