"""
Migrate Gemini AI → Mistral AI across the codebase.

Changes made:
1. app/config.py              – GEMINI_API_KEY → MISTRAL_API_KEY
2. requirements.txt           – google-generativeai → mistralai
3. .env.example               – GEMINI_API_KEY → MISTRAL_API_KEY
4. app/services/ai_assessment.py       – full migration
5. app/services/ai_assessment_fixed.py – full migration
"""
import pathlib, re

ROOT = pathlib.Path(r'c:\Users\SO\OneDrive - Higher Education Commission\Desktop\prot-final')

# ── helpers ──────────────────────────────────────────────────────────────────

def swap(path: pathlib.Path, old: str, new: str, count: int = -1):
    txt = path.read_text(encoding='utf-8')
    if count == -1:
        updated = txt.replace(old, new)
    else:
        updated = txt.replace(old, new, count)
    if updated == txt:
        print(f'  WARN: no match for {repr(old[:60])} in {path.name}')
    else:
        path.write_text(updated, encoding='utf-8')
        n = txt.count(old)
        print(f'  OK: replaced {n}x in {path.name}')

# ── 1. app/config.py ─────────────────────────────────────────────────────────
print('==> app/config.py')
cfg = ROOT / 'app' / 'config.py'
swap(cfg,
     "GEMINI_API_KEY = os.environ.get('GEMINI_API_KEY') or 'your-gemini-api-key-here'  # Get from environment or use default",
     "MISTRAL_API_KEY = os.environ.get('MISTRAL_API_KEY') or 'your-mistral-api-key-here'  # Get from environment or use default")

# ── 2. requirements.txt ───────────────────────────────────────────────────────
print('==> requirements.txt')
req = ROOT / 'requirements.txt'
swap(req, 'google-generativeai==0.3.2', 'mistralai>=1.0.0')

# ── 3. .env.example ──────────────────────────────────────────────────────────
print('==> .env.example')
env = ROOT / '.env.example'
swap(env, 'GEMINI_API_KEY=', 'MISTRAL_API_KEY=')
swap(env,
     '# AI Keys (optional; app has robust fallbacks)\nMISTRAL_API_KEY=\nOPENAI_API_KEY=  # For Whisper transcription (optional but highly recommended for best accuracy)',
     '# AI Keys (optional; app has robust fallbacks)\nMISTRAL_API_KEY=       # Get your key from https://console.mistral.ai/\nOPENAI_API_KEY=  # For Whisper transcription (optional but highly recommended for best accuracy)')

# ── Mistral wrapper block (injected once, used in both service files) ─────────
MISTRAL_WRAPPER = '''\
try:
    from mistralai import Mistral as _MistralClient
    _MISTRAL_AVAILABLE = True
except ImportError:
    _MISTRAL_AVAILABLE = False
    _MistralClient = None


class _MistralResponse:
    """Wraps a Mistral chat response to expose a `.text` attribute."""
    def __init__(self, content: str):
        self.text = content


class MistralModelWrapper:
    """Drop-in replacement for `genai.GenerativeModel`.

    Exposes `generate_content(content)` so all existing callers keep working
    without modification. `content` may be either:
      - a plain string prompt, or
      - a list [prompt_str, {\'mime_type\': ..., \'data\': ...}]  (legacy Gemini
        multimodal – the audio part is silently dropped because Mistral is
        text-only; the text prompt is still sent as usual).
    """

    def __init__(self, client, model_name: str):
        self._client = client
        self._model_name = model_name

    def generate_content(self, content) -> _MistralResponse:
        if isinstance(content, list):
            # Extract the text part from a multimodal list; ignore binary blobs
            text_parts = [p for p in content if isinstance(p, str)]
            prompt = ' '.join(text_parts) if text_parts else ''
        else:
            prompt = str(content)

        response = self._client.chat.complete(
            model=self._model_name,
            messages=[{'role': 'user', 'content': prompt}],
        )
        text = response.choices[0].message.content
        return _MistralResponse(text)

'''

# ── 4. app/services/ai_assessment.py ─────────────────────────────────────────
print('==> app/services/ai_assessment.py')
ai_path = ROOT / 'app' / 'services' / 'ai_assessment.py'
ai_txt = ai_path.read_text(encoding='utf-8')

# 4a. Replace the Gemini import with the Mistral wrapper block
ai_txt = ai_txt.replace(
    'import google.generativeai as genai\n',
    MISTRAL_WRAPPER,
    1
)

