core/tools.py

import os
import re
import csv
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from typing import Optional, List

import pytz
from langchain.schema import Document, HumanMessage, SystemMessage
from langchain.tools import tool
from .retrievers import hybrid_search
from .validation import validate_medical_answer
from .github_storage import get_github_storage
from .context_enrichment import enrich_retrieved_documents
from .config import logger
from langchain_openai import ChatOpenAI

# Canonical provider names - used for "all providers" queries
CANONICAL_PROVIDERS = ["ASCO", "ESMO", "IASLC", "NCCN", "NICE"]

# Global configuration for medical_guidelines_knowledge_tool retrieval and enrichment
# These control the number of documents retrieved and context pages added
# When no provider is specified, these settings are applied PER PROVIDER (5 providers total)
# When a specific provider is given, these settings apply to that single provider
TOOL_K_VECTOR = 5  # Number of documents to retrieve using vector search (per provider)
TOOL_K_BM25 = 1    # Number of documents to retrieve using BM25 search (per provider)
TOOL_PAGES_BEFORE = 1  # Number of pages to include before each top result
TOOL_PAGES_AFTER = 1   # Number of pages to include after each top result
TOOL_MAX_ENRICHED = 1  # Maximum number of top documents to enrich with context (per provider)

# Global variables to store context for validation
_last_question = None  # Stores the tool query
_last_documents = None

TOOL_MAX_WORKERS = max(2, min(8, (os.cpu_count() or 4)))
_tool_executor = ThreadPoolExecutor(max_workers=TOOL_MAX_WORKERS)


# Map lowercase variants and full names to canonical provider codes
_PROVIDER_ALIASES = {
    # NCCN
    "nccn": "NCCN",
    "national comprehensive cancer network": "NCCN",
    "nccn guidelines": "NCCN",
    # ESMO
    "esmo": "ESMO",
    "european society for medical oncology": "ESMO",
    "esmo guidelines": "ESMO",
    # ASCO
    "asco": "ASCO",
    "american society of clinical oncology": "ASCO",
    "asco guidelines": "ASCO",
    # NICE
    "nice": "NICE",
    "national institute for health and care excellence": "NICE",
    "nice guidelines": "NICE",
    # IASLC
    "iaslc": "IASLC",
    "international association for the study of lung cancer": "IASLC",
    "iaslc guidelines": "IASLC",
    # Manus (custom provider)
    "manus": "Manus",
    "by manus": "Manus",
}


def _normalize_provider(provider: Optional[str], query: str) -> Optional[str]:
    """Normalize provider name from explicit parameter or query text.
    
    Handles:
    - Exact canonical matches (e.g., "NCCN", "nccn")
    - Full name aliases (e.g., "National Comprehensive Cancer Network")
    - Provider mentions in text (e.g., "according to NCCN guidelines")
    
    Args:
        provider: Explicit provider name if specified
        query: Query text that may contain provider reference
        
    Returns:
        Canonical provider name (NCCN, ASCO, ESMO, NICE, Manus) or None
    """
    # Try explicit provider first, then fall back to query text
    text = provider if provider else query
    if not text:
        return None
    
    t = text.lower()
    
    # Quick direct hits for canonical providers
    for canon in CANONICAL_PROVIDERS:
        if re.search(rf"\b{re.escape(canon.lower())}\b", t):
            return canon
    
    # Alias-based detection
    for alias, canon in _PROVIDER_ALIASES.items():
        if alias in t:
            return canon
    
    # If explicit provider didn't match, try query text as fallback
    if provider and provider != query:
        return _normalize_provider(None, query)
    
    return None


def clear_text(text: str) -> str:
    """Clean and normalize text by removing markdown and excess whitespace.
    
    Operations:
    - Convert [title](url) -> title (url)
    - Remove images ![alt](url)
    - Strip code fences, backticks, and markdown emphasis
    - Collapse multiple newlines and spaces
    - Trim whitespace
    
    Args:
        text: Raw text to clean
        
    Returns:
        Cleaned text string
    """
    if not text:
        return ""
    t = text
    # Normalize newlines
    t = t.replace("\r\n", "\n").replace("\r", "\n")
    # Links: keep title and URL
    t = re.sub(r"\[([^\]]+)\]\(([^)]+)\)", r"\1 (\2)", t)
    # Images: drop entirely
    t = re.sub(r"!\[[^\]]*\]\([^)]*\)", "", t)
    # Remove headers/quotes markers at line starts
    t = re.sub(r"(?m)^[>\s]*#{1,6}\s*", "", t)
    # Remove backticks/code fences and emphasis
    t = t.replace("```", "").replace("`", "")
    t = t.replace("**", "").replace("*", "").replace("_", "")
    # Collapse spaces before newlines
    t = re.sub(r"[ \t]+\n", "\n", t)
    # Collapse multiple newlines and spaces
    t = re.sub(r"\n{3,}", "\n\n", t)
    t = re.sub(r"[ \t]{2,}", " ", t)
    # Trim and truncate
    t = t.strip()
    return t

