import sqlite3
import pandas as pd
from pathlib import Path
from typing import Optional, List, Tuple
from contextlib import contextmanager
import os

class DatabaseManager:
    """Manages SQLite database for storing OHLCV data."""

    def __init__(self, db_path: str = "./data/cache/market_data.db"):
        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.initialize_db()

    @contextmanager
    def get_connection(self):
        """Context manager for database connection."""
        conn = sqlite3.connect(self.db_path)
        try:
            yield conn
        finally:
            conn.close()

    def initialize_db(self):
        """Create tables if they don't exist."""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # Create OHLCV table
            # Composite primary key ensures uniqueness for a specific candle
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS ohlcv (
                    timestamp DATETIME,
                    symbol TEXT,
                    timeframe TEXT,
                    open REAL,
                    high REAL,
                    low REAL,
                    close REAL,
                    volume REAL,
                    PRIMARY KEY (symbol, timeframe, timestamp)
                )
            """)
            
            # Create index for faster querying by time range
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_ohlcv_time 
                ON ohlcv (symbol, timeframe, timestamp)
            """)
            
            conn.commit()

    def save_ohlcv(self, df: pd.DataFrame, symbol: str, timeframe: str):
        """Save OHLCV data to database (Upsert)."""
        if df.empty:
            return

        # Ensure timestamp is datetime
        df = df.copy()
        if not pd.api.types.is_datetime64_any_dtype(df['timestamp']):
            df['timestamp'] = pd.to_datetime(df['timestamp'])
            
        # Convert to string for SQLite compatibility (ISO 8601)
        df['timestamp'] = df['timestamp'].dt.strftime('%Y-%m-%d %H:%M:%S')
        
        # Add metadata columns if they don't exist
        if 'symbol' not in df.columns:
            df['symbol'] = symbol
        if 'timeframe' not in df.columns:
            df['timeframe'] = timeframe
        
        # Prepare data for insertion (ensure order)
        records = df[['timestamp', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume']].to_records(index=False)
        
        with self.get_connection() as conn:
            cursor = conn.cursor()
            # efficiently insert or replace
            # sqlite3 adapter handles datetime objects correctly
            cursor.executemany("""
                INSERT OR REPLACE INTO ohlcv 
                (timestamp, symbol, timeframe, open, high, low, close, volume)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, records)
            conn.commit()

    def get_ohlcv(self, symbol: str, timeframe: str, start_date: Optional[str] = None, end_date: Optional[str] = None) -> pd.DataFrame:
        """Retrieve OHLCV data from database."""
        query = """
            SELECT timestamp, open, high, low, close, volume 
            FROM ohlcv 
            WHERE symbol = ? AND timeframe = ?
        """
        params = [symbol, timeframe]

        if start_date:
            query += " AND timestamp >= ?"
            params.append(start_date)
        
        if end_date:
            # Handle 'YYYY-MM-DD' vs full timestamp
            if len(str(end_date)) == 10:  # YYYY-MM-DD
                 query += " AND timestamp < date(?, '+1 day')"
            else:
                 query += " AND timestamp <= ?"
            params.append(end_date)
            
        query += " ORDER BY timestamp ASC"

        with self.get_connection() as conn:
            df = pd.read_sql_query(query, conn, params=params, parse_dates=['timestamp'])
            
        return df

    def get_latest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]:
        """Get the latest timestamp for a symbol/timeframe."""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT MAX(timestamp) 
                FROM ohlcv 
                WHERE symbol = ? AND timeframe = ?
            """, (symbol, timeframe))
            result = cursor.fetchone()
            
        return pd.to_datetime(result[0]) if result[0] else None

    def get_earliest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]:
        """Get the earliest timestamp for a symbol/timeframe."""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT MIN(timestamp) 
                FROM ohlcv 
                WHERE symbol = ? AND timeframe = ?
            """, (symbol, timeframe))
            result = cursor.fetchone()
            
        return pd.to_datetime(result[0]) if result[0] else None