# 4b. Replace _configure_genai body (inside AIAssessment class)
old_cfg_ai = '''\
    def _configure_genai(self):
        # Configure Gemini API only when needed
        if not hasattr(self, 'genai_configured'):
            api_key = current_app.config['GEMINI_API_KEY']
            
            # Check if the API key is valid (not empty and not the default placeholder)
            if not api_key or api_key == 'your-gemini-api-key-here':
                logger.error("Invalid or missing Gemini API key. Please set a valid GEMINI_API_KEY in your .env file.")
                self.genai_configured = False
                return False
                
            try:
                genai.configure(api_key=api_key)
                self.genai_configured = True
                logger.debug("Gemini API configured successfully")
                return True
            except Exception as e:
                logger.error(f"Error configuring Gemini API: {str(e)}")
                self.genai_configured = False
                return False
        
        return self.genai_configured'''

new_cfg_ai = '''\
    def _configure_genai(self):
        # Configure Mistral API only when needed
        if not hasattr(self, 'genai_configured'):
            if not _MISTRAL_AVAILABLE:
                logger.error("mistralai package not installed. Run: pip install mistralai")
                self.genai_configured = False
                return False

            api_key = current_app.config.get('MISTRAL_API_KEY', '')

            # Check if the API key is valid (not empty and not the default placeholder)
            if not api_key or api_key == 'your-mistral-api-key-here':
                logger.error("Invalid or missing Mistral API key. Please set a valid MISTRAL_API_KEY in your .env file.")
                self.genai_configured = False
                return False

            try:
                self._mistral_client = _MistralClient(api_key=api_key)
                self.genai_configured = True
                logger.debug("Mistral API configured successfully")
                return True
            except Exception as e:
                logger.error(f"Error configuring Mistral API: {str(e)}")
                self.genai_configured = False
                return False

        return self.genai_configured'''

if old_cfg_ai in ai_txt:
    ai_txt = ai_txt.replace(old_cfg_ai, new_cfg_ai, 1)
    print('  OK: replaced _configure_genai in ai_assessment.py')
else:
    print('  WARN: _configure_genai block not found in ai_assessment.py')

# 4c. Replace get_model body
old_model_ai = '''\
    def get_model(self):
        """Lazy initialization of Gemini model"""
        if self.model is None:
            if not self._configure_genai():
                logger.error("Failed to configure Gemini API")
                return None
            # List of models to try in order of preference for multimodal support
            models_to_try = [
                'gemini-2.5-flash',          # Latest fast multimodal model
                'gemini-2.5-pro',            # Latest pro multimodal model
                'gemini-2.0-flash',          # Stable multimodal model
                'gemini-flash-latest',       # Latest flash model
                'gemini-pro-latest'          # Latest pro model fallback
            ]
            for model_name in models_to_try:
                try:
                    logger.info(f"Attempting to initialize model: {model_name}")
                    self.model = genai.GenerativeModel(model_name)
                    
                    # Test the model
                    test_response = self.model.generate_content("Hello, world!")
                    logger.info(f"Successfully initialized and tested model: {model_name}")
                    return self.model
                    
                except Exception as e:
                    logger.warning(f"Failed to initialize model {model_name}: {str(e)}")
                    self.model = None  # Reset for next attempt
            
            # If we get here, all models failed
            logger.error("All Gemini models failed to initialize")
            return None
                
        return self.model'''

new_model_ai = '''\
    def get_model(self):
        """Lazy initialization of Mistral model wrapper."""
        if self.model is None:
            if not self._configure_genai():
                logger.error("Failed to configure Mistral API")
                return None
            # Models to try in order of preference (cost/speed vs quality)
            models_to_try = [
                'mistral-small-latest',   # Fast, cost-effective
                'mistral-medium-latest',  # Balanced
                'mistral-large-latest',   # Highest quality fallback
            ]
            for model_name in models_to_try:
                try:
                    logger.info(f"Attempting to initialize Mistral model: {model_name}")
                    wrapper = MistralModelWrapper(self._mistral_client, model_name)
                    # Quick smoke-test
                    test_response = wrapper.generate_content("Hello, world!")
                    logger.info(f"Successfully initialized and tested Mistral model: {model_name}")
                    self.model = wrapper
                    return self.model
                except Exception as e:
                    logger.warning(f"Failed to initialize Mistral model {model_name}: {str(e)}")
                    self.model = None  # Reset for next attempt

            logger.error("All Mistral models failed to initialize")
            return None

        return self.model'''