def _format_docs_with_citations(docs: List[Document], group_by_provider: bool = False) -> str:
    """Format documents with citations.
    
    Args:
        docs: List of documents to format
        group_by_provider: If True, group results by provider with headers
        
    Returns:
        Formatted string with document citations
    """
    if not docs:
        return "No results."
    
    if group_by_provider:
        return _format_grouped_by_provider(docs)
    
    parts = []
    for i, d in enumerate(docs, start=1):
        meta = d.metadata or {}
        citation = _build_citation(i, meta, d.page_content)
        parts.append(citation)
    
    return "\n\n".join(parts)


def _build_citation(index: int, metadata: dict, content: str, include_provider: bool = True) -> str:
    """Build a single citation string with clean formatting.
    
    Args:
        index: Result number
        metadata: Document metadata
        content: Document content
        include_provider: Whether to include provider in metadata line
        
    Returns:
        Formatted citation string
    """
    source = metadata.get("source", "unknown")
    if isinstance(source, str):
        source = os.path.splitext(source)[0]
    page = metadata.get("page_number", "?")
    provider = metadata.get("provider", "unknown")
    # Ensure source displays as "<PROVIDER> <file-stem>" unless it already starts with the provider
    if isinstance(source, str) and isinstance(provider, str) and provider and source:
        if not source.lower().startswith(provider.lower()):
            source = f"{provider} {source}"
    disease = metadata.get("disease", "unknown")
    is_context = metadata.get("context_enrichment", False)
    
    snippet = clear_text(content)
    
    # Build citation header
    citation = f"📄 Result {index}:\n"
    
    # Build metadata line
    metadata_parts = []
    if include_provider:
        metadata_parts.append(f"Provider: {provider}")
    metadata_parts.append(f"Disease: {disease}")
    metadata_parts.append(f"Source: {source}")
    metadata_parts.append(f"Page: {page}")
    
    citation += " | ".join(metadata_parts)
    
    if is_context:
        citation += " [CONTEXT PAGE]"
    
    citation += f"\n\n{snippet}\n"
    return citation


def _document_to_dict(doc: Document) -> dict:
    """Convert a Document to a dictionary for storage.
    
    Args:
        doc: Document to convert
        
    Returns:
        Dictionary with document metadata and content
    """
    return {
        "doc_id": getattr(doc, 'id', None),
        "source": doc.metadata.get("source", "unknown"),
        "provider": doc.metadata.get("provider", "unknown"),
        "page_number": doc.metadata.get("page_number", "unknown"),
        "disease": doc.metadata.get("disease", "unknown"),
        "context_enrichment": doc.metadata.get("context_enrichment", False),
        "enriched": doc.metadata.get("enriched", False),
        "pages_included": doc.metadata.get("pages_included", []),
        "primary_page": doc.metadata.get("primary_page"),
        "context_pages_before": doc.metadata.get("context_pages_before"),
        "context_pages_after": doc.metadata.get("context_pages_after"),
        "content": doc.page_content
    }


def _format_grouped_by_provider(docs: List[Document]) -> str:
    """Format results grouped by provider for multi-provider queries.
    
    Args:
        docs: List of documents to format
        
    Returns:
        Formatted string with documents grouped by provider
    """
    if not docs:
        return "No results found from any guideline provider."
    
    # Group documents by provider
    provider_groups = {}
    for doc in docs:
        provider = doc.metadata.get("provider", "unknown")
        if provider not in provider_groups:
            provider_groups[provider] = []
        provider_groups[provider].append(doc)
    
    # Format header
    parts = [
        f"\n{'='*70}",
        f"MULTI-PROVIDER SEARCH RESULTS",
        f"Retrieved information from {len(provider_groups)} guideline provider(s)",
        f"{'='*70}\n"
    ]
    
    # Format each provider's results
    for idx, provider in enumerate(sorted(provider_groups.keys()), start=1):
        provider_docs = provider_groups[provider]
        
        # Provider header
        parts.append(f"\n{'─'*70}")
        parts.append(f"🏥 PROVIDER {idx}: {provider} ({len(provider_docs)} result{'s' if len(provider_docs) != 1 else ''})")
        parts.append(f"{'─'*70}\n")
        
        # Format each document for this provider
        for i, doc in enumerate(provider_docs, start=1):
            meta = doc.metadata or {}
            # Use _build_citation but without provider in metadata (already in header)
            citation = _build_citation(i, meta, doc.page_content, include_provider=False)
            parts.append(citation)
            
            # Add separator between results (except after last one)
            if i < len(provider_docs):
                parts.append("")
    
    return "\n".join(parts)


