import sqlite3
import datetime
from typing import List, Optional, Dict, Any

class Database:
    def __init__(self, db_name="exercise_diary.db"):
        self.db_name = db_name
        self.init_db()

    def get_connection(self):
        return sqlite3.connect(self.db_name)

    def init_db(self):
        """Initialize the database tables."""
        conn = self.get_connection()
        cursor = conn.cursor()

        # Daily Log Table (One entry per day)
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS daily_logs (
                date TEXT PRIMARY KEY,
                weight REAL,
                waist REAL,
                sleep_hours REAL,
                sleep_quality INTEGER,
                mood INTEGER,
                notes TEXT
            )
        ''')

        # Exercises Table (Library of available exercises)
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS exercises (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL UNIQUE,
                description TEXT,
                type TEXT DEFAULT 'reps' -- 'reps' or 'time'
            )
        ''')

        # Pre-populate some basic calisthenics exercises if empty
        cursor.execute("SELECT count(*) FROM exercises")
        if cursor.fetchone()[0] == 0:
            basic_exercises = [
                ("Pushups", "Standard pushups", "reps"),
                ("Squats", "Bodyweight squats", "reps"),
                ("Plank", "Static hold", "time"),
                ("Lunges", "Alternating lunges", "reps"),
                ("Crunches", "Abdominal crunches", "reps"),
                ("Burpees", "Full body cardio/strength", "reps"),
                ("Wall Sit", "Static leg hold", "time")
            ]
            cursor.executemany("INSERT INTO exercises (name, description, type) VALUES (?, ?, ?)", basic_exercises)

        # Workout Sessions Table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS workout_sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                date TEXT NOT NULL,
                notes TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')

        # Workout Sets Table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS workout_sets (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id INTEGER,
                exercise_id INTEGER,
                reps INTEGER,
                duration_seconds INTEGER,
                rpe INTEGER, -- 1-10 intensity
                variation TEXT, -- e.g., "Knee pushups"
                FOREIGN KEY (session_id) REFERENCES workout_sessions (id),
                FOREIGN KEY (exercise_id) REFERENCES exercises (id)
            )
        ''')

        # User Profile Table (Single row for settings)
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS user_profile (
                id INTEGER PRIMARY KEY CHECK (id = 1),
                height_cm REAL,
                goal_weight_kg REAL,
                total_xp INTEGER DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')

        # Unlocked Achievements Table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS unlocked_achievements (
                achievement_id TEXT PRIMARY KEY,
                unlocked_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')

        conn.commit()
        conn.close()

    # --- User Profile Operations ---
    def get_user_profile(self) -> Optional[Dict]:
        conn = self.get_connection()
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM user_profile WHERE id = 1")
        row = cursor.fetchone()
        
        # Ensure total_xp exists (migration for existing db)
        if row and 'total_xp' not in row.keys():
            cursor.execute("ALTER TABLE user_profile ADD COLUMN total_xp INTEGER DEFAULT 0")
            conn.commit()
            cursor.execute("SELECT * FROM user_profile WHERE id = 1")
            row = cursor.fetchone()

        conn.close()
        return dict(row) if row else None

    def save_user_profile(self, height_cm: float, goal_weight_kg: float):
        conn = self.get_connection()
        cursor = conn.cursor()
        # Check if exists to preserve XP
        cursor.execute("SELECT total_xp FROM user_profile WHERE id = 1")
        row = cursor.fetchone()
        current_xp = row[0] if row else 0
        
        cursor.execute("INSERT OR REPLACE INTO user_profile (id, height_cm, goal_weight_kg, total_xp) VALUES (1, ?, ?, ?)", (height_cm, goal_weight_kg, current_xp))
        conn.commit()
        conn.close()

    def add_xp(self, amount: int):
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("UPDATE user_profile SET total_xp = total_xp + ? WHERE id = 1", (amount,))
        conn.commit()
        conn.close()

    def unlock_achievement(self, achievement_id: str):
        conn = self.get_connection()
        cursor = conn.cursor()
        try:
            cursor.execute("INSERT INTO unlocked_achievements (achievement_id) VALUES (?)", (achievement_id,))
            conn.commit()
            return True # Newly unlocked
        except sqlite3.IntegrityError:
            return False # Already unlocked
        finally:
            conn.close()

    def get_achievements(self) -> List[str]:
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT achievement_id FROM unlocked_achievements")
        rows = cursor.fetchall()
        conn.close()
        return [r[0] for r in rows]

    # --- Daily Log Operations ---
    def save_daily_log(self, date: str, data: Dict[str, Any]):
        """Insert or update a daily log."""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        # Check if exists
        cursor.execute("SELECT date FROM daily_logs WHERE date = ?", (date,))
        exists = cursor.fetchone()

        if exists:
            # Dynamic update based on keys provided
            set_clause = ", ".join([f"{key} = ?" for key in data.keys()])
            values = list(data.values()) + [date]
            cursor.execute(f"UPDATE daily_logs SET {set_clause} WHERE date = ?", values)
        else:
            columns = ", ".join(["date"] + list(data.keys()))
            placeholders = ", ".join(["?"] * (len(data) + 1))
            values = [date] + list(data.values())
            cursor.execute(f"INSERT INTO daily_logs ({columns}) VALUES ({placeholders})", values)
        
        conn.commit()
        conn.close()

    def get_daily_log(self, date: str) -> Optional[Dict]:
        conn = self.get_connection()
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM daily_logs WHERE date = ?", (date,))
        row = cursor.fetchone()
        conn.close()
        if row:
            return dict(row)
        return None

    def get_log_history(self, limit=30) -> List[Dict]:
        """Get the last N days of logs."""
        conn = self.get_connection()
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM daily_logs ORDER BY date DESC LIMIT ?", (limit,))
        rows = cursor.fetchall()
        conn.close()
        return [dict(row) for row in rows]

    # --- Exercise Operations ---
    def get_all_exercises(self) -> List[Dict]:
        conn = self.get_connection()
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM exercises ORDER BY name")
        rows = cursor.fetchall()
        conn.close()
        return [dict(row) for row in rows]

    def add_exercise(self, name, description="", type="reps"):
        conn = self.get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("INSERT INTO exercises (name, description, type) VALUES (?, ?, ?)", (name, description, type))
            conn.commit()
            return True
        except sqlite3.IntegrityError:
            return False # Name already exists
        finally:
            conn.close()

    # --- Workout Operations ---
    def create_workout_session(self, date: str, notes: str = "") -> int:
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("INSERT INTO workout_sessions (date, notes) VALUES (?, ?)", (date, notes))
        session_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return session_id

    def add_set(self, session_id: int, exercise_id: int, reps: int = 0, duration: int = 0, rpe: int = 0, variation: str = ""):
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute('''
            INSERT INTO workout_sets (session_id, exercise_id, reps, duration_seconds, rpe, variation)
            VALUES (?, ?, ?, ?, ?, ?)
        ''', (session_id, exercise_id, reps, duration, rpe, variation))
        conn.commit()
        conn.close()

    def delete_set(self, set_id: int):
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("DELETE FROM workout_sets WHERE id = ?", (set_id,))
        conn.commit()
        conn.close()

    def get_workouts_by_date(self, date: str) -> List[Dict]:
        conn = self.get_connection()
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        # Get sessions
        cursor.execute("SELECT * FROM workout_sessions WHERE date = ?", (date,))
        sessions = [dict(row) for row in cursor.fetchall()]
        
        # Get sets for each session
        for session in sessions:
            cursor.execute('''
                SELECT s.*, e.name as exercise_name 
                FROM workout_sets s 
                JOIN exercises e ON s.exercise_id = e.id 
                WHERE s.session_id = ?
            ''', (session['id'],))
            session['sets'] = [dict(row) for row in cursor.fetchall()]
            
        conn.close()
        return sessions