if old_model_ai in ai_txt:
    ai_txt = ai_txt.replace(old_model_ai, new_model_ai, 1)
    print('  OK: replaced get_model in ai_assessment.py')
else:
    print('  WARN: get_model block not found in ai_assessment.py')

# 4d. Replace the multimodal Gemini-specific block inside _ai_transcribe_and_assess
old_multimodal = '''\
                    logger.info("Sending audio to AI model for transcription+assessment")
                    
                    # Create proper multimodal content for Gemini
                    try:
                        import google.generativeai as genai
                        
                        # Create the content with both text prompt and audio
                        content = [
                            prompt,
                            {
                                'mime_type': 'audio/wav',
                                'data': audio_bytes
                            }
                        ]
                        
                        response = model.generate_content(content)
                    except Exception as api_error:
                        logger.error(f"Error calling Gemini API with multimodal content: {api_error}")
                        # Try fallback approach
                        try:
                            response = model.generate_content(prompt)
                            logger.warning("Falling back to text-only Gemini call due to multimodal error")
                        except Exception as fallback_error:
                            logger.error(f"Both multimodal and text-only Gemini calls failed: {fallback_error}")
                            return None'''

new_multimodal = '''\
                    logger.info("Sending prompt to Mistral API for speaking assessment (text-only)")

                    # Mistral is text-only; we cannot pass raw audio bytes.
                    # The MistralModelWrapper.generate_content() handles this gracefully:
                    # if a list is passed it extracts the text part and ignores binary data.
                    try:
                        response = model.generate_content(prompt)
                    except Exception as api_error:
                        logger.error(f"Error calling Mistral API: {api_error}")
                        return None'''

if old_multimodal in ai_txt:
    ai_txt = ai_txt.replace(old_multimodal, new_multimodal, 1)
    print('  OK: replaced multimodal block in ai_assessment.py')
else:
    print('  WARN: multimodal block not found in ai_assessment.py')

# 4e. Replace remaining "Gemini" string references in log messages / comments
ai_txt = ai_txt.replace(
    'logger.warning("Gemini API not available, using fallback speaking assessment")',
    'logger.warning("Mistral API not available, using fallback speaking assessment")'
)
ai_txt = ai_txt.replace(
    'Toggle AI-based transcription. When enabled, audio bytes are sent to configured Gemini model.',
    'Toggle AI-based transcription. When enabled, speaking assessment is sent to the configured Mistral model.'
)

# 4f. Inline import inside _ai_transcribe_and_assess no longer needed
ai_txt = ai_txt.replace('import google.generativeai as genai\n        \n        # Create the content with both text prompt and audio', '')

ai_path.write_text(ai_txt, encoding='utf-8')
print(f'  Wrote ai_assessment.py ({len(ai_txt.splitlines())} lines)')

# ── 5. app/services/ai_assessment_fixed.py ───────────────────────────────────
print('==> app/services/ai_assessment_fixed.py')
fix_path = ROOT / 'app' / 'services' / 'ai_assessment_fixed.py'
fix_txt = fix_path.read_text(encoding='utf-8')

# 5a. Replace import
fix_txt = fix_txt.replace(
    'import google.generativeai as genai\n',
    MISTRAL_WRAPPER,
    1
)

# 5b. Replace _configure_genai
old_cfg_fix = '''\
    def _configure_genai(self):
        # Configure Gemini API only when needed
        if not hasattr(self, 'genai_configured'):
            api_key = current_app.config['GEMINI_API_KEY']
            
            # Check if the API key is valid (not empty and not the default placeholder)
            if not api_key or api_key == 'your-gemini-api-key-here':
                logger.error("Invalid or missing Gemini API key. Please set a valid GEMINI_API_KEY in your .env file.")
                self.genai_configured = False
                return False
                
            try:
                genai.configure(api_key=api_key)
                self.genai_configured = True
                logger.debug("Gemini API configured successfully")
                return True
            except Exception as e:
                logger.error(f"Error configuring Gemini API: {str(e)}")
                self.genai_configured = False
                return False
        
        return self.genai_configured'''