@tool
def medical_guidelines_knowledge_tool(query: str, provider: Optional[str] = None) -> str:
    """
    Retrieve comprehensive medical guideline knowledge with enriched context.
    Includes surrounding pages (before/after) for top results to provide complete clinical context.
    
    Behavior:
    - If provider is specified (e.g., "NCCN", "ASCO", "ESMO", "NICE"): queries only that provider
    - If provider is NOT specified: queries ALL providers separately and combines results for balanced representation
    
    This ensures that results include perspectives from all guideline providers, preventing any single
    provider from dominating the results due to higher similarity scores.
    
    Returns detailed text with full metadata and contextual information for expert analysis.
    """
    global _last_question, _last_documents
    try:
        # Store question for validation context
        _last_question = query
        
        # Normalize provider name from either explicit arg or query text
        normalized_provider = _normalize_provider(provider, query)
        
        # If no specific provider is given, query ALL providers separately to ensure balanced representation
        # This prevents one provider from dominating results due to higher similarity scores
        if not normalized_provider:
            logger.info("No specific provider - querying each provider separately for balanced results")
            all_enriched_docs = []
            
            for prov in CANONICAL_PROVIDERS:
                try:
                    logger.info(f"Querying provider: {prov}")
                    docs = hybrid_search(query=query, provider=prov, k_vector=TOOL_K_VECTOR, k_bm25=TOOL_K_BM25, use_query_expansion=False)
                    
                    if docs:
                        # Enrich documents for this provider
                        enriched = enrich_retrieved_documents(
                            documents=docs,
                            pages_before=TOOL_PAGES_BEFORE,
                            pages_after=TOOL_PAGES_AFTER,
                            max_enriched=TOOL_MAX_ENRICHED
                        )
                        all_enriched_docs.extend(enriched)
                        logger.info(f"Retrieved {len(docs)} documents from {prov}")
                except Exception as e:
                    logger.warning(f"Error querying {prov}: {str(e)}")
                    # Continue with other providers even if one fails
                    continue
            
            if not all_enriched_docs:
                return "No results found from any guideline provider for this query."
            
            # Store all documents for validation
            _last_documents = [_document_to_dict(doc) for doc in all_enriched_docs]
            
            # Format results grouped by provider for better readability
            return _format_docs_with_citations(all_enriched_docs, group_by_provider=True)
        
        # Single specific provider query
        logger.info(f"Querying specific provider: {normalized_provider}")
        docs = hybrid_search(query=query, provider=normalized_provider, 
                           k_vector=TOOL_K_VECTOR, k_bm25=TOOL_K_BM25,
                           use_query_expansion=False)
        
        # Enrich top documents with surrounding pages for richer context
        enriched_docs = enrich_retrieved_documents(
            documents=docs,
            pages_before=TOOL_PAGES_BEFORE,
            pages_after=TOOL_PAGES_AFTER,
            max_enriched=TOOL_MAX_ENRICHED
        )
        
        # Count context pages added
        context_pages_count = sum(1 for doc in enriched_docs if doc.metadata.get("context_enrichment", False))
        logger.info(f"Retrieved {len(docs)} documents, added {context_pages_count} context pages")
        
        # Store documents for validation context
        _last_documents = [_document_to_dict(doc) for doc in enriched_docs]
        
        return _format_docs_with_citations(enriched_docs)
    except Exception as e:
        logger.error(f"Retrieval error: {str(e)}")
        return f"Retrieval error: {str(e)}"



@tool
def compare_providers_tool(query: str, provider_a: str, provider_b: str) -> str:
    """
    Compare guideline answers between two providers (e.g., provider_a="NCCN", provider_b="ESMO").
    Retrieves provider-filtered results independently, then returns a structured text block suited for comparison.
    Output includes citations (source file, page number, provider, disease) for each side.
    """
    try:
        canon_a = _normalize_provider(provider_a, query) or provider_a
        canon_b = _normalize_provider(provider_b, query) or provider_b

        a_future = _tool_executor.submit(hybrid_search, query, canon_a, 5, 5)
        b_future = _tool_executor.submit(hybrid_search, query, canon_b, 5, 5)

        a_docs = a_future.result()
        b_docs = b_future.result()

        format_a_future = _tool_executor.submit(_format_docs_with_citations, a_docs)
        format_b_future = _tool_executor.submit(_format_docs_with_citations, b_docs)

        a_text = format_a_future.result()
        b_text = format_b_future.result()

        return (
            f"Comparison for query: {query}\n\n"
            f"Provider A: {canon_a}\n"
            f"{'-'*40}\n"
            f"{a_text}\n\n"
            f"Provider B: {canon_b}\n"
            f"{'-'*40}\n"
            f"{b_text}\n"
        )
    except Exception as e:
        return f"Comparison retrieval error: {str(e)}"



