#!/usr/bin/env python3
"""
LLM Signal Logger
=================

Logs ML signals and LLM decisions to PostgreSQL for:
- Feedback and performance tracking
- Fine-tuning dataset generation
- Trade result correlation

Uses the ml.llm_signals table in PostgreSQL.

Author: ML-Specialist (NEXUS v4.0)
Version: 1.0.0
Created: 2026-01-25
"""

import os
import json
import uuid
from datetime import datetime
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, asdict
from pathlib import Path

import pandas as pd
from loguru import logger

try:
    from sqlalchemy import create_engine, text
    from sqlalchemy.pool import QueuePool
    HAS_SQLALCHEMY = True
except ImportError:
    HAS_SQLALCHEMY = False
    logger.warning("SQLAlchemy not available - database logging disabled")


# ============================================================
# DATA CLASSES
# ============================================================

@dataclass
class TradeResult:
    """Result of a trade for logging"""
    signal_id: str
    result: str  # "WIN", "LOSS", "BREAKEVEN", "PARTIAL"
    pnl: float  # P&L in USD
    pnl_pct: float  # P&L as percentage
    entry_price: float
    exit_price: float
    stop_loss_hit: bool
    take_profit_hit: bool
    duration_minutes: int
    exit_reason: str  # "TP1", "TP2", "SL", "MANUAL", "TRAILING"
    notes: Optional[str] = None


# ============================================================
# SIGNAL LOGGER CLASS
# ============================================================

