"""
LLM Service for intelligent ranking and scoring of clinical trials.
Supports Hugging Face models including DeepSeek-V3.2.
"""
import os
import logging
from typing import List, Dict, Optional
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

logger = logging.getLogger(__name__)

class LLMService:
    """Service for interacting with Hugging Face LLM models"""
    
    def __init__(self, model_name: Optional[str] = None, use_api: bool = False, api_token: Optional[str] = None):
        """
        Initialize LLM service
        
        Args:
            model_name: Hugging Face model identifier (e.g., 'deepseek-ai/DeepSeek-V3.2')
                       If None, uses DEEPSEEK_MODEL env var or defaults to DeepSeek-V3.2
            use_api: If True, use Hugging Face Inference API instead of local model
            api_token: Hugging Face API token (required if use_api=True)
        """
        self.model_name = model_name or os.environ.get('DEEPSEEK_MODEL', 'deepseek-ai/DeepSeek-V3.2')
        self.use_api = use_api or os.environ.get('USE_HF_API', 'false').lower() == 'true'
        self.api_token = api_token or os.environ.get('HUGGINGFACE_API_TOKEN', '')
        
        self.tokenizer = None
        self.model = None
        self.pipeline = None
        
        if not self.use_api:
            self._load_local_model()
        else:
            if not self.api_token:
                logger.warning("Hugging Face API token not provided. Set HUGGINGFACE_API_TOKEN env var.")
    
    def _load_local_model(self):
        """Load model locally using transformers"""
        try:
            logger.info(f"Loading model: {self.model_name}")
            
            # Check if CUDA is available
            device = "cuda" if torch.cuda.is_available() else "cpu"
            logger.info(f"Using device: {device}")
            
            # Load tokenizer and model
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name,
                trust_remote_code=True
            )
            
            # Load model with appropriate settings
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                trust_remote_code=True,
                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                device_map="auto" if device == "cuda" else None,
                low_cpu_mem_usage=True
            )
            
            if device == "cpu":
                self.model = self.model.to(device)
            
            # Create pipeline for easier text generation
            self.pipeline = pipeline(
                "text-generation",
                model=self.model,
                tokenizer=self.tokenizer,
                device=0 if device == "cuda" else -1,
                torch_dtype=torch.float16 if device == "cuda" else torch.float32
            )
            
            logger.info(f"Model {self.model_name} loaded successfully")
            
        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            raise
    
    def rank_studies(self, studies: List[Dict], ranking_terms: str) -> List[Dict]:
        """
        Rank studies based on relevance to ranking terms using LLM
        
        Args:
            studies: List of study dictionaries
            ranking_terms: Terms to use for ranking (e.g., "KRAS mutation, immunotherapy")
        
        Returns:
            List of studies sorted by relevance score (highest first), with ranking_reasoning added
        """
        if not ranking_terms or not ranking_terms.strip():
            return studies
        
        if not studies:
            return studies
        
        try:
            # Score each study
            scored_studies = []
            for study in studies:
                score, reasoning = self._score_study(study, ranking_terms)
                study_with_score = study.copy()
                study_with_score['relevance_score'] = score
                study_with_score['ranking_reasoning'] = reasoning
                scored_studies.append(study_with_score)
            
            # Sort by score (highest first)
            scored_studies.sort(key=lambda x: x.get('relevance_score', 0), reverse=True)
            
            return scored_studies
            
        except Exception as e:
            logger.error(f"Error ranking studies: {str(e)}")
            # Return original studies if ranking fails
            return studies
    
    def _score_study(self, study: Dict, ranking_terms: str) -> tuple:
        """
        Score a single study's relevance to ranking terms and get reasoning
        
        Args:
            study: Study dictionary
            ranking_terms: Terms to match against
        
        Returns:
            Tuple of (relevance score (0.0 to 1.0), reasoning explanation)
        """
        try:
            # Build context from study
            study_text = self._build_study_context(study)
            
            # Create prompt for scoring with reasoning
            prompt = f"""You are a medical research assistant. Rate the relevance of this clinical trial to the search terms on a scale of 0.0 to 1.0.

Search terms: {ranking_terms}

Clinical Trial:
{study_text}

Provide your response in this exact format:
SCORE: [number between 0.0 and 1.0]
REASONING: [brief explanation of why this score was assigned, focusing on how the study matches or doesn't match the search terms]"""

            if self.use_api:
                score, reasoning = self._score_with_reasoning_api(prompt)
            else:
                score, reasoning = self._score_with_reasoning_local(prompt)
            
            # Ensure score is between 0 and 1
            score = max(0.0, min(1.0, float(score)))
            
            return score, reasoning
            
        except Exception as e:
            logger.error(f"Error scoring study {study.get('nctId', 'unknown')}: {str(e)}")
            return 0.0, "Unable to generate reasoning due to an error."
    
    def _build_study_context(self, study: Dict) -> str:
        """Build a text context from study data"""
        parts = []
        
        if study.get('title'):
            parts.append(f"Title: {study['title']}")
        
        if study.get('briefSummary'):
            parts.append(f"Summary: {study['briefSummary'][:500]}")  # Limit summary length
        
        if study.get('conditions'):
            parts.append(f"Conditions: {', '.join(study['conditions'])}")
        
        if study.get('inclusionCriteria'):
            inclusion_text = ' '.join(study['inclusionCriteria'][:3])  # First 3 criteria
            parts.append(f"Inclusion Criteria: {inclusion_text[:300]}")
        
        return "\n".join(parts)
    
    def _score_with_local_model(self, prompt: str) -> float:
        """Score using local model (legacy method)"""
        try:
            # Generate response
            outputs = self.pipeline(
                prompt,
                max_new_tokens=10,
                temperature=0.1,
                do_sample=False,
                return_full_text=False
            )
            
            # Extract score from response
            response_text = outputs[0]['generated_text'].strip()
            
            # Try to extract a number from the response
            import re
            numbers = re.findall(r'\d+\.?\d*', response_text)
            if numbers:
                score = float(numbers[0])
                # Normalize if it's > 1 (might be percentage or 0-100 scale)
                if score > 1.0:
                    score = score / 100.0
                return score
            
            return 0.5  # Default score if parsing fails
            
        except Exception as e:
            logger.error(f"Error in local model scoring: {str(e)}")
            return 0.5
    
    def _score_with_reasoning_local(self, prompt: str) -> tuple:
        """Score with reasoning using local model"""
        try:
            # Generate response with more tokens for reasoning
            outputs = self.pipeline(
                prompt,
                max_new_tokens=150,
                temperature=0.3,
                do_sample=True,
                return_full_text=False
            )
            
            # Extract response text
            response_text = outputs[0]['generated_text'].strip()
            
            # Parse score and reasoning
            import re
            score_match = re.search(r'SCORE:\s*([\d.]+)', response_text, re.IGNORECASE)
            reasoning_match = re.search(r'REASONING:\s*(.+?)(?=SCORE:|$)', response_text, re.IGNORECASE | re.DOTALL)
            
            score = 0.5
            if score_match:
                score = float(score_match.group(1))
                if score > 1.0:
                    score = score / 100.0
                score = max(0.0, min(1.0, score))
            
            reasoning = "No specific reasoning provided."
            if reasoning_match:
                reasoning = reasoning_match.group(1).strip()
            elif not score_match:
                # Fallback: try to extract any number as score
                numbers = re.findall(r'\d+\.?\d*', response_text)
                if numbers:
                    score = float(numbers[0])
                    if score > 1.0:
                        score = score / 100.0
                    score = max(0.0, min(1.0, score))
                reasoning = response_text[:200] if len(response_text) > 0 else "Unable to parse reasoning."
            
            return score, reasoning
            
        except Exception as e:
            logger.error(f"Error in local model scoring with reasoning: {str(e)}")
            return 0.5, "Unable to generate reasoning due to an error."
    
    def _score_with_api(self, prompt: str) -> float:
        """Score using Hugging Face Inference API (legacy method)"""
        try:
            import requests
            
            # Try router endpoint first, fallback to inference API
            api_url = f"https://api-inference.huggingface.co/models/{self.model_name}"
            headers = {"Authorization": f"Bearer {self.api_token}"}
            
            payload = {
                "inputs": prompt,
                "parameters": {
                    "max_new_tokens": 10,
                    "temperature": 0.1,
                    "return_full_text": False
                }
            }
            
            response = requests.post(api_url, headers=headers, json=payload, timeout=30)
            response.raise_for_status()
            
            result = response.json()
            
            # Extract generated text
            if isinstance(result, list) and len(result) > 0:
                generated_text = result[0].get('generated_text', '')
            else:
                generated_text = str(result)
            
            # Extract score
            import re
            numbers = re.findall(r'\d+\.?\d*', generated_text)
            if numbers:
                score = float(numbers[0])
                if score > 1.0:
                    score = score / 100.0
                return score
            
            return 0.5
            
        except Exception as e:
            logger.error(f"Error in API scoring: {str(e)}")
            return 0.5
    
    def _score_with_reasoning_api(self, prompt: str) -> tuple:
        """Score with reasoning using Hugging Face Inference API"""
        try:
            import requests
            
            # Use router API endpoint (new format)
            api_url = f"https://router.huggingface.co/v1/models/{self.model_name}/generate"
            headers = {
                "Authorization": f"Bearer {self.api_token}",
                "Content-Type": "application/json"
            }
            
            payload = {
                "inputs": prompt,
                "parameters": {
                    "max_new_tokens": 150,
                    "temperature": 0.3,
                    "return_full_text": False
                }
            }
            
            response = requests.post(api_url, headers=headers, json=payload, timeout=60)
            
            # If router endpoint fails, try alternative format
            if response.status_code == 404:
                # Try OpenAI-compatible format
                api_url = f"https://router.huggingface.co/v1/chat/completions"
                payload = {
                    "model": self.model_name,
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": 150,
                    "temperature": 0.3
                }
                response = requests.post(api_url, headers=headers, json=payload, timeout=60)
            
            response.raise_for_status()
            
            result = response.json()
            
            # Extract generated text
            if isinstance(result, list) and len(result) > 0:
                generated_text = result[0].get('generated_text', '')
            else:
                generated_text = str(result)
            
            # Parse score and reasoning
            import re
            score_match = re.search(r'SCORE:\s*([\d.]+)', generated_text, re.IGNORECASE)
            reasoning_match = re.search(r'REASONING:\s*(.+?)(?=SCORE:|$)', generated_text, re.IGNORECASE | re.DOTALL)
            
            score = 0.5
            if score_match:
                score = float(score_match.group(1))
                if score > 1.0:
                    score = score / 100.0
                score = max(0.0, min(1.0, score))
            
            reasoning = "No specific reasoning provided."
            if reasoning_match:
                reasoning = reasoning_match.group(1).strip()
            elif not score_match:
                # Fallback: try to extract any number as score
                numbers = re.findall(r'\d+\.?\d*', generated_text)
                if numbers:
                    score = float(numbers[0])
                    if score > 1.0:
                        score = score / 100.0
                    score = max(0.0, min(1.0, score))
                reasoning = generated_text[:200] if len(generated_text) > 0 else "Unable to parse reasoning."
            
            return score, reasoning
            
        except Exception as e:
            logger.error(f"Error in API scoring with reasoning: {str(e)}")
            return 0.5, "Unable to generate reasoning due to an error."


# Global LLM service instance (lazy loaded)
_llm_service = None

def get_llm_service() -> Optional[LLMService]:
    """Get or create LLM service instance"""
    global _llm_service
    
    if _llm_service is None:
        try:
            use_api = os.environ.get('USE_HF_API', 'false').lower() == 'true'
            _llm_service = LLMService(use_api=use_api)
        except Exception as e:
            logger.error(f"Failed to initialize LLM service: {str(e)}")
            return None
    
    return _llm_service