@tool
def get_current_datetime_tool() -> str:
    """
    Returns the current date, time, and day of the week for Egypt (Africa/Cairo).
    This is the only reliable source for date and time information. Use this tool
    whenever a user asks about 'today', 'now', or any other time-sensitive query.
    The output is always in English and in standard 12-hour format.
    """
    try:
        # Define the timezone for Egypt
        egypt_tz = pytz.timezone('Africa/Cairo')
        
        # Get the current time in that timezone
        now_egypt = datetime.now(egypt_tz)
        
        # Manual mapping to ensure English output regardless of system locale
        days_en = {
            0: "Monday", 1: "Tuesday", 2: "Wednesday", 3: "Thursday",
            4: "Friday", 5: "Saturday", 6: "Sunday"
        }
        months_en = {
            1: "January", 2: "February", 3: "March", 4: "April",
            5: "May", 6: "June", 7: "July", 8: "August",
            9: "September", 10: "October", 11: "November", 12: "December"
        }
        
        # Get English names using manual mapping
        day_name = days_en[now_egypt.weekday()]
        month_name = months_en[now_egypt.month]
        day = now_egypt.day
        year = now_egypt.year
        
        # Format time manually to avoid locale issues
        hour = now_egypt.hour
        minute = now_egypt.minute
        
        # Convert to 12-hour format
        if hour == 0:
            hour_12 = 12
            period = "AM"
        elif hour < 12:
            hour_12 = hour
            period = "AM"
        elif hour == 12:
            hour_12 = 12
            period = "PM"
        else:
            hour_12 = hour - 12
            period = "PM"
        
        time_str = f"{hour_12:02d}:{minute:02d} {period}"
        
        # Create the final string
        return f"Current date and time in Egypt: {day_name}, {month_name} {day}, {year} at {time_str}"
    
    except Exception as e:
        return f"Error getting current datetime: {str(e)}"


def _get_llm_safe(temperature: float = 0.0, model: str = "gpt-4o") -> Optional[ChatOpenAI]:
    """Create a ChatOpenAI client if API key is available.
    
    Args:
        temperature: Model temperature (0.0 = deterministic)
        model: OpenAI model name
        
    Returns:
        ChatOpenAI instance or None if unavailable
    """
    try:
        # ChatOpenAI will read OPENAI_API_KEY from env as in validation.py
        return ChatOpenAI(model=model, temperature=temperature, max_tokens=512, request_timeout=30)
    except Exception:
        return None

def _is_side_effect_report_llm(user_input: str) -> Optional[bool]:
    """Classify if input is an adverse drug reaction report using LLM.
    
    Args:
        user_input: User's text input
        
    Returns:
        True if side effect report, False if not, None if uncertain/unavailable
    """
    llm = _get_llm_safe()
    if not llm:
        return None
    try:
        system = SystemMessage(content=(
            "You are a medical triage classifier. Decide if the user's text is a report of an adverse drug reaction (side effect) about a medication.\n"
            "Criteria: mentions a medication/drug and symptoms or adverse effects experienced by a patient.\n"
            "Respond with exactly one token: yes or no."
        ))
        human = HumanMessage(content=user_input[:1500])
        resp = llm.invoke([system, human])
        ans = (resp.content or "").strip().lower()
        if ans.startswith("yes"):
            return True
        if ans.startswith("no"):
            return False
        return None
    except Exception:
        return None


