# src/data/storage.py
"""
Módulo para almacenamiento persistente de datos en PostgreSQL y caché en Redis
"""
import os
import pandas as pd

from sqlalchemy import (
    create_engine, Column, String, Float,
    DateTime, Integer, Index, text
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.dialects.postgresql import insert as pg_insert
from datetime import datetime, timedelta, timezone
from typing import Optional
import redis
from ..utils.logger import log

Base = declarative_base()


class OHLCV(Base):
    __tablename__ = 'ohlcv'

    id = Column(Integer, primary_key=True)
    timestamp = Column(DateTime, nullable=False)
    symbol = Column(String(20), nullable=False)
    timeframe = Column(String(10), nullable=False)

    open = Column(Float, nullable=False)
    high = Column(Float, nullable=False)
    low = Column(Float, nullable=False)
    close = Column(Float, nullable=False)
    volume = Column(Float, nullable=False)

    returns = Column(Float)
    log_returns = Column(Float)
    adx = Column(Float)

    __table_args__ = (
        Index('idx_symbol_tf_ts', 'symbol', 'timeframe', 'timestamp', unique=True),
    )


class StorageManager:
    """
    Gestor de almacenamiento con PostgreSQL y Redis
    """

    def __init__(
        self,
        db_host: str,
        db_port: int,
        db_name: str,
        db_user: str,
        db_password: str,
        redis_host: str = 'localhost',
        redis_port: int = 6379,
        redis_db: int = 0
    ):
        self.db_url = (
            f"postgresql://{db_user}:{db_password}"
            f"@{db_host}:{db_port}/{db_name}"
        )
        self.engine = create_engine(self.db_url, echo=False)

        Base.metadata.create_all(self.engine)

        Session = sessionmaker(bind=self.engine)
        self.session = Session()

        log.success("Conectado a PostgreSQL")

        try:
            self.redis_client = redis.Redis(
                host=redis_host,
                port=redis_port,
                db=redis_db,
                decode_responses=True
            )
            self.redis_client.ping()
            log.success("Conectado a Redis")
        except Exception:
            self.redis_client = None
            log.warning("Redis no disponible")

    # ------------------------------------------------------------------

    @classmethod
    def from_env(cls):
        """
        Crea StorageManager leyendo configuración desde variables de entorno.
        """
        return cls(
            db_host=os.getenv("DB_HOST"),
            db_port=int(os.getenv("DB_PORT", "5432")),
            db_name=os.getenv("DB_NAME"),
            db_user=os.getenv("DB_USER"),
            db_password=os.getenv("DB_PASSWORD"),
            redis_host=os.getenv("REDIS_HOST", "localhost"),
            redis_port=int(os.getenv("REDIS_PORT", "6379")),
            redis_db=int(os.getenv("REDIS_DB", "0")),
        )

    # ------------------------------------------------------------------

    def _timeframe_delta(self, timeframe: str) -> timedelta:
        if timeframe.endswith("m"):
            return timedelta(minutes=int(timeframe.replace("m", "")))
        if timeframe.endswith("h"):
            return timedelta(hours=int(timeframe.replace("h", "")))
        if timeframe.endswith("d"):
            return timedelta(days=int(timeframe.replace("d", "")))
        raise NotImplementedError(f"Timeframe no soportado: {timeframe}")

    # ------------------------------------------------------------------

    def save_ohlcv(self, df: pd.DataFrame) -> int:
        """
        Guarda datos OHLCV en PostgreSQL con UPSERT:
        - Inserta filas nuevas
        - Ignora duplicados (ON CONFLICT DO NOTHING)
        """
        if df is None or df.empty:
            log.warning("DataFrame vacío, nada que guardar")
            return 0

        df_to_save = df.reset_index()

        if df_to_save.columns[0] != 'timestamp':
            df_to_save.rename(
                columns={df_to_save.columns[0]: 'timestamp'},
                inplace=True
            )

        allowed_columns = [
            'timestamp', 'symbol', 'timeframe',
            'open', 'high', 'low', 'close', 'volume',
            'returns', 'log_returns', 'adx'
        ]
        df_to_save = df_to_save[
            [c for c in allowed_columns if c in df_to_save.columns]
        ]

        ts = pd.to_datetime(df_to_save["timestamp"], errors="coerce")
        if getattr(ts.dt, "tz", None) is not None:
            ts = ts.dt.tz_convert("UTC").dt.tz_localize(None)
        df_to_save["timestamp"] = ts

        df_to_save = df_to_save.dropna(subset=["timestamp"])

        if df_to_save.empty:
            log.warning("Tras normalizar timestamps, no quedan filas para guardar")
            return 0

        records = df_to_save.to_dict(orient="records")
        log.info(f"Guardando {len(records)} registros (UPSERT / ignore duplicates)")

        inserted_total = 0
        table = OHLCV.__table__

        chunk_size = 2000
        with self.engine.begin() as conn:
            for i in range(0, len(records), chunk_size):
                chunk = records[i:i + chunk_size]

                stmt = pg_insert(table).values(chunk).on_conflict_do_nothing(
                    index_elements=["symbol", "timeframe", "timestamp"]
                )
                res = conn.execute(stmt)
                inserted_total += (res.rowcount or 0)

        log.success(f"Guardadas {inserted_total} filas nuevas (duplicados ignorados)")
        return inserted_total

    # ------------------------------------------------------------------

    def load_ohlcv(
        self,
        symbol: str,
        timeframe: str,
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None,
        use_cache: bool = True
    ) -> pd.DataFrame:

        query = """
            SELECT *
            FROM ohlcv
            WHERE symbol = :symbol AND timeframe = :timeframe
        """

        params = {'symbol': symbol, 'timeframe': timeframe}

        if start_date:
            query += " AND timestamp >= :start_date"
            params['start_date'] = start_date

        if end_date:
            query += " AND timestamp <= :end_date"
            params['end_date'] = end_date

        query += " ORDER BY timestamp ASC"

        with self.engine.connect() as conn:
            df = pd.read_sql(text(query), conn, params=params)

        if df.empty:
            log.warning(f"No se encontraron datos para {symbol} {timeframe}")
            return df

        df['timestamp'] = pd.to_datetime(df['timestamp'])
        df.set_index('timestamp', inplace=True)

        log.info(f"Cargados {len(df)} registros de {symbol} {timeframe}")
        return df

    # ------------------------------------------------------------------

    def get_latest_timestamp(self, symbol: str, timeframe: str) -> Optional[datetime]:
        query = """
            SELECT MAX(timestamp) AS last_timestamp
            FROM ohlcv
            WHERE symbol = :symbol AND timeframe = :timeframe
        """

        with self.engine.connect() as conn:
            result = conn.execute(
                text(query),
                {'symbol': symbol, 'timeframe': timeframe}
            ).fetchone()

        return result[0] if result and result[0] else None

    # ------------------------------------------------------------------

    def get_available_data(self) -> pd.DataFrame:
        query = """
            SELECT
                symbol,
                timeframe,
                MIN(timestamp) AS first_date,
                MAX(timestamp) AS last_date,
                COUNT(*) AS record_count
            FROM ohlcv
            GROUP BY symbol, timeframe
            ORDER BY symbol, timeframe
        """

        with self.engine.connect() as conn:
            return pd.read_sql(text(query), conn)

    # ------------------------------------------------------------------

    def sync_ohlcv(
        self,
        *,
        symbol: str,
        timeframe: str,
        downloader,
    ) -> int:

        last_ts = self.get_latest_timestamp(symbol, timeframe)
        now = datetime.now(timezone.utc)
        delta = self._timeframe_delta(timeframe)

        if timeframe.endswith("m"):
            expected_last = (
                now.replace(second=0, microsecond=0) - delta
            )
        elif timeframe.endswith("h"):
            expected_last = (
                now.replace(minute=0, second=0, microsecond=0) - delta
            )
        elif timeframe.endswith("d"):
            expected_last = (
                now.replace(hour=0, minute=0, second=0, microsecond=0) - delta
            )
        else:
            raise NotImplementedError(f"Timeframe no soportado: {timeframe}")

        if last_ts is None:
            log.warning("DB vacía → descarga completa")
            df = downloader.download_full(symbol, timeframe)
            return self.save_ohlcv(df)

        last_ts = last_ts.replace(tzinfo=timezone.utc)

        if last_ts >= expected_last:
            log.info("OHLCV ya está actualizado")
            return 0

        start = last_ts + delta

        log.info(f"Descargando OHLCV faltante:\n{start} → {expected_last}")

        df = downloader.download_range(
            symbol=symbol,
            timeframe=timeframe,
            start=start,
            end=expected_last,
            storage=self,
        )

        return self.save_ohlcv(df)

    # ------------------------------------------------------------------

    def load_ohlcv_tail(
        self,
        *,
        symbol: str,
        timeframe: str,
        end_ts: datetime,
        limit: int,
    ) -> pd.DataFrame:

        query = """
            SELECT *
            FROM ohlcv
            WHERE symbol = :symbol
              AND timeframe = :timeframe
              AND timestamp < :end_ts
            ORDER BY timestamp DESC
            LIMIT :limit
        """

        with self.engine.connect() as conn:
            df = pd.read_sql(
                text(query),
                conn,
                params={
                    "symbol": symbol,
                    "timeframe": timeframe,
                    "end_ts": end_ts,
                    "limit": limit,
                },
            )

        if df.empty:
            return df

        df["timestamp"] = pd.to_datetime(df["timestamp"])
        df = df.sort_values("timestamp")
        df.set_index("timestamp", inplace=True)
        return df

    # ------------------------------------------------------------------

    def close(self):
        self.session.close()
        self.engine.dispose()
        if self.redis_client:
            self.redis_client.close()
        log.info("Conexiones cerradas")
