import asyncio
import time
from typing import List, Optional

import numpy as np
import pandas as pd

from src.helpers.prompt_extends import PromptExtendsHelper
from src.modules.corrector_text.repository import CorrectorTextRepository
from src.modules.translate.repository import TranslateRepository
from src.modules.translate.service import TranslateService
from src.modules.sentimental_analyst.models import (
    SentimentalAnalystTextModel,
    SentimentalAnalystAnswerModel,
)


def _default_sentiment_helper_name():
    return 'pysentimiento/robertuito-sentiment-analysis'


class SentimentalAnalystRepository:
    def __init__(self):
        self._prompt_helper = None
        self._helper_name = _default_sentiment_helper_name()
        self._last_used = None

        # Servicios ligeros cargados de inmediato
        self._corrector = CorrectorTextRepository()
        self._translator = TranslateService(repository=TranslateRepository())

    def _load_helper(self):
        """Carga perezosa del PromptExtendsHelper si es necesario, y reprograma descarga."""
        if self._prompt_helper is None:
            self._prompt_helper = PromptExtendsHelper.get_instance(self._helper_name)
        # Actualizar uso y temporizador
        self._last_used = time.time()

    def close(self) -> None:
        PromptExtendsHelper.release_instance(self._helper_name)
        self._prompt_helper = None

    def calculate_average_scores(self, results):
        scores = {}
        for result in results:
            for item in result:
                lbl = item['label'].lower()
                scores.setdefault(lbl, []).append(item['score'])
        avg = {k: (sum(v)/len(v) if v else 0.0) for k, v in scores.items()}
        # Renombrar claves
        out = {}
        if 'negative' in avg:
            out['neg'] = avg.pop('negative')
        if 'neutral' in avg:
            out['neu'] = avg.pop('neutral')
        else:
            out['neu'] = avg.get('neu', 0.0)
        if 'positive' in avg:
            out['pos'] = avg.pop('positive')
        # Mantener otras claves si las hubiera
        out.update(avg)
        return out

    def sentimental_analyst_text(self, chunk):
        self._load_helper()
        return self._prompt_helper.sentiment_pipeline(chunk, batch_size=8)

    def sentimental_analyst_vander(self, comment: str):
        self._load_helper()
        sentiment_results = []
        for chunk in self._prompt_helper.split_text(comment, 100):
            sentiment_results.extend(self.sentimental_analyst_text(chunk))
        puntajes = self.calculate_average_scores(sentiment_results)
        # Cálculo promedio
        puntajes['average'] = puntajes['neg'] * 25 + puntajes['neu'] * 60 + puntajes['pos'] * 100
        # Determinar resultado
        max_attr = max(('neg', 'neu', 'pos'), key=lambda k: puntajes.get(k, 0.0))
        puntajes['result'] = 'invalid' if puntajes['average'] <= 0 else max_attr
        return puntajes

    def resultAnalystText(self, df: pd.DataFrame):
        df = df.apply(self.translateText, axis=1)
        df = df.apply(self._corrector.validate_query_answer, axis=1)
        df = self.sentimentlAnalyst(df)
        df = df.apply(self.formatRow, axis=1)
        return df.drop(columns=['before']).reset_index(drop=True)

    async def _translate_column_async(self, df: pd.DataFrame, key: str = 'text') -> pd.DataFrame:
        """Traduce async la columna `key` al español y guarda el original en 'before'."""
        texts = df[key].tolist()
        results = await asyncio.gather(*[
            self._translator.translate(text=t, language='es') for t in texts
        ])
        df['before'] = df[key]
        df[key] = [r.translated_text for r in results]
        return df

    async def resultAnalystTextAsync(
        self, df: pd.DataFrame, validate: Optional[bool] = True
    ) -> List[SentimentalAnalystTextModel]:
        df['validated'] = validate
        # 1. Traducción async en bloque (evita llamar async dentro de pandas apply)
        df = await self._translate_column_async(df, key='text')
        # 2. Validación sync mediante CorrectorTextRepository
        df = df.apply(self._corrector.validate_query_answer, axis=1)
        # 3. Análisis de sentimiento async
        df = await self.sentimentlAnalystAsync(df)
        # 4. Restaurar texto original
        df['text'] = df['before']
        df = df.drop(columns=['before']).reset_index(drop=True)
        return [
            SentimentalAnalystTextModel(
                text=row['text'],
                neg=row['neg'],
                neu=row['neu'],
                pos=row['pos'],
                average=row['average'],
                result=row['result'],
                is_valid=bool(row['is_valid']),
                validated=bool(row['validated']),
            )
            for _, row in df.iterrows()
        ]

    async def resultAnalystAnswerAsync(
        self, df: pd.DataFrame, validate: Optional[bool] = True
    ) -> List[SentimentalAnalystAnswerModel]:
        df['validated'] = validate
        # 1. Traducción async en bloque
        df = await self._translate_column_async(df, key='answer')
        # 2. Validación sync
        df = df.apply(self._corrector.validate_query_answer, axis=1)
        # 3. Análisis de sentimiento async
        df = await self.sentimentlAnalystAsync(df, 'answer')
        # 4. Restaurar texto original
        df['answer'] = df['before']
        df = df.drop(columns=['before']).reset_index(drop=True)
        return [
            SentimentalAnalystAnswerModel(
                answer=row['answer'],
                neg=row['neg'],
                neu=row['neu'],
                pos=row['pos'],
                average=row['average'],
                result=row['result'],
                is_valid=bool(row['is_valid']),
                validated=bool(row['validated']),
            )
            for _, row in df.iterrows()
        ]

    def resultAnalystAnswer(self, df: pd.DataFrame):
        df = df.apply(self.translateText, axis=1, key='answer')
        df = df.apply(self._corrector.validate_query_answer, axis=1)
        df = self.sentimentlAnalyst(df, 'answer')
        df = df.apply(self.formatRow, axis=1, key='answer')
        return df.drop(columns=['before']).reset_index(drop=True)

    def translateText(self, row, key: str = 'text'):
        row['before'] = row.get(key)
        row[key] = self._translator.translate(row[key], 'es')
        return row

    def formatRow(self, row, key: str = 'text'):
        row[key] = row.get('before')
        row['before'] = np.nan
        return row

    def sentimentlAnalyst(self, df: pd.DataFrame, key='text'):
        new_cols = df.apply(self.get_sentiment, axis=1).apply(pd.Series)
        return pd.concat([df, new_cols], axis=1)

    async def sentimentlAnalystAsync(self, df: pd.DataFrame, key='text'):
        loop = asyncio.get_event_loop()
        sem = asyncio.Semaphore(150)

        async def task(row):
            async with sem:
                return await self.get_sentiment_async(row, key)

        tasks = [loop.create_task(task(row)) for _, row in df.iterrows()]
        results = await asyncio.gather(*tasks)
        new_cols = pd.DataFrame(results, index=df.index)
        return pd.concat([df, new_cols], axis=1)

    def get_sentiment(self, row):
        if row.get('is_valid'):
            return self.sentimental_analyst_vander(row.get('text'))
        return {'neg': 0, 'neu': 0, 'pos': 0, 'average': 0, 'result': 'invalid'}

    async def get_sentiment_async(self, row, key='text'):
        if row.get('is_valid') and isinstance(row.get(key), str):
            return self.sentimental_analyst_vander(row[key])
        return {'neg': 0, 'neu': 0, 'pos': 0, 'average': 0, 'result': 'invalid'}

    def categorize(self, score):
        if score <= 0:
            return 'invalid'
        if score <= 25:
            return 'neg'
        if score <= 60:
            return 'neu'
        return 'pos'