@tool
def side_effect_recording_tool(user_input: str) -> str:
    """
    Detects when a doctor reports or mentions discovering a side effect related to a drug.
    First asks for missing critical information (drug name, side effects) and optional details
    (patient_age, patient_gender, dosage, duration, severity). If user cannot provide optional
    information, saves the report with NaN values for unknown data.
    
    This tool should be used when the input contains:
    - Reports of adverse drug reactions or side effects
    - Patient experiencing unexpected symptoms after medication
    - Drug-related complications or adverse events
    - Medical professionals reporting medication issues
    
    Args:
        user_input (str): The doctor's input describing the side effect or adverse reaction
        
    Returns:
        str: Interactive form for collecting missing information or confirmation of data recording
    """
    try:
        # LLM classification (preferred), with keyword fallback to preserve behavior
        side_effect_keywords = [
            'side effect', 'adverse reaction', 'adverse event', 'drug reaction',
            'medication reaction', 'allergic reaction', 'complication', 'toxicity',
            'intolerance', 'hypersensitivity', 'contraindication', 'withdrawal',
            'overdose', 'poisoning', 'drug-induced', 'medication-induced',
            'experienced after taking', 'developed after', 'caused by medication',
            'drug-related', 'medication-related', 'pharmaceutical reaction',
            'kidney problems', 'liver problems', 'heart problems', 'breathing problems',
            'skin problems', 'stomach problems', 'nausea', 'vomiting', 'diarrhea',
            'headache', 'dizziness', 'fatigue', 'weakness', 'rash', 'swelling',
            'pain', 'fever', 'cough', 'infection', 'bleeding', 'bruising',
            'has these', 'has serious', 'causes', 'resulted in', 'led to',
            'problems with', 'issues with', 'complications from'
        ]
        input_lower = user_input.lower().strip()
        llm_decision = _is_side_effect_report_llm(user_input)
        
        # Check for special commands first
        if input_lower in ['save report', 'save', 'submit report', 'submit']:
            # Create minimal data for saving
            extracted_data = _extract_side_effect_data(user_input)
            return _save_side_effect_report(extracted_data)
        
        if input_lower in ['cancel', 'cancel report', 'abort']:
            return "**Side Effect Report Cancelled**\n\nThe adverse drug reaction report has been cancelled and no data was saved."
        
        # Check if this is a follow-up with additional information or user saying they can't provide info
        if _is_followup_response(user_input):
            # For follow-up responses, we need to get the base data from somewhere
            # Since we don't have session state, treat this as a new report
            extracted_data = _extract_side_effect_data(user_input)
            return _process_followup_response(user_input, extracted_data)
        
        # Combine LLM decision with keyword fallback to avoid behavior regression
        keyword_detected = any(keyword in input_lower for keyword in side_effect_keywords)
        contains_side_effect = (llm_decision is True) or (llm_decision is not False and keyword_detected)
        
        if not contains_side_effect:
            return "This input does not appear to contain a side effect report. If you are reporting an adverse drug reaction, please include specific details about the medication and symptoms."
        
        # Extract information using pattern matching and keyword analysis
        extracted_data = _extract_side_effect_data(user_input)
        
        # Check if we have the critical information (drug name and side effects)
        missing_critical = _identify_missing_information(extracted_data)
        
        if missing_critical:
            # Missing critical info, ask for it first
            return _generate_information_request(extracted_data, missing_critical)
        else:
            # Have critical info, now ask for optional information
            missing_optional = _identify_missing_optional_information(extracted_data)
            if missing_optional:
                return _generate_optional_information_request(extracted_data, missing_optional)
            else:
                # Have all available info, save the report
                return _save_side_effect_report(extracted_data)
        
    except Exception as e:
        return f"Error processing side effect report: {str(e)}. Please ensure your report includes drug name and symptoms."


def _is_followup_response(user_input: str) -> bool:
    """Check if the input appears to be a follow-up response with additional information."""
    followup_indicators = [
        'patient age:', 'age:', 'gender:', 'dosage:', 'dose:', 'duration:', 
        'severity:', 'outcome:', 'additional:', 'reporter:', 'notes:',
        'male', 'female', 'years old', 'mg', 'ml', 'tablets', 'capsules',
        'mild', 'moderate', 'severe', 'recovered', 'ongoing', 'hospitalized',
        # Add indicators for when user can't provide info
        "can't provide", "cannot provide", "don't have", "do not have", 
        "not available", "unavailable", "missing", "no information",
        "just save", "save them", "save it", "save anyway", "will not provide",
        "won't provide", "don't know", "unknown", "not sure"
    ]
    input_lower = user_input.lower()
    return any(indicator in input_lower for indicator in followup_indicators)


def _process_followup_response(user_input: str, base_data: dict) -> str:
    """Process follow-up response and update the extracted data."""
    
    # Check if user is indicating they can't provide information
    cant_provide_indicators = [
        "can't provide", "cannot provide", "don't have", "do not have", 
        "not available", "unknown", "unavailable", "missing", "no information",
        "will not provide", "won't provide", "save anyway", "just save"
    ]
    
    input_lower = user_input.lower()
    if any(indicator in input_lower for indicator in cant_provide_indicators):
        # User can't provide additional info, save with what we have
        return _save_side_effect_report(base_data)
    
    # Extract additional information from the follow-up
    additional_data = _extract_side_effect_data(user_input)
    
    # Merge with base data, prioritizing new information
    merged_data = base_data.copy()
    for key, value in additional_data.items():
        if value and value != 'NaN' and str(value).strip():
            merged_data[key] = value
    
    # Check if there are still critical missing fields (only drug_name and side_effects are truly critical)
    critical_missing = []
    truly_critical_fields = ['drug_name', 'side_effects']
    
    for field in truly_critical_fields:
        value = merged_data.get(field, '')
        if not value or value == 'NaN' or not str(value).strip():
            critical_missing.append(field)
    
    # If critical information is missing, ask for it
    if critical_missing:
        return _generate_information_request(merged_data, [(field, field.replace('_', ' ').title()) for field in critical_missing])
    
    # Always save automatically after processing follow-up information
    # This ensures we save after any follow-up response, whether complete or partial
    return _save_side_effect_report(merged_data)


