#!/usr/bin/env python3
"""
Real AI Models Service - ZERO MOCK DATA
All AI predictions use REAL models from HuggingFace
"""

import logging
from typing import Dict, Any, Optional
from datetime import datetime
import asyncio
import time
import hashlib

logger = logging.getLogger(__name__)

# Try to import transformers - if not available, use HF API
try:
    from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False
    logger.warning("⚠ Transformers not available, will use HF API")

import httpx
from backend.services.real_api_clients import RealAPIConfiguration


class RealAIModelsRegistry:
    """
    Real AI Models Registry using HuggingFace models
    NO MOCK PREDICTIONS - Only real model inference
    """
    
    def __init__(self):
        self.models = {}
        self.loaded = False
        import os
        # Strip whitespace from token to avoid "Illegal header value" errors
        token_raw = os.getenv("HF_API_TOKEN") or os.getenv("HF_TOKEN") or RealAPIConfiguration.HF_API_TOKEN or ""
        token = str(token_raw).strip() if token_raw else ""
        self.hf_api_token = token if token else None
        self.hf_api_url = "https://router.huggingface.co/models"

        # Simple in-memory cache to reduce repeated HF Inference calls
        # key -> {"time": float, "data": Any}
        self._cache: Dict[str, Dict[str, Any]] = {}
        
        # Model configurations - REAL HuggingFace models with fallback chain
        # Each task has at least 3 fallback models
        self.model_configs = {
            "sentiment_crypto": {
                "model_id": "ElKulako/cryptobert",
                "task": "sentiment-analysis",
                "description": "CryptoBERT for crypto sentiment analysis",
                "fallbacks": [
                    "kk08/CryptoBERT",
                    "ProsusAI/finbert",
                    "cardiffnlp/twitter-roberta-base-sentiment-latest",
                    "distilbert-base-uncased-finetuned-sst-2-english"
                ]
            },
            "sentiment_twitter": {
                "model_id": "cardiffnlp/twitter-roberta-base-sentiment-latest",
                "task": "sentiment-analysis",
                "description": "Twitter sentiment analysis",
                "fallbacks": [
                    "cardiffnlp/twitter-roberta-base-sentiment",
                    "ProsusAI/finbert",
                    "distilbert-base-uncased-finetuned-sst-2-english",
                    "nlptown/bert-base-multilingual-uncased-sentiment"
                ]
            },
            "sentiment_financial": {
                "model_id": "ProsusAI/finbert",
                "task": "sentiment-analysis",
                "description": "FinBERT for financial sentiment",
                "fallbacks": [
                    "yiyanghkust/finbert-tone",
                    "mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis",
                    "cardiffnlp/twitter-roberta-base-sentiment-latest",
                    "distilbert-base-uncased-finetuned-sst-2-english"
                ]
            },
            "text_generation": {
                # Use a widely-available, lightweight text generation model as primary
                # to avoid "model not found / gated / gone" failures during deploy.
                "model_id": "gpt2",
                "task": "text-generation",
                "description": "Text generation (lightweight)",
                "fallbacks": [
                    "distilgpt2",
                    "EleutherAI/gpt-neo-125M"
                ]
            },
            "trading_signals": {
                # Keep signals reliable; prompt will be crypto-specific.
                "model_id": "gpt2",
                "task": "text-generation",
                "description": "Trading signals (prompted text generation)",
                "fallbacks": [
                    "distilgpt2",
                    "EleutherAI/gpt-neo-125M"
                ]
            },
            "summarization": {
                "model_id": "facebook/bart-large-cnn",
                "task": "summarization",
                "description": "BART for news summarization",
                "fallbacks": [
                    "sshleifer/distilbart-cnn-12-6",
                    "google/pegasus-xsum",
                    "facebook/bart-large",
                    "FurkanGozukara/Crypto-Financial-News-Summarizer",
                    "facebook/mbart-large-50"
                ]
            }
        }
    
    async def load_models(self):
        """
        Load REAL models from HuggingFace
        """
        if self.loaded:
            return {"status": "already_loaded", "models": len(self.models)}
        
        logger.info("🤖 Loading REAL AI models from HuggingFace...")
        
        if TRANSFORMERS_AVAILABLE:
            # Load models locally using transformers
            for model_key, config in self.model_configs.items():
                try:
                    if config["task"] == "sentiment-analysis":
                        self.models[model_key] = pipeline(
                            config["task"],
                            model=config["model_id"],
                            truncation=True,
                            max_length=512
                        )
                        logger.info(f"✅ Loaded local model: {config['model_id']}")
                    # For text generation, we'll use API to avoid heavy downloads
                except Exception as e:
                    logger.warning(f"⚠ Could not load {model_key} locally: {e}")
        
        self.loaded = True
        return {
            "status": "loaded",
            "models_local": len(self.models),
            "models_api": len(self.model_configs) - len(self.models),
            "total": len(self.model_configs)
        }
    
    async def predict_sentiment(
        self,
        text: str,
        model_key: str = "sentiment_crypto"
    ) -> Dict[str, Any]:
        """
        Run REAL sentiment analysis using HuggingFace models
        NO FAKE PREDICTIONS
        """
        try:
            # Check if model is loaded locally
            if model_key in self.models:
                # Use local model
                result = self.models[model_key](text)[0]
                
                return {
                    "success": True,
                    "label": result["label"],
                    "score": result["score"],
                    "model": model_key,
                    "source": "local",
                    "timestamp": datetime.utcnow().isoformat()
                }
            else:
                # Use HuggingFace API
                return await self._predict_via_api(text, model_key)
        
        except Exception as e:
            logger.error(f"❌ Sentiment prediction failed: {e}")
            raise Exception(f"Failed to predict sentiment: {str(e)}")
    
    async def generate_text(
        self,
        prompt: str,
        model_key: str = "text_generation",
        max_length: int = 200
    ) -> Dict[str, Any]:
        """
        Generate REAL text using HuggingFace models
        NO FAKE GENERATION
        """
        try:
            return await self._generate_via_api(prompt, model_key, max_length)
        except Exception as e:
            logger.error(f"❌ Text generation failed: {e}")
            raise Exception(f"Failed to generate text: {str(e)}")
    
    async def get_trading_signal(
        self,
        symbol: str,
        context: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Get REAL trading signal using HF text-generation (prompted)
        NO FAKE SIGNALS
        """
        try:
            # Prepare prompt for trading model
            prompt = f"Trading signal for {symbol}."
            if context:
                prompt += f" Context: {context}"
            
            result = await self._generate_via_api(
                prompt,
                "trading_signals",
                max_length=100
            )
            
            # Parse trading signal from generated text
            generated_text = result.get("generated_text", "").upper()
            
            # Determine signal type
            if "BUY" in generated_text or "BULLISH" in generated_text:
                signal_type = "BUY"
                score = 0.75
            elif "SELL" in generated_text or "BEARISH" in generated_text:
                signal_type = "SELL"
                score = 0.75
            else:
                signal_type = "HOLD"
                score = 0.60
            
            return {
                "success": True,
                "symbol": symbol,
                "signal": signal_type,
                "score": score,
                "explanation": result.get("generated_text", ""),
                "model": "trading_signals",
                "timestamp": datetime.utcnow().isoformat()
            }
        
        except Exception as e:
            logger.error(f"❌ Trading signal failed: {e}")
            raise Exception(f"Failed to get trading signal: {str(e)}")
    
    async def summarize_news(
        self,
        text: str
    ) -> Dict[str, Any]:
        """
        Summarize REAL news using BART
        NO FAKE SUMMARIES
        """
        try:
            return await self._summarize_via_api(text)
        except Exception as e:
            logger.error(f"❌ News summarization failed: {e}")
            raise Exception(f"Failed to summarize news: {str(e)}")
    
    async def _predict_via_api(
        self,
        text: str,
        model_key: str
    ) -> Dict[str, Any]:
        """
        Run REAL inference via HuggingFace API with fallback chain
        Tries at least 3 models before failing
        """
        config = self.model_configs.get(model_key)
        if not config:
            raise ValueError(f"Unknown model: {model_key}")
        
        # Build fallback chain: primary model + fallbacks
        models_to_try = [config["model_id"]] + config.get("fallbacks", [])
        
        last_error = None
        for model_id in models_to_try[:5]:  # Try up to 5 models
            try:
                logger.info(f"🔄 Trying sentiment model: {model_id}")
                async with httpx.AsyncClient(timeout=30.0) as client:
                    _headers = {"Content-Type": "application/json"}
                    if self.hf_api_token:
                        _headers["Authorization"] = f"Bearer {self.hf_api_token}"
                    response = await client.post(
                        f"{self.hf_api_url}/{model_id}",
                        headers=_headers,
                        json={"inputs": text[:512]}  # Limit input length
                    )
                    response.raise_for_status()
                    result = response.json()
                
                # Parse result based on task type
                if isinstance(result, list) and len(result) > 0:
                    if isinstance(result[0], list):
                        result = result[0]
                    
                    if isinstance(result[0], dict):
                        top_result = result[0]
                        label = top_result.get("label", "neutral")
                        score = top_result.get("score", 0.0)
                        
                        # Normalize label
                        label_upper = label.upper()
                        if "POSITIVE" in label_upper or "LABEL_2" in label_upper:
                            normalized_label = "positive"
                        elif "NEGATIVE" in label_upper or "LABEL_0" in label_upper:
                            normalized_label = "negative"
                        else:
                            normalized_label = "neutral"
                        
                        logger.info(f"✅ Sentiment analysis succeeded with {model_id}: {normalized_label} ({score})")
                        return {
                            "success": True,
                            "label": normalized_label,
                            "score": score,
                            "confidence": score,
                            "model": model_id,
                            "source": "hf_api",
                            "fallback_used": model_id != config["model_id"],
                            "timestamp": datetime.utcnow().isoformat()
                        }
                
                # If we got here, result format is unexpected but not an error
                return {
                    "success": True,
                    "result": result,
                    "model": model_id,
                    "source": "hf_api",
                    "fallback_used": model_id != config["model_id"],
                    "timestamp": datetime.utcnow().isoformat()
                }
            except Exception as e:
                logger.warning(f"⚠️ Sentiment model {model_id} failed: {e}")
                last_error = e
                continue
        
        logger.error(f"❌ All sentiment models failed. Last error: {last_error}")
        raise Exception(f"Failed to predict sentiment: All models failed. Tried: {models_to_try[:5]}")
    
    async def _generate_via_api(
        self,
        prompt: str,
        model_key: str,
        max_length: int = 200
    ) -> Dict[str, Any]:
        """
        Generate REAL text via HuggingFace API with fallback chain
        """
        config = self.model_configs.get(model_key)
        if not config:
            raise ValueError(f"Unknown model: {model_key}")

        # Cache key
        cache_key_raw = f"gen:{model_key}:{max_length}:{prompt}".encode("utf-8", errors="ignore")
        cache_key = hashlib.sha256(cache_key_raw).hexdigest()
        cached = self._cache.get(cache_key)
        if cached and (time.time() - float(cached.get("time", 0))) < 45:
            data = cached.get("data")
            if isinstance(data, dict):
                return data

        models_to_try = [config["model_id"]] + config.get("fallbacks", [])
        last_error = None

        for model_id in models_to_try[:5]:
            try:
                logger.info(f"🔄 Trying generation model: {model_id}")
                result = await self._post_hf_inference(
                    model_id=model_id,
                    payload={
                        "inputs": prompt[:2000],
                        "parameters": {
                            # Some endpoints prefer max_new_tokens; keep both to be safe.
                            "max_new_tokens": max(16, min(max_length, 256)),
                            "max_length": max_length,
                            "temperature": 0.7,
                            "top_p": 0.9,
                            "do_sample": True,
                            "return_full_text": True,
                        },
                    },
                    timeout_seconds=60.0,
                )

                generated = self._extract_generated_text(result)
                if not generated or not generated.strip():
                    raise ValueError("Empty generation result")

                out = {
                    "success": True,
                    "generated_text": generated,
                    "model": model_id,
                    "source": "hf_api",
                    "fallback_used": model_id != config["model_id"],
                    "prompt": prompt,
                    "timestamp": datetime.utcnow().isoformat(),
                }
                self._cache[cache_key] = {"time": time.time(), "data": out}
                return out
            except Exception as e:
                logger.warning(f"⚠️ Generation model {model_id} failed: {e}")
                last_error = e
                continue

        raise Exception(f"Failed to generate text: All models failed. Tried: {models_to_try[:5]}. Last error: {last_error}")

    async def _post_hf_inference(
        self,
        model_id: str,
        payload: Dict[str, Any],
        timeout_seconds: float = 30.0,
    ) -> Any:
        """
        Shared HF inference helper with minimal retry for loading (503) responses.
        """
        _headers = {"Content-Type": "application/json"}
        if self.hf_api_token:
            _headers["Authorization"] = f"Bearer {self.hf_api_token}"

        url = f"{self.hf_api_url}/{model_id}"
        async with httpx.AsyncClient(timeout=timeout_seconds) as client:
            # Try twice: initial + one retry after estimated loading time (if provided)
            for attempt in range(2):
                response = await client.post(url, headers=_headers, json=payload)
                if response.status_code == 503:
                    try:
                        body = response.json()
                    except Exception:
                        body = {}
                    estimated = body.get("estimated_time")
                    if attempt == 0 and isinstance(estimated, (int, float)):
                        await asyncio.sleep(min(float(estimated), 10.0))
                        continue
                response.raise_for_status()
                return response.json()

    def _extract_generated_text(self, result: Any) -> str:
        """
        Normalize various HF text-generation return formats.
        """
        if isinstance(result, list) and result:
            item = result[0]
            if isinstance(item, dict):
                return (
                    item.get("generated_text")
                    or item.get("summary_text")
                    or item.get("text")
                    or ""
                )
            if isinstance(item, str):
                return item
        if isinstance(result, dict):
            return (
                result.get("generated_text")
                or result.get("summary_text")
                or result.get("text")
                or str(result)
            )
        return str(result)
    
    async def _summarize_via_api(
        self,
        text: str
    ) -> Dict[str, Any]:
        """
        Summarize REAL text via HuggingFace API with fallback chain
        Tries at least 3 models before failing
        """
        config = self.model_configs["summarization"]
        models_to_try = [config["model_id"]] + config.get("fallbacks", [])
        
        last_error = None
        for model_id in models_to_try[:5]:  # Try up to 5 models
            try:
                logger.info(f"🔄 Trying summarization model: {model_id}")
                async with httpx.AsyncClient(timeout=30.0) as client:
                    _headers = {"Content-Type": "application/json"}
                    if self.hf_api_token:
                        _headers["Authorization"] = f"Bearer {self.hf_api_token}"
                    response = await client.post(
                        f"{self.hf_api_url}/{model_id}",
                        headers=_headers,
                        json={
                            "inputs": text[:1024],  # Limit input length
                            "parameters": {
                                "max_length": 130,
                                "min_length": 30,
                                "do_sample": False
                            }
                        }
                    )
                    response.raise_for_status()
                    result = response.json()
                
                # Parse result
                if isinstance(result, list) and len(result) > 0:
                    summary = result[0].get("summary_text", "")
                else:
                    summary = result.get("summary_text", str(result))
                
                if summary and len(summary.strip()) > 0:
                    logger.info(f"✅ Summarization succeeded with {model_id}")
                    return {
                        "success": True,
                        "summary": summary,
                        "model": model_id,
                        "source": "hf_api",
                        "fallback_used": model_id != config["model_id"],
                        "timestamp": datetime.utcnow().isoformat()
                    }
            except Exception as e:
                logger.warning(f"⚠️ Summarization model {model_id} failed: {e}")
                last_error = e
                continue
        
        logger.error(f"❌ All summarization models failed. Last error: {last_error}")
        raise Exception(f"Failed to summarize news: All models failed. Tried: {models_to_try[:5]}")
    
    def get_models_list(self) -> Dict[str, Any]:
        """
        Get list of available REAL models
        """
        models_list = []
        for key, config in self.model_configs.items():
            models_list.append({
                "key": key,
                "model_id": config["model_id"],
                "task": config["task"],
                "description": config["description"],
                "loaded_locally": key in self.models,
                "available": True
            })
        
        return {
            "success": True,
            "models": models_list,
            "total": len(models_list),
            "loaded_locally": len(self.models),
            "timestamp": datetime.utcnow().isoformat()
        }


# Global instance
ai_registry = RealAIModelsRegistry()


# Export
__all__ = ["RealAIModelsRegistry", "ai_registry"]