from flask import Flask, request, jsonify
from flask_cors import CORS
import requests
import re
import logging
import os

app = Flask(__name__)
CORS(app)

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ClinicalTrials.gov API base URL
CLINICALTRIALS_API_BASE = "https://clinicaltrials.gov/api/v2/studies"

# Import LLM service (optional - will gracefully degrade if not available)
try:
    from llm_service import get_llm_service
    LLM_AVAILABLE = True
except ImportError:
    logger.warning("LLM service not available. Ranking will be disabled.")
    LLM_AVAILABLE = False
    get_llm_service = None


@app.route('/api/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({"status": "healthy"}), 200


@app.route('/api/llm-info', methods=['GET'])
def get_llm_info():
    """Get information about the LLM service"""
    try:
        if not LLM_AVAILABLE:
            return jsonify({
                "available": False,
                "model": None,
                "using_api": False
            }), 200
        
        llm_service = get_llm_service()
        if llm_service:
            return jsonify({
                "available": True,
                "model": llm_service.model_name,
                "using_api": llm_service.use_api
            }), 200
        else:
            return jsonify({
                "available": False,
                "model": None,
                "using_api": False
            }), 200
    except Exception as e:
        logger.error(f"Error getting LLM info: {str(e)}")
        return jsonify({
            "available": False,
            "model": None,
            "using_api": False,
            "error": str(e)
        }), 200


@app.route('/api/search', methods=['GET'])
def search_trials():
    """
    Search clinical trials based on query parameters
    Query params:
        - query: search keywords
        - country: filter by country
        - status: filter by recruitment status (RECRUITING, NOT_YET_RECRUITING, COMPLETED, etc.)
        - rankingTerms: optional terms to use for LLM-based ranking of results
        - page: page number (default: 1)
        - pageSize: results per page (default: 10)
    """
    try:
        # Get query parameters
        query = request.args.get('query', '')
        country = request.args.get('country', '')
        status = request.args.get('status', '')
        ranking_terms = request.args.get('rankingTerms', '')
        page = int(request.args.get('page', 1))
        page_size = int(request.args.get('pageSize', 10))
        
        # Build API request parameters
        params = {
            'format': 'json',
            'pageSize': page_size
        }
        
        # Add page token if not first page
        if page > 1:
            params['pageToken'] = str(page)
        
        # Build query filter
        query_parts = []
        if query:
            # Search in multiple fields for better results
            query_parts.append(query)
        
        if country:
            query_parts.append(f'AREA[LocationCountry]{country}')
            
        if status:
            query_parts.append(f'AREA[OverallStatus]{status}')
        
        if query_parts:
            params['query.term'] = ' AND '.join(query_parts)
        
        # Make request to ClinicalTrials.gov API
        response = requests.get(CLINICALTRIALS_API_BASE, params=params, timeout=30)
        response.raise_for_status()
        
        data = response.json()
        
        # Extract and format relevant information
        studies = []
        if 'studies' in data:
            for study in data['studies']:
                protocol_section = study.get('protocolSection', {})
                identification = protocol_section.get('identificationModule', {})
                nct_id = identification.get('nctId', '')
                organization = identification.get('organization', {})
                sponsor = organization.get('fullName', '') if organization else ''
                status_module = protocol_section.get('statusModule', {})
                description = protocol_section.get('descriptionModule', {})
                conditions = protocol_section.get('conditionsModule', {})
                locations_module = protocol_section.get('contactsLocationsModule', {})
                
                # Fetch full study details to get ALL locations and eligibility criteria
                locations = []
                inclusion_criteria = []
                exclusion_criteria = []
                detail_fetch_success = False
                detail_data = None
                if nct_id:
                    try:
                        # Use individual study endpoint to get all locations and eligibility
                        detail_url = f"https://clinicaltrials.gov/api/v2/studies/{nct_id}"
                        detail_response = requests.get(detail_url, params={'format': 'json'}, timeout=60)
                        detail_response.raise_for_status()
                        detail_data = detail_response.json()
                        detail_protocol = detail_data.get('protocolSection', {})
                        detail_locations_module = detail_protocol.get('contactsLocationsModule', {})
                        if 'locations' in detail_locations_module:
                            all_locations = detail_locations_module['locations']
                            if isinstance(all_locations, list) and len(all_locations) > 0:
                                for loc in all_locations:
                                    locations.append({
                                        'facility': loc.get('facility', ''),
                                        'city': loc.get('city', ''),
                                        'country': loc.get('country', '')
                                    })
                                detail_fetch_success = True
                    except Exception as e:
                        # Log the error for debugging
                        logger.warning(f"Failed to fetch detail data for {nct_id}: {str(e)}")
                        # Fallback to search results if detail fetch fails
                        # Detail fetch is optional - search results will be used
                        pass
                
                # Extract eligibility criteria from detail data if available
                if detail_data:
                    detail_protocol = detail_data.get('protocolSection', {})
                    eligibility_module = detail_protocol.get('eligibilityModule', {})
                    eligibility_text = eligibility_module.get('eligibilityCriteria', '')
                    
                    if eligibility_text:
                        # Parse inclusion and exclusion criteria
                        # Find inclusion and exclusion sections (case insensitive)
                        # Handle variations like:
                        # - "Inclusion Criteria:" (standard)
                        # - "Inclusion Criteria (for...):" (with parenthetical text)
                        # - "Key inclusion criteria" (no colon, different wording)
                        # - "Inclusion criteria" (no colon)
                        inclusion_patterns = [
                            r'(?:Key\s+)?Inclusion\s+Criteria[^:\n]*:?\s*(.*?)(?=(?:Key\s+)?Exclusion\s+Criteria[^:\n]*:?|$)',
                            r'Inclusion\s+Criteria[^:\n]*:?\s*(.*?)(?=Exclusion\s+Criteria[^:\n]*:?|$)',
                        ]
                        exclusion_patterns = [
                            r'(?:Key\s+)?Exclusion\s+Criteria[^:\n]*:?\s*(.*?)$',
                            r'Exclusion\s+Criteria[^:\n]*:?\s*(.*?)$',
                        ]
                        
                        inclusion_match = None
                        exclusion_match = None
                        for pattern in inclusion_patterns:
                            inclusion_match = re.search(pattern, eligibility_text, re.IGNORECASE | re.DOTALL)
                            if inclusion_match:
                                break
                        
                        for pattern in exclusion_patterns:
                            exclusion_match = re.search(pattern, eligibility_text, re.IGNORECASE | re.DOTALL)
                            if exclusion_match:
                                break
                        
                        if inclusion_match:
                            inclusion_text = inclusion_match.group(1).strip()
                            # Split by lines and clean up
                            inclusion_lines = [line.strip() for line in inclusion_text.split('\n') if line.strip()]
                            # Remove bullet points and clean
                            inclusion_criteria = [re.sub(r'^[\*\-\•]\s*', '', line).strip() for line in inclusion_lines if line.strip()]
                            # Filter out empty strings and very short lines (likely formatting artifacts)
                            inclusion_criteria = [item for item in inclusion_criteria if item and len(item) > 3]
                        else:
                            # If no explicit "Inclusion Criteria:" header, try to find criteria in the text
                            # Some studies might have different formatting
                            logger.debug(f"No inclusion criteria header found for {nct_id}, trying alternative parsing")
                        
                        if exclusion_match:
                            exclusion_text = exclusion_match.group(1).strip()
                            # Split by lines and clean up
                            exclusion_lines = [line.strip() for line in exclusion_text.split('\n') if line.strip()]
                            # Remove bullet points and clean
                            exclusion_criteria = [re.sub(r'^[\*\-\•]\s*', '', line).strip() for line in exclusion_lines if line.strip()]
                            # Filter out empty strings and very short lines
                            exclusion_criteria = [item for item in exclusion_criteria if item and len(item) > 3]
                        else:
                            logger.debug(f"No exclusion criteria header found for {nct_id}, trying alternative parsing")
                    else:
                        logger.debug(f"No eligibility text found for {nct_id}")
                else:
                    # Try to get eligibility from search results as fallback
                    eligibility_module = protocol_section.get('eligibilityModule', {})
                    eligibility_text = eligibility_module.get('eligibilityCriteria', '')
                    if eligibility_text:
                        # Same parsing logic as above
                        # Handle variations like:
                        # - "Inclusion Criteria:" (standard)
                        # - "Inclusion Criteria (for...):" (with parenthetical text)
                        # - "Key inclusion criteria" (no colon, different wording)
                        # - "Inclusion criteria" (no colon)
                        inclusion_patterns = [
                            r'(?:Key\s+)?Inclusion\s+Criteria[^:\n]*:?\s*(.*?)(?=(?:Key\s+)?Exclusion\s+Criteria[^:\n]*:?|$)',
                            r'Inclusion\s+Criteria[^:\n]*:?\s*(.*?)(?=Exclusion\s+Criteria[^:\n]*:?|$)',
                        ]
                        exclusion_patterns = [
                            r'(?:Key\s+)?Exclusion\s+Criteria[^:\n]*:?\s*(.*?)$',
                            r'Exclusion\s+Criteria[^:\n]*:?\s*(.*?)$',
                        ]
                        
                        inclusion_match = None
                        exclusion_match = None
                        for pattern in inclusion_patterns:
                            inclusion_match = re.search(pattern, eligibility_text, re.IGNORECASE | re.DOTALL)
                            if inclusion_match:
                                break
                        
                        for pattern in exclusion_patterns:
                            exclusion_match = re.search(pattern, eligibility_text, re.IGNORECASE | re.DOTALL)
                            if exclusion_match:
                                break
                        
                        if inclusion_match:
                            inclusion_text = inclusion_match.group(1).strip()
                            inclusion_lines = [line.strip() for line in inclusion_text.split('\n') if line.strip()]
                            inclusion_criteria = [re.sub(r'^[\*\-\•]\s*', '', line).strip() for line in inclusion_lines if line.strip()]
                            inclusion_criteria = [item for item in inclusion_criteria if item and len(item) > 3]
                        
                        if exclusion_match:
                            exclusion_text = exclusion_match.group(1).strip()
                            exclusion_lines = [line.strip() for line in exclusion_text.split('\n') if line.strip()]
                            exclusion_criteria = [re.sub(r'^[\*\-\•]\s*', '', line).strip() for line in exclusion_lines if line.strip()]
                            exclusion_criteria = [item for item in exclusion_criteria if item and len(item) > 3]
                
                # If detail fetch failed or returned no locations, use search results as fallback
                if len(locations) == 0 and 'locations' in locations_module:
                    all_locations = locations_module.get('locations', [])
                    if isinstance(all_locations, list):
                        for loc in all_locations:
                            locations.append({
                                'facility': loc.get('facility', ''),
                                'city': loc.get('city', ''),
                                'country': loc.get('country', '')
                            })
                
                # Sort locations: searched country first, then others
                if country and locations:
                    search_country_lower = country.lower().strip()
                    def location_sort_key(loc):
                        loc_country = loc.get('country', '').lower().strip()
                        # Exact match gets priority 0, partial match gets 1, no match gets 2
                        if loc_country == search_country_lower:
                            return 0
                        elif search_country_lower in loc_country or loc_country in search_country_lower:
                            return 1
                        else:
                            return 2
                    locations.sort(key=location_sort_key)
                
                study_info = {
                    'nctId': nct_id,
                    'title': identification.get('officialTitle') or identification.get('briefTitle', ''),
                    'sponsor': sponsor,
                    'status': status_module.get('overallStatus', ''),
                    'conditions': conditions.get('conditions', []),
                    'briefSummary': description.get('briefSummary', ''),
                    'locations': locations,  # Show all locations
                    'inclusionCriteria': inclusion_criteria,
                    'exclusionCriteria': exclusion_criteria,
                    'lastUpdateDate': status_module.get('lastUpdateSubmitDate', '')
                }
                studies.append(study_info)
        
        # Apply LLM-based ranking if ranking terms provided
        llm_model_name = None
        if ranking_terms and ranking_terms.strip() and LLM_AVAILABLE:
            try:
                llm_service = get_llm_service()
                if llm_service:
                    llm_model_name = llm_service.model_name
                    logger.info(f"Ranking {len(studies)} studies using terms: {ranking_terms}")
                    studies = llm_service.rank_studies(studies, ranking_terms)
                    logger.info("Ranking completed")
                else:
                    logger.warning("LLM service not initialized. Skipping ranking.")
            except Exception as e:
                logger.error(f"Error during LLM ranking: {str(e)}")
                # Continue without ranking if LLM fails
        
        # ClinicalTrials.gov API v2 doesn't reliably provide totalCount in the response
        # Use the number of studies found in this response if totalCount is missing or 0
        # Note: This shows the count for the current page, not the total across all pages
        api_total_count = data.get('totalCount', 0)
        total_count = api_total_count if api_total_count > 0 else len(studies)
        
        return jsonify({
            'studies': studies,
            'totalCount': total_count,
            'page': page,
            'pageSize': page_size,
            'ranked': bool(ranking_terms and ranking_terms.strip() and LLM_AVAILABLE),
            'llmModel': llm_model_name
        }), 200
        
    except requests.RequestException as e:
        return jsonify({"error": f"Error fetching data from ClinicalTrials.gov: {str(e)}"}), 500
    except Exception as e:
        return jsonify({"error": f"Internal server error: {str(e)}"}), 500


@app.route('/api/study/<nct_id>', methods=['GET'])
def get_study_details(nct_id):
    """Get detailed information for a specific study"""
    try:
        url = f"{CLINICALTRIALS_API_BASE}/{nct_id}"
        response = requests.get(url, params={'format': 'json'}, timeout=30)
        response.raise_for_status()
        
        return jsonify(response.json()), 200
        
    except requests.RequestException as e:
        return jsonify({"error": f"Error fetching study details: {str(e)}"}), 500
    except Exception as e:
        return jsonify({"error": f"Internal server error: {str(e)}"}), 500


if __name__ == '__main__':
    import os
    debug_mode = os.environ.get('FLASK_DEBUG', 'False').lower() == 'true'
    port = int(os.environ.get('PORT', 5001))
    app.run(debug=debug_mode, host='0.0.0.0', port=port)