def _identify_missing_information(extracted_data: dict) -> list:
    """Identify which critical information is missing from the extracted data."""
    missing = []
    
    # Only truly critical fields - drug name and side effects
    critical_fields = {
        'drug_name': 'Drug/Medication Name',
        'side_effects': 'Side Effects/Symptoms'
    }
    
    for field, display_name in critical_fields.items():
        value = extracted_data.get(field, '')
        if not value or value == 'NaN' or not value.strip():
            missing.append((field, display_name))
    
    return missing


def _identify_missing_optional_information(extracted_data: dict) -> list:
    """Identify which optional information is missing from the extracted data."""
    missing = []
    
    # Optional fields that we should ask for
    optional_fields = {
        'patient_age': 'Patient Age',
        'patient_gender': 'Patient Gender', 
        'dosage': 'Medication Dosage',
        'duration': 'Treatment Duration',
        'severity': 'Severity Level'
    }
    
    for field, display_name in optional_fields.items():
        value = extracted_data.get(field, '')
        if not value or value == 'NaN' or not value.strip():
            missing.append((field, display_name))
    
    return missing


def _generate_information_request(extracted_data: dict, missing_info: list) -> str:
    """Generate a medical-professional request for missing critical information."""
    
    # Only ask for truly critical missing information
    critical_missing = []
    for field, display_name in missing_info:
        if field in ['drug_name', 'side_effects']:
            critical_missing.append((field, display_name))
    
    if not critical_missing:
        # No critical info missing, save the report
        return _save_side_effect_report(extracted_data)
    
    # Create a concise request for only critical missing information
    response = "**Adverse Drug Reaction Report**\n\n"
    
    if any(field == 'drug_name' for field, _ in critical_missing):
        response += "Please specify the **medication/drug name** involved in this adverse reaction.\n\n"
    
    if any(field == 'side_effects' for field, _ in critical_missing):
        response += "Please describe the **side effects or symptoms** experienced.\n\n"
    
    response += "**Note**: All other details (age, gender, dosage, etc.) are optional. If you cannot provide them, I'll save the report with the available information."
    
    return response.strip()


def _generate_optional_information_request(extracted_data: dict, missing_optional: list) -> str:
    """Generate a request for optional information that would enhance the side effect report."""
    
    # Show what we already have
    response = "**Adverse Drug Reaction Report**\n\n"
    response += "**Recorded Information:**\n"
    
    if extracted_data.get('drug_name') and extracted_data['drug_name'] != 'NaN':
        response += f"- **Drug:** {extracted_data['drug_name']}\n"
    
    if extracted_data.get('side_effects') and extracted_data['side_effects'] != 'NaN':
        response += f"- **Side Effects:** {extracted_data['side_effects']}\n"
    
    response += "\n**Additional Information (Optional):**\n"
    response += "To enhance this report, please provide any of the following details if available:\n\n"
    
    for field, display_name in missing_optional:
        if field == 'patient_age':
            response += "- **Patient Age:** (e.g., 45 years old)\n"
        elif field == 'patient_gender':
            response += "- **Patient Gender:** (Male/Female)\n"
        elif field == 'dosage':
            response += "- **Dosage:** (e.g., 10mg daily, 2 tablets)\n"
        elif field == 'duration':
            response += "- **Duration:** (e.g., 3 months, 2 weeks)\n"
        elif field == 'severity':
            response += "- **Severity:** (Mild/Moderate/Severe)\n"
    
    response += "\n**Note:** If you don't have this information or cannot provide it, just reply with \"I don't have that information\" or \"save anyway\" and I'll save the report with the available data."
    
    return response.strip()