new_cfg_fix = '''\
    def _configure_genai(self):
        # Configure Mistral API only when needed
        if not hasattr(self, 'genai_configured'):
            if not _MISTRAL_AVAILABLE:
                logger.error("mistralai package not installed. Run: pip install mistralai")
                self.genai_configured = False
                return False

            api_key = current_app.config.get('MISTRAL_API_KEY', '')

            if not api_key or api_key == 'your-mistral-api-key-here':
                logger.error("Invalid or missing Mistral API key. Please set a valid MISTRAL_API_KEY in your .env file.")
                self.genai_configured = False
                return False

            try:
                self._mistral_client = _MistralClient(api_key=api_key)
                self.genai_configured = True
                logger.debug("Mistral API configured successfully")
                return True
            except Exception as e:
                logger.error(f"Error configuring Mistral API: {str(e)}")
                self.genai_configured = False
                return False

        return self.genai_configured'''

if old_cfg_fix in fix_txt:
    fix_txt = fix_txt.replace(old_cfg_fix, new_cfg_fix, 1)
    print('  OK: replaced _configure_genai in ai_assessment_fixed.py')
else:
    print('  WARN: _configure_genai block not found in ai_assessment_fixed.py')

# 5c. Replace get_model
old_model_fix = '''\
    def get_model(self):
        """Lazy initialization of Gemini model"""
        if self.model is None:
            if not self._configure_genai():
                logger.error("Failed to configure Gemini API")
                return None
            # List of models to try in order of preference
            models_to_try = [
                'gemini-1.5-flash',   # First choice
                'gemini-pro',         # Second choice
                'gemini-1.0-pro'      # Third choice
            ]
            for model_name in models_to_try:
                try:
                    logger.info(f"Attempting to initialize model: {model_name}")
                    self.model = genai.GenerativeModel(model_name)
                    
                    # Test the model
                    test_response = self.model.generate_content("Hello, world!")
                    logger.info(f"Successfully initialized and tested model: {model_name}")
                    return self.model
                    
                except Exception as e:
                    logger.warning(f"Failed to initialize model {model_name}: {str(e)}")
                    self.model = None  # Reset for next attempt
            
            # If we get here, all models failed
            logger.error("All Gemini models failed to initialize")
            return None
                
        return self.model'''

new_model_fix = '''\
    def get_model(self):
        """Lazy initialization of Mistral model wrapper."""
        if self.model is None:
            if not self._configure_genai():
                logger.error("Failed to configure Mistral API")
                return None
            models_to_try = [
                'mistral-small-latest',
                'mistral-medium-latest',
                'mistral-large-latest',
            ]
            for model_name in models_to_try:
                try:
                    logger.info(f"Attempting to initialize Mistral model: {model_name}")
                    wrapper = MistralModelWrapper(self._mistral_client, model_name)
                    test_response = wrapper.generate_content("Hello, world!")
                    logger.info(f"Successfully initialized Mistral model: {model_name}")
                    self.model = wrapper
                    return self.model
                except Exception as e:
                    logger.warning(f"Failed to initialize Mistral model {model_name}: {str(e)}")
                    self.model = None

            logger.error("All Mistral models failed to initialize")
            return None

        return self.model'''

if old_model_fix in fix_txt:
    fix_txt = fix_txt.replace(old_model_fix, new_model_fix, 1)
    print('  OK: replaced get_model in ai_assessment_fixed.py')
else:
    print('  WARN: get_model block not found in ai_assessment_fixed.py')

# 5d. Replace remaining Gemini strings in ai_assessment_fixed.py
fix_txt = fix_txt.replace(
    "logger.error(\"Failed to configure Gemini API. Using fallback vocabulary quiz.\")",
    "logger.error(\"Failed to configure Mistral API. Using fallback vocabulary quiz.\")"
)
fix_txt = fix_txt.replace(
    "logger.error(\"Failed to get AI model. Using fallback vocabulary quiz.\")",
    "logger.error(\"Failed to get Mistral AI model. Using fallback vocabulary quiz.\")"
)

fix_path.write_text(fix_txt, encoding='utf-8')
print(f'  Wrote ai_assessment_fixed.py ({len(fix_txt.splitlines())} lines)')

print('\n==> Done. Summary:')
print('  - GEMINI_API_KEY → MISTRAL_API_KEY  (config.py, .env.example)')
print('  - google-generativeai → mistralai>=1.0.0  (requirements.txt)')
print('  - Gemini SDK calls replaced with MistralModelWrapper  (both service files)')
print('  - Audio multimodal block replaced with text-only Mistral call  (ai_assessment.py)')