class LLMSignalLogger:
    """
    Logger for ML signals and LLM decisions.

    Persists signals to PostgreSQL for:
    - Performance tracking and feedback
    - Trade result correlation
    - Fine-tuning dataset generation

    Usage:
        logger = LLMSignalLogger()

        # Log a signal
        signal_id = logger.log_signal(
            signal=formatted_signal,
            decision=trading_decision
        )

        # Update with trade result
        logger.log_trade_result(signal_id, trade_result)

        # Export for fine-tuning
        logger.export_for_finetuning("training_data.jsonl")
    """

    # Default connection parameters
    DEFAULT_HOST = "localhost"
    DEFAULT_PORT = 5432
    DEFAULT_DATABASE = "trading_platform"
    DEFAULT_USER = "trading_user"
    DEFAULT_PASSWORD = "trading_dev_2026"

    def __init__(
        self,
        host: Optional[str] = None,
        port: Optional[int] = None,
        database: Optional[str] = None,
        user: Optional[str] = None,
        password: Optional[str] = None,
        create_table: bool = True
    ):
        """
        Initialize the signal logger.

        Args:
            host: Database host (default: localhost)
            port: Database port (default: 5432)
            database: Database name (default: trading_platform)
            user: Database user (default: trading_user)
            password: Database password
            create_table: Create table if it doesn't exist
        """
        self.host = host or os.getenv("DB_HOST", self.DEFAULT_HOST)
        self.port = port or int(os.getenv("DB_PORT", self.DEFAULT_PORT))
        self.database = database or os.getenv("DB_NAME", self.DEFAULT_DATABASE)
        self.user = user or os.getenv("DB_USER", self.DEFAULT_USER)
        self.password = password or os.getenv("DB_PASSWORD", self.DEFAULT_PASSWORD)

        self._engine = None
        self._connected = False

        if HAS_SQLALCHEMY:
            self._connect()
            if create_table and self._connected:
                self._ensure_table_exists()
        else:
            logger.warning("SQLAlchemy not available - using file-based fallback")
            self._fallback_dir = Path("logs/llm_signals")
            self._fallback_dir.mkdir(parents=True, exist_ok=True)

    def _connect(self) -> None:
        """Establish database connection"""
        try:
            connection_string = (
                f"postgresql://{self.user}:{self.password}"
                f"@{self.host}:{self.port}/{self.database}"
            )

            self._engine = create_engine(
                connection_string,
                poolclass=QueuePool,
                pool_size=5,
                max_overflow=10,
                pool_pre_ping=True,
                echo=False
            )

            # Test connection
            with self._engine.connect() as conn:
                conn.execute(text("SELECT 1"))

            self._connected = True
            logger.info(f"Connected to PostgreSQL: {self.host}:{self.port}/{self.database}")

        except Exception as e:
            logger.error(f"Failed to connect to database: {e}")
            self._connected = False

    def _ensure_table_exists(self) -> None:
        """Create the ml.llm_signals table if it doesn't exist"""
        create_schema_sql = "CREATE SCHEMA IF NOT EXISTS ml;"

        create_table_sql = """
        CREATE TABLE IF NOT EXISTS ml.llm_signals (
            id SERIAL PRIMARY KEY,
            signal_id UUID NOT NULL UNIQUE,
            symbol VARCHAR(20) NOT NULL,
            timestamp TIMESTAMPTZ NOT NULL,
            ml_prediction JSONB,
            strategy_predictions JSONB,
            llm_prompt TEXT,
            llm_response TEXT,
            decision JSONB,
            trade_result JSONB,
            created_at TIMESTAMPTZ DEFAULT NOW(),
            updated_at TIMESTAMPTZ DEFAULT NOW()
        );

        CREATE INDEX IF NOT EXISTS idx_llm_signals_symbol ON ml.llm_signals(symbol);
        CREATE INDEX IF NOT EXISTS idx_llm_signals_timestamp ON ml.llm_signals(timestamp);
        CREATE INDEX IF NOT EXISTS idx_llm_signals_created_at ON ml.llm_signals(created_at);
        """

        try:
            with self._engine.connect() as conn:
                conn.execute(text(create_schema_sql))
                conn.execute(text(create_table_sql))
                conn.commit()
            logger.info("ml.llm_signals table ready")
        except Exception as e:
            logger.error(f"Failed to create table: {e}")

    def log_signal(
        self,
        signal: Any,
        decision: Any,
        llm_prompt: Optional[str] = None,
        llm_response: Optional[str] = None
    ) -> str:
        """
        Log a signal and its decision.

        Args:
            signal: FormattedSignal or dict with signal data
            decision: TradingDecision or dict with decision data
            llm_prompt: The prompt sent to the LLM (optional)
            llm_response: Raw LLM response (optional)

        Returns:
            signal_id (UUID string) for later reference
        """
        signal_id = str(uuid.uuid4())

        # Convert to dicts if needed
        signal_dict = signal.to_dict() if hasattr(signal, "to_dict") else dict(signal)
        decision_dict = decision.to_dict() if hasattr(decision, "to_dict") else dict(decision)

        # Extract key fields
        symbol = signal_dict.get("symbol", "UNKNOWN")
        timestamp = signal_dict.get("timestamp", datetime.utcnow().isoformat())

        # Separate ML prediction and strategy predictions
        ml_prediction = {
            "direction": signal_dict.get("direction"),
            "magnitude": signal_dict.get("magnitude_value"),
            "confidence": signal_dict.get("confidence_value"),
            "delta_high": signal_dict.get("delta_high"),
            "delta_low": signal_dict.get("delta_low"),
            "volatility_regime": signal_dict.get("volatility_regime"),
            "market_phase": signal_dict.get("market_phase"),
        }

        strategy_predictions = signal_dict.get("strategies", [])

        if self._connected and self._engine:
            return self._log_to_database(
                signal_id=signal_id,
                symbol=symbol,
                timestamp=timestamp,
                ml_prediction=ml_prediction,
                strategy_predictions=strategy_predictions,
                llm_prompt=llm_prompt,
                llm_response=llm_response,
                decision=decision_dict
            )
        else:
            return self._log_to_file(
                signal_id=signal_id,
                symbol=symbol,
                timestamp=timestamp,
                ml_prediction=ml_prediction,
                strategy_predictions=strategy_predictions,
                llm_prompt=llm_prompt,
                llm_response=llm_response,
                decision=decision_dict
            )

    def _log_to_database(
        self,
        signal_id: str,
        symbol: str,
        timestamp: str,
        ml_prediction: Dict,
        strategy_predictions: List,
        llm_prompt: Optional[str],
        llm_response: Optional[str],
        decision: Dict
    ) -> str:
        """Log signal to PostgreSQL database"""
        insert_sql = """
        INSERT INTO ml.llm_signals (
            signal_id, symbol, timestamp, ml_prediction,
            strategy_predictions, llm_prompt, llm_response, decision
        ) VALUES (
            :signal_id, :symbol, :timestamp, :ml_prediction::jsonb,
            :strategy_predictions::jsonb, :llm_prompt, :llm_response, :decision::jsonb
        )
        """

        try:
            with self._engine.connect() as conn:
                conn.execute(
                    text(insert_sql),
                    {
                        "signal_id": signal_id,
                        "symbol": symbol,
                        "timestamp": timestamp,
                        "ml_prediction": json.dumps(ml_prediction),
                        "strategy_predictions": json.dumps(strategy_predictions),
                        "llm_prompt": llm_prompt,
                        "llm_response": llm_response,
                        "decision": json.dumps(decision)
                    }
                )
                conn.commit()

            logger.debug(f"Logged signal {signal_id} for {symbol}")
            return signal_id

        except Exception as e:
            logger.error(f"Failed to log signal: {e}")
            # Fall back to file
            return self._log_to_file(
                signal_id, symbol, timestamp, ml_prediction,
                strategy_predictions, llm_prompt, llm_response, decision
            )

    def _log_to_file(
        self,
        signal_id: str,
        symbol: str,
        timestamp: str,
        ml_prediction: Dict,
        strategy_predictions: List,
        llm_prompt: Optional[str],
        llm_response: Optional[str],
        decision: Dict
    ) -> str:
        """Fallback: log signal to JSONL file"""
        record = {
            "signal_id": signal_id,
            "symbol": symbol,
            "timestamp": timestamp,
            "ml_prediction": ml_prediction,
            "strategy_predictions": strategy_predictions,
            "llm_prompt": llm_prompt,
            "llm_response": llm_response,
            "decision": decision,
            "created_at": datetime.utcnow().isoformat()
        }

        filepath = self._fallback_dir / f"signals_{datetime.utcnow().strftime('%Y%m%d')}.jsonl"

        with open(filepath, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, default=str) + "\n")

        logger.debug(f"Logged signal {signal_id} to file")
        return signal_id

    def log_trade_result(
        self,
        signal_id: str,
        result: TradeResult
    ) -> bool:
        """
        Update a signal with its trade result.

        Args:
            signal_id: The signal_id returned from log_signal
            result: TradeResult with outcome data

        Returns:
            True if successful
        """
        result_dict = asdict(result)

        if self._connected and self._engine:
            update_sql = """
            UPDATE ml.llm_signals
            SET trade_result = :trade_result::jsonb,
                updated_at = NOW()
            WHERE signal_id = :signal_id
            """

            try:
                with self._engine.connect() as conn:
                    result_obj = conn.execute(
                        text(update_sql),
                        {
                            "signal_id": signal_id,
                            "trade_result": json.dumps(result_dict)
                        }
                    )
                    conn.commit()

                if result_obj.rowcount > 0:
                    logger.info(f"Updated signal {signal_id} with trade result: {result.result}")
                    return True
                else:
                    logger.warning(f"Signal {signal_id} not found for update")
                    return False

            except Exception as e:
                logger.error(f"Failed to update trade result: {e}")
                return False
        else:
            # File fallback - append to a results file
            result_record = {
                "signal_id": signal_id,
                "trade_result": result_dict,
                "updated_at": datetime.utcnow().isoformat()
            }

            filepath = self._fallback_dir / "trade_results.jsonl"
            with open(filepath, "a", encoding="utf-8") as f:
                f.write(json.dumps(result_record, default=str) + "\n")

            return True

    def get_signals_for_training(
        self,
        symbol: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
        only_with_results: bool = True
    ) -> pd.DataFrame:
        """
        Get signals for training/fine-tuning.

        Args:
            symbol: Filter by symbol (optional)
            start_date: Start date filter (YYYY-MM-DD)
            end_date: End date filter (YYYY-MM-DD)
            only_with_results: Only include signals with trade results

        Returns:
            DataFrame with signals and results
        """
        if not self._connected or not self._engine:
            logger.warning("Database not connected - returning empty DataFrame")
            return pd.DataFrame()

        query = """
        SELECT
            signal_id,
            symbol,
            timestamp,
            ml_prediction,
            strategy_predictions,
            llm_prompt,
            llm_response,
            decision,
            trade_result,
            created_at
        FROM ml.llm_signals
        WHERE 1=1
        """

        params = {}

        if symbol:
            query += " AND symbol = :symbol"
            params["symbol"] = symbol

        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

        if only_with_results:
            query += " AND trade_result IS NOT NULL"

        query += " ORDER BY timestamp ASC"

        try:
            df = pd.read_sql(text(query), self._engine, params=params)
            logger.info(f"Retrieved {len(df)} signals for training")
            return df
        except Exception as e:
            logger.error(f"Failed to get signals: {e}")
            return pd.DataFrame()

    def export_for_finetuning(
        self,
        output_path: str,
        format: str = "jsonl",
        symbol: Optional[str] = None,
        only_successful: bool = False
    ) -> Path:
        """
        Export signals in a format suitable for LLM fine-tuning.

        Args:
            output_path: Output file path
            format: Output format ("jsonl", "openai", "anthropic")
            symbol: Filter by symbol (optional)
            only_successful: Only include winning trades

        Returns:
            Path to the output file
        """
        df = self.get_signals_for_training(symbol=symbol, only_with_results=True)

        if df.empty:
            logger.warning("No signals to export")
            return Path(output_path)

        # Filter successful if requested
        if only_successful:
            df = df[df["trade_result"].apply(
                lambda x: x.get("result") == "WIN" if isinstance(x, dict) else False
            )]

        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        if format == "openai":
            return self._export_openai_format(df, output_path)
        elif format == "anthropic":
            return self._export_anthropic_format(df, output_path)
        else:
            return self._export_jsonl_format(df, output_path)

    def _export_jsonl_format(self, df: pd.DataFrame, output_path: Path) -> Path:
        """Export in generic JSONL format"""
        with open(output_path, "w", encoding="utf-8") as f:
            for _, row in df.iterrows():
                record = {
                    "id": str(row["signal_id"]),
                    "symbol": row["symbol"],
                    "timestamp": str(row["timestamp"]),
                    "prompt": row["llm_prompt"],
                    "response": row["llm_response"],
                    "decision": row["decision"],
                    "result": row["trade_result"]
                }
                f.write(json.dumps(record, default=str) + "\n")

        logger.info(f"Exported {len(df)} signals to {output_path}")
        return output_path

    def _export_openai_format(self, df: pd.DataFrame, output_path: Path) -> Path:
        """Export in OpenAI fine-tuning format"""
        from .prompts.trading_decision import SYSTEM_PROMPT

        with open(output_path, "w", encoding="utf-8") as f:
            for _, row in df.iterrows():
                messages = [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": row["llm_prompt"] or ""},
                    {"role": "assistant", "content": row["llm_response"] or ""}
                ]
                record = {"messages": messages}
                f.write(json.dumps(record, default=str) + "\n")

        logger.info(f"Exported {len(df)} signals in OpenAI format to {output_path}")
        return output_path

    def _export_anthropic_format(self, df: pd.DataFrame, output_path: Path) -> Path:
        """Export in Anthropic fine-tuning format"""
        from .prompts.trading_decision import SYSTEM_PROMPT

        with open(output_path, "w", encoding="utf-8") as f:
            for _, row in df.iterrows():
                record = {
                    "system": SYSTEM_PROMPT,
                    "messages": [
                        {"role": "user", "content": row["llm_prompt"] or ""},
                        {"role": "assistant", "content": row["llm_response"] or ""}
                    ]
                }
                f.write(json.dumps(record, default=str) + "\n")

        logger.info(f"Exported {len(df)} signals in Anthropic format to {output_path}")
        return output_path

    def get_statistics(
        self,
        symbol: Optional[str] = None,
        days: int = 30
    ) -> Dict[str, Any]:
        """
        Get statistics about logged signals.

        Args:
            symbol: Filter by symbol (optional)
            days: Number of days to look back

        Returns:
            Dictionary with statistics
        """
        if not self._connected or not self._engine:
            return {"error": "Database not connected"}

        query = """
        SELECT
            COUNT(*) as total_signals,
            COUNT(trade_result) as signals_with_results,
            COUNT(CASE WHEN (trade_result->>'result')::text = 'WIN' THEN 1 END) as wins,
            COUNT(CASE WHEN (trade_result->>'result')::text = 'LOSS' THEN 1 END) as losses,
            AVG((trade_result->>'pnl')::float) as avg_pnl
        FROM ml.llm_signals
        WHERE timestamp >= NOW() - INTERVAL '{days} days'
        """

        params = {"days": days}

        if symbol:
            query = query.replace("WHERE", "WHERE symbol = :symbol AND")
            params["symbol"] = symbol

        try:
            with self._engine.connect() as conn:
                result = conn.execute(text(query.format(days=days)), params)
                row = result.fetchone()

            if row:
                total = row[0] or 0
                with_results = row[1] or 0
                wins = row[2] or 0
                losses = row[3] or 0
                avg_pnl = row[4] or 0

                return {
                    "total_signals": total,
                    "signals_with_results": with_results,
                    "wins": wins,
                    "losses": losses,
                    "win_rate": wins / with_results if with_results > 0 else 0,
                    "avg_pnl": avg_pnl,
                    "period_days": days
                }
            else:
                return {"total_signals": 0}

        except Exception as e:
            logger.error(f"Failed to get statistics: {e}")
            return {"error": str(e)}


# ============================================================
# EXPORTS
# ============================================================

__all__ = [
    'LLMSignalLogger',
    'TradeResult',
]