def _save_side_effect_report(extracted_data: dict) -> str:
    """Save the side effect report to CSV file."""
    try:
        # Ensure all fields have values (use 'NaN' for empty fields)
        fieldnames = [
            'timestamp', 'drug_name', 'side_effects', 'patient_age', 
            'patient_gender', 'dosage', 'duration', 'severity', 
            'outcome', 'additional_details', 'reporter_info', 'raw_input'
        ]
        
        # Fill missing fields with 'NaN' and ensure proper data types
        for field in fieldnames:
            value = extracted_data.get(field, '')
            if not value or value == '' or not str(value).strip():
                extracted_data[field] = 'NaN'
            else:
                # Ensure the value is properly formatted
                extracted_data[field] = str(value).strip()
        
        # Save to GitHub repository (fallback to local if needed)
        github_storage = get_github_storage()
        success = github_storage.save_side_effects_report(extracted_data)
        if not success:
            csv_filename = "side_effects_reports.csv"
            csv_path = os.path.join(os.getcwd(), csv_filename)
            file_exists = os.path.exists(csv_path)
            with open(csv_path, 'a', newline='', encoding='utf-8') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                if not file_exists:
                    writer.writeheader()
                writer.writerow(extracted_data)
            storage_location = "locally to side_effects_reports.csv (GitHub upload failed)"
        else:
            storage_location = "to GitHub cloud repository"

        # Generate confirmation message
        drug_name = extracted_data.get('drug_name', 'NaN')
        side_effects = extracted_data.get('side_effects', 'NaN')
        report_id = extracted_data['timestamp'].replace(':', '').replace('-', '').replace(' ', '_')
        
        # Create a summary of provided vs missing information
        provided_info = []
        missing_info = []
        
        info_fields = {
            'drug_name': 'Drug/Medication',
            'side_effects': 'Side Effects', 
            'patient_age': 'Patient Age',
            'patient_gender': 'Patient Gender',
            'dosage': 'Dosage',
            'duration': 'Duration',
            'severity': 'Severity',
            'outcome': 'Outcome'
        }
        
        for field, display_name in info_fields.items():
            value = extracted_data.get(field, 'NaN')
            if value and value != 'NaN':
                provided_info.append(f"- **{display_name}:** {value}")
            else:
                missing_info.append(display_name)
        
        confirmation = f"""
**✅ Adverse Drug Reaction Report Saved**

**Report ID:** {report_id}

**Documented Information:**
{chr(10).join(provided_info) if provided_info else '- Basic side effect report recorded'}

**Pharmacovigilance Status:** Report successfully saved {storage_location} for regulatory review.

**Clinical Recommendations:**
- Monitor patient for symptom progression
- Consider dose adjustment or alternative therapy if appropriate
- Document in patient medical record
- Report serious reactions to pharmacovigilance authorities

How can I assist you further with clinical guidance for this case?
"""
        
        return confirmation.strip()
        
    except Exception as e:
        return f"Error saving side effect report: {str(e)}"


def _extract_side_effect_data_with_llm(user_input: str) -> dict:
    """
    Extract structured data from side effect report text using LLM-based extraction.
    
    Args:
        user_input (str): Raw input text containing side effect report
        
    Returns:
        dict: Structured data extracted from the input
    """
    import json
    # Get current timestamp
    egypt_tz = pytz.timezone('Africa/Cairo')
    current_time = datetime.now(egypt_tz).strftime('%Y-%m-%d %H:%M:%S')
    # Initialize extracted data with defaults
    extracted_data = {
        'timestamp': current_time,
        'drug_name': 'NaN',
        'side_effects': 'NaN',
        'patient_age': 'NaN',
        'patient_gender': 'NaN',
        'dosage': 'NaN',
        'duration': 'NaN',
        'severity': 'NaN',
        'outcome': 'NaN',
        'additional_details': 'NaN',
        'reporter_info': 'NaN',
        'raw_input': user_input[:500]
    }
    llm = _get_llm_safe()
    if llm:
        try:
            system = SystemMessage(content=(
                "Extract medical side effect information. Return ONLY a JSON object with these exact fields: "
                "drug_name, side_effects, patient_age, patient_gender, dosage, duration, severity, outcome. "
                "If missing/unclear, use 'NaN'."
            ))
            human = HumanMessage(content=user_input[:2000])
            response = llm.invoke([system, human])
            text = (response.content or "").strip()
            # Try parse; if fails, fallback regex
            try:
                extracted_json = json.loads(text)
            except json.JSONDecodeError:
                extracted_json = _extract_with_improved_regex(user_input)
        except Exception:
            extracted_json = _extract_with_improved_regex(user_input)
    else:
        extracted_json = _extract_with_improved_regex(user_input)
    # Update extracted_data
    for key, value in extracted_json.items():
        if key in extracted_data and value and str(value).strip() and str(value).strip().lower() != 'nan':
            extracted_data[key] = str(value).strip()
    return extracted_data


