"""
Zero-Shot Predictor for Influencer and Business Classification

This module orchestrates the prediction pipeline:
1. Create prompts
2. Generate responses with LLaMA
3. Parse JSON outputs
4. Return probabilities
"""

import logging
from typing import Dict

from .base_model import LLaMAModel
from ..utils.prompts import (
    create_influencer_niche_prompt,
    create_influencer_popularity_prompt,
    create_business_niche_prompt,
    create_business_popularity_prompt
)
from ..utils.parser import parse_json_response, validate_and_fix_probabilities
from ..constants import NICHE_CLASSES, POPULARITY_CLASSES

logger = logging.getLogger(__name__)


class ModelPredictor:
    """
    Zero-shot predictor using LLaMA 3.2-1B
    
    Handles both influencer and business predictions without fine-tuning.
    """
    
    def __init__(self):
        logger.info("Initializing ModelPredictor...")
        self.model = LLaMAModel()
        logger.info("ModelPredictor initialized successfully!")
    
    def predict_influencer(
        self,
        biography: str,
        recent_captions: str,
        recent_hashtags: str,
        followers: int,
        following: int,
        posts_count: int,
        verified: bool,
        avg_likes: float,
        avg_comments: float,
        avg_video_views: float,
        engagement_rate: float
    ) -> Dict:
        """
        Predict niche and popularity for an influencer
        
        Returns:
            {
                "niche_probabilities": {...},
                "popularity_probabilities": {...},
                "processed_data": {...}
            }
        """
        logger.info("Predicting influencer niche and popularity...")
        
        # Predict niche
        niche_probs = self._predict_niche(biography, recent_captions, recent_hashtags)
        
        # Predict popularity
        popularity_probs = self._predict_popularity(
            followers, engagement_rate, verified, avg_likes, avg_comments
        )
        
        # Create processed data snapshot
        processed_data = {
            "biography_length": len(biography) if biography else 0,
            "captions_length": len(recent_captions) if recent_captions else 0,
            "hashtags_length": len(recent_hashtags) if recent_hashtags else 0,
            "followers": followers,
            "engagement_rate": engagement_rate,
            "verified": verified,
            "top_niche": max(niche_probs, key=niche_probs.get),
            "top_popularity": max(popularity_probs, key=popularity_probs.get)
        }
        
        return {
            "niche_probabilities": niche_probs,
            "popularity_probabilities": popularity_probs,
            "processed_data": processed_data
        }
    
    def predict_business(
        self,
        business_description: str,
        budget_min: float,
        budget_max: float
    ) -> Dict:
        """
        Predict niche and popularity requirements for a business
        
        Returns:
            {
                "niche_probabilities": {...},
                "popularity_probabilities": {...},
                "processed_data": {...}
            }
        """
        logger.info("Predicting business requirements...")
        
        # Predict niche from description
        niche_probs = self._predict_business_niche(business_description)
        
        # Predict popularity from budget
        popularity_probs = self._predict_business_popularity(budget_min, budget_max)
        
        # Create processed data snapshot
        processed_data = {
            "description_length": len(business_description),
            "budget_min": budget_min,
            "budget_max": budget_max,
            "budget_avg": (budget_min + budget_max) / 2,
            "top_niche": max(niche_probs, key=niche_probs.get),
            "top_popularity": max(popularity_probs, key=popularity_probs.get)
        }
        
        return {
            "niche_probabilities": niche_probs,
            "popularity_probabilities": popularity_probs,
            "processed_data": processed_data
        }
    
    def _predict_niche(self, biography: str, recent_captions: str, recent_hashtags: str) -> Dict[str, float]:
        """Predict influencer niche using zero-shot classification"""
        prompt = create_influencer_niche_prompt(biography, recent_captions, recent_hashtags)
        response = self.model.generate(prompt, max_new_tokens=200, temperature=0.2)

        logger.info(f"Niche prediction raw response: {response[:500]}")  # Log first 500 chars

        parsed = parse_json_response(response)
        if parsed:
            logger.info(f"Niche prediction parsed successfully: {parsed}")
        else:
            logger.error(f"Failed to parse niche prediction response: {response}")

        return validate_and_fix_probabilities(parsed, NICHE_CLASSES, default_class="Lifestyle")
    
    def _predict_popularity(
        self, followers: int, engagement_rate: float, verified: bool,
        avg_likes: float, avg_comments: float
    ) -> Dict[str, float]:
        """Predict influencer popularity tier using zero-shot classification"""
        prompt = create_influencer_popularity_prompt(
            followers, engagement_rate, verified, avg_likes, avg_comments
        )
        response = self.model.generate(prompt, max_new_tokens=150, temperature=0.2)

        logger.info(f"Popularity prediction raw response: {response[:300]}")

        parsed = parse_json_response(response)
        if parsed:
            logger.info(f"Popularity prediction parsed successfully: {parsed}")
        else:
            logger.error(f"Failed to parse popularity prediction response: {response}")

        return validate_and_fix_probabilities(parsed, POPULARITY_CLASSES, default_class="Micro")
    
    def _predict_business_niche(self, business_description: str) -> Dict[str, float]:
        """Predict business niche requirements using zero-shot classification"""
        prompt = create_business_niche_prompt(business_description)
        response = self.model.generate(prompt, max_new_tokens=150, temperature=0.1)

        logger.debug(f"Business niche prediction response: {response}")

        parsed = parse_json_response(response)
        return validate_and_fix_probabilities(parsed, NICHE_CLASSES, default_class="Business")

    def _predict_business_popularity(self, budget_min: float, budget_max: float) -> Dict[str, float]:
        """Predict business popularity tier requirements using zero-shot classification"""
        prompt = create_business_popularity_prompt(budget_min, budget_max)
        response = self.model.generate(prompt, max_new_tokens=100, temperature=0.1)

        logger.debug(f"Business popularity prediction response: {response}")

        parsed = parse_json_response(response)
        return validate_and_fix_probabilities(parsed, POPULARITY_CLASSES, default_class="Micro")

    def is_loaded(self) -> bool:
        """Check if model is loaded"""
        return self.model.is_loaded()

    def get_loaded_adapters(self) -> list:
        """Get list of loaded adapters (empty for zero-shot)"""
        return []  # No adapters in zero-shot mode

