"""
Data persistence module using SQLite
"""

import sqlite3
import json
from typing import List, Optional, Dict, Any
from datetime import datetime
from io import StringIO
import pandas as pd

from openfund_agent.models.trading import Trade
from openfund_agent.models.market_data import MarketData


class DataStore:
    """
    Data storage manager using SQLite for persistence

    Manages storage and retrieval of trading records and market data.

    Attributes:
        db_path: Path to SQLite database file
        conn: SQLite connection object
    """

    def __init__(self, db_path: str):
        """
        Initialize data store with SQLite database

        Args:
            db_path: Path to SQLite database file
        """
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.conn.row_factory = sqlite3.Row
        self._init_tables()

    def _init_tables(self) -> None:
        """
        Initialize database tables

        Creates tables for trades and market data if they don't exist.
        """
        cursor = self.conn.cursor()

        # Create trades table
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS trades (
                trade_id TEXT PRIMARY KEY,
                symbol TEXT NOT NULL,
                side TEXT NOT NULL,
                entry_time INTEGER NOT NULL,
                entry_price REAL NOT NULL,
                exit_time INTEGER,
                exit_price REAL,
                quantity REAL NOT NULL,
                stop_loss REAL NOT NULL,
                take_profit REAL NOT NULL,
                status TEXT NOT NULL,
                pnl REAL,
                fees REAL NOT NULL,
                strategy_name TEXT NOT NULL,
                metadata TEXT,
                created_at INTEGER NOT NULL
            )
        """
        )

        # Create indexes for trades table
        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_trades_symbol 
            ON trades(symbol)
        """
        )

        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_trades_entry_time 
            ON trades(entry_time)
        """
        )

        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_trades_status 
            ON trades(status)
        """
        )

        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_trades_strategy 
            ON trades(strategy_name)
        """
        )

        # Create market_data table
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS market_data (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                symbol TEXT NOT NULL,
                timeframe TEXT NOT NULL,
                klines_data TEXT NOT NULL,
                indicators TEXT NOT NULL,
                timestamp INTEGER NOT NULL,
                created_at INTEGER NOT NULL
            )
        """
        )

        # Create indexes for market_data table
        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_market_data_symbol 
            ON market_data(symbol)
        """
        )

        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_market_data_timeframe 
            ON market_data(timeframe)
        """
        )

        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_market_data_timestamp 
            ON market_data(timestamp)
        """
        )

        self.conn.commit()

    def save_trade(self, trade: Trade) -> None:
        """
        Save trade record to database

        Args:
            trade: Trade object to save

        Raises:
            sqlite3.Error: If database operation fails
        """
        cursor = self.conn.cursor()

        # Serialize metadata to JSON
        metadata_json = json.dumps(trade.metadata) if trade.metadata else "{}"

        # Get current timestamp
        created_at = int(datetime.now().timestamp())

        cursor.execute(
            """
            INSERT OR REPLACE INTO trades (
                trade_id, symbol, side, entry_time, entry_price,
                exit_time, exit_price, quantity, stop_loss, take_profit,
                status, pnl, fees, strategy_name, metadata, created_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
            (
                trade.trade_id,
                trade.symbol,
                trade.side,
                trade.entry_time,
                trade.entry_price,
                trade.exit_time,
                trade.exit_price,
                trade.quantity,
                trade.stop_loss,
                trade.take_profit,
                trade.status,
                trade.pnl,
                trade.fees,
                trade.strategy_name,
                metadata_json,
                created_at,
            ),
        )

        self.conn.commit()

    def get_trades(
        self, start_time: int, end_time: int, symbol: Optional[str] = None
    ) -> List[Trade]:
        """
        Query trade records from database

        Args:
            start_time: Start timestamp (inclusive)
            end_time: End timestamp (inclusive)
            symbol: Optional symbol filter

        Returns:
            List of Trade objects matching the criteria

        Raises:
            sqlite3.Error: If database operation fails
        """
        cursor = self.conn.cursor()

        if symbol:
            cursor.execute(
                """
                SELECT * FROM trades
                WHERE entry_time >= ? AND entry_time <= ?
                AND symbol = ?
                ORDER BY entry_time DESC
            """,
                (start_time, end_time, symbol),
            )
        else:
            cursor.execute(
                """
                SELECT * FROM trades
                WHERE entry_time >= ? AND entry_time <= ?
                ORDER BY entry_time DESC
            """,
                (start_time, end_time),
            )

        rows = cursor.fetchall()

        trades = []
        for row in rows:
            # Deserialize metadata from JSON
            metadata = json.loads(row["metadata"]) if row["metadata"] else {}

            trade = Trade(
                trade_id=row["trade_id"],
                symbol=row["symbol"],
                side=row["side"],
                entry_time=row["entry_time"],
                entry_price=row["entry_price"],
                exit_time=row["exit_time"],
                exit_price=row["exit_price"],
                quantity=row["quantity"],
                stop_loss=row["stop_loss"],
                take_profit=row["take_profit"],
                status=row["status"],
                pnl=row["pnl"],
                fees=row["fees"],
                strategy_name=row["strategy_name"],
                metadata=metadata,
            )
            trades.append(trade)

        return trades

    def save_market_data(self, data: MarketData) -> None:
        """
        Save market data to database

        Args:
            data: MarketData object to save

        Raises:
            sqlite3.Error: If database operation fails
        """
        cursor = self.conn.cursor()

        # Serialize klines DataFrame to JSON
        klines_json = data.klines.to_json(orient="records", date_format="iso")

        # Serialize indicators to JSON
        indicators_json = json.dumps(data.indicators)

        # Get current timestamp
        created_at = int(datetime.now().timestamp())

        cursor.execute(
            """
            INSERT INTO market_data (
                symbol, timeframe, klines_data, indicators, timestamp, created_at
            ) VALUES (?, ?, ?, ?, ?, ?)
        """,
            (data.symbol, data.timeframe, klines_json, indicators_json, data.timestamp, created_at),
        )

        self.conn.commit()

    def get_market_data(
        self,
        symbol: str,
        timeframe: str,
        start_time: Optional[int] = None,
        end_time: Optional[int] = None,
        limit: int = 100,
    ) -> List[MarketData]:
        """
        Query market data from database

        Args:
            symbol: Trading pair symbol
            timeframe: Timeframe (e.g., "1h", "4h")
            start_time: Optional start timestamp filter
            end_time: Optional end timestamp filter
            limit: Maximum number of records to return

        Returns:
            List of MarketData objects matching the criteria

        Raises:
            sqlite3.Error: If database operation fails
        """
        cursor = self.conn.cursor()

        query = """
            SELECT * FROM market_data
            WHERE symbol = ? AND timeframe = ?
        """
        params = [symbol, timeframe]

        if start_time is not None:
            query += " AND timestamp >= ?"
            params.append(start_time)

        if end_time is not None:
            query += " AND timestamp <= ?"
            params.append(end_time)

        query += " ORDER BY timestamp DESC LIMIT ?"
        params.append(limit)

        cursor.execute(query, params)
        rows = cursor.fetchall()

        market_data_list = []
        for row in rows:
            # Deserialize klines from JSON
            klines_df = pd.read_json(StringIO(row["klines_data"]), orient="records")

            # Deserialize indicators from JSON
            indicators = json.loads(row["indicators"])

            market_data = MarketData(
                symbol=row["symbol"],
                timeframe=row["timeframe"],
                klines=klines_df,
                indicators=indicators,
                timestamp=row["timestamp"],
            )
            market_data_list.append(market_data)

        return market_data_list

    def close(self) -> None:
        """
        Close database connection
        """
        if self.conn:
            self.conn.close()

    def __enter__(self):
        """Context manager entry"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit"""
        self.close()