def _extract_with_improved_regex(user_input: str) -> dict:
    """
    Improved regex-based extraction with better duration handling.
    """
    extracted = {
        'drug_name': 'NaN',
        'side_effects': 'NaN',
        'patient_age': 'NaN',
        'patient_gender': 'NaN',
        'dosage': 'NaN',
        'duration': 'NaN',
        'severity': 'NaN',
        'outcome': 'NaN'
    }
    
    input_lower = user_input.lower()
    
    # Extract drug names with improved patterns
    drug_patterns = [
        r'\b(afinitor|cisplatin|afatinib|imatinib|dasatinib|nilotinib|bosutinib|ponatinib|bevacizumab|cetuximab|trastuzumab)\b',
        r'(?:found that|that)\s+([A-Za-z]{4,20})\s+(?:has|have)',
        r'([A-Za-z]{4,20})\s+(?:has|have)\s+(?:these\s+)?(?:side\s+effects?|adverse\s+effects?)',
        r'(?:taking|prescribed|given|on)\s+([A-Za-z][A-Za-z0-9\s\-]{2,20}?)(?:\s+(?:mg|mcg|g|ml)|\s+for|\.|,)',
        r'(?:side effects?|adverse effects?)\s+(?:of|from)\s+([A-Za-z][A-Za-z0-9\s\-]{2,20}?)(?:\s|,|\.|;)'
    ]
    
    for pattern in drug_patterns:
        matches = re.findall(pattern, user_input, re.IGNORECASE)
        if matches:
            drug_name = matches[0].strip()
            if len(drug_name) > 2 and not drug_name.lower() in ['that', 'these', 'those', 'found', 'have']:
                extracted['drug_name'] = drug_name
                break
    
    # Extract side effects
    symptom_patterns = [
        r'(?:side effects?|symptoms?|adverse effects?)\s*[:\-]?\s*([^.!?]+?)(?:\.|!|\?|patient|$)',
        r'(?:has|have)\s+(?:these\s+)?(?:side effects?[:\s]+)?([A-Za-z][^.!?]*?)(?:\.|!|\?|patient|$)',
        r'(?:experienced|developed|suffered|had)\s+([^.!?]+?)(?:\.|!|\?|after|following|$)'
    ]
    
    for pattern in symptom_patterns:
        matches = re.findall(pattern, user_input, re.IGNORECASE)
        if matches:
            symptoms = matches[0].strip()
            if len(symptoms) > 3:
                extracted['side_effects'] = symptoms
                break
    
    # Extract patient age with better patterns
    age_patterns = [
        r'patient\'?s?\s+age\s*[:\-]?\s*(\d{1,3})',
        r'age\s*[:\-]?\s*(\d{1,3})',
        r'(\d{1,3})\s*(?:years?\s+old|y/?o)',
        r'aged\s+(\d{1,3})'
    ]
    
    for pattern in age_patterns:
        matches = re.findall(pattern, user_input, re.IGNORECASE)
        if matches:
            age = int(matches[0])
            if 0 <= age <= 120:
                extracted['patient_age'] = str(age)
                break
    
    # Extract patient gender
    if re.search(r'\b(?:male|man|boy|gentleman|he|his|him)\b', input_lower):
        extracted['patient_gender'] = 'Male'
    elif re.search(r'\b(?:female|woman|girl|lady|she|her)\b', input_lower):
        extracted['patient_gender'] = 'Female'
    
    # Extract dosage
    dosage_patterns = [
        r'(?:medication\s+)?dosage\s*[:\-]?\s*([\d\.]+\s*(?:mg|mcg|g|ml|units?|tablets?|capsules?))',
        r'dosage\s*[:\-]?\s*([\d\.]+\s*ml)',
        r'(\d+(?:\.\d+)?\s*(?:mg|mcg|g|ml|units?|tablets?|capsules?))'
    ]
    
    for pattern in dosage_patterns:
        matches = re.findall(pattern, user_input, re.IGNORECASE)
        if matches:
            extracted['dosage'] = matches[0].strip()
            break
    
    # Extract duration with improved patterns
    duration_patterns = [
        r'treatment\s+duration\s*[:\-]?\s*(\d+\s*(?:days?|weeks?|months?|years?))',
        r'duration\s*[:\-]?\s*(\d+\s*(?:days?|weeks?|months?|years?))',
        r'(?:for|over|during)\s+(\d+\s*(?:days?|weeks?|months?|years?))',
        r'(\d+\s*(?:days?|weeks?|months?|years?))\s+(?:of\s+)?(?:treatment|therapy)',
        r'(?:lasted|continuing for|ongoing for)\s+(\d+\s*(?:days?|weeks?|months?|years?))'
    ]
    
    for pattern in duration_patterns:
        matches = re.findall(pattern, user_input, re.IGNORECASE)
        if matches:
            extracted['duration'] = matches[0].strip()
            break
    
    # Extract severity
    severity_keywords = {
        'mild': ['mild', 'slight', 'minor', 'light'],
        'moderate': ['moderate', 'medium', 'noticeable'],
        'severe': ['severe', 'serious', 'major', 'significant', 'intense', 'extreme']
    }
    
    for severity, keywords in severity_keywords.items():
        if any(keyword in input_lower for keyword in keywords):
            extracted['severity'] = severity.capitalize()
            break
    
    # Extract outcome
    outcome_keywords = {
        'recovered': ['recovered', 'resolved', 'better', 'improved'],
        'ongoing': ['ongoing', 'continuing', 'persistent', 'current status: ongoing'],
        'worsened': ['worsened', 'deteriorated', 'worse'],
        'hospitalized': ['hospitalized', 'admitted', 'emergency']
    }
    
    for outcome, keywords in outcome_keywords.items():
        if any(keyword in input_lower for keyword in keywords):
            extracted['outcome'] = outcome.capitalize()
            break
    
    return extracted


# Alias for backward compatibility
_extract_side_effect_data = _extract_side_effect_data_with_llm