import sqlite3
import datetime
import pandas as pd
import os

DB_FILE = os.environ.get("DB_FILE", "work_log.db")

# Ensure the directory for the database file exists
db_dir = os.path.dirname(DB_FILE)
if db_dir:
    os.makedirs(db_dir, exist_ok=True)

def get_connection():
    conn = sqlite3.connect(DB_FILE)
    return conn

def init_db():
    conn = get_connection()
    c = conn.cursor()
    
    # Create table with new schema
    c.execute('''
        CREATE TABLE IF NOT EXISTS work_days (
            date TEXT PRIMARY KEY,
            hours REAL DEFAULT 8.0,
            note TEXT,
            type TEXT DEFAULT 'work',
            course_name TEXT,
            travel_costs INTEGER DEFAULT 0
        )
    ''')
    
    # Migration: Check if columns exist, if not add them
    c.execute("PRAGMA table_info(work_days)")
    columns = [info[1] for info in c.fetchall()]
    
    if 'type' not in columns:
        c.execute("ALTER TABLE work_days ADD COLUMN type TEXT DEFAULT 'work'")
    if 'course_name' not in columns:
        c.execute("ALTER TABLE work_days ADD COLUMN course_name TEXT")
    if 'travel_costs' not in columns:
        c.execute("ALTER TABLE work_days ADD COLUMN travel_costs INTEGER DEFAULT 0")
        
    conn.commit()
    conn.close()

def update_day(date_str, day_type="work", course_name=None, hours=8.0, travel_costs=False):
    """Updates or inserts a day. If course_name is None, it won't overwrite existing unless it's a new entry."""
    conn = get_connection()
    c = conn.cursor()
    
    travel_val = 1 if travel_costs else 0
    
    # Check if exists
    c.execute("SELECT date FROM work_days WHERE date = ?", (date_str,))
    data = c.fetchone()
    
    if data:
        # Update existing
        c.execute("""
            UPDATE work_days 
            SET type = ?, course_name = ?, hours = ?, travel_costs = ?
            WHERE date = ?
        """, (day_type, course_name, hours, travel_val, date_str))
    else:
        # Insert new
        c.execute("""
            INSERT INTO work_days (date, hours, type, course_name, travel_costs) 
            VALUES (?, ?, ?, ?, ?)
        """, (date_str, hours, day_type, course_name, travel_val))
        
    conn.commit()
    conn.close()

def remove_day(date_str):
    """Removes a day from the database."""
    conn = get_connection()
    c = conn.cursor()
    c.execute("DELETE FROM work_days WHERE date = ?", (date_str,))
    conn.commit()
    conn.close()

def get_all_days():
    """Returns a list of all work days."""
    conn = get_connection()
    df = pd.read_sql_query("SELECT * FROM work_days", conn)
    conn.close()
    return df