import os
import argparse
import psycopg2
import yaml
from urllib.parse import urlparse

def get_config(config_path):
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def count_events(data):
    return data.count('BEGIN:VEVENT')

def import_events():
    parser = argparse.ArgumentParser(description='Import CalDAV events from files to database.')
    parser.add_argument('--config', default='config.yaml', help='Path to config.yaml')
    parser.add_argument('--input', default='export', help='Input directory or file')
    parser.add_argument('--user', help='Target user (filter or override)')
    parser.add_argument('--calendar', help='Target calendar (filter or override)')
    
    args = parser.parse_args()
    
    cfg = get_config(args.config)
    db_url = cfg.get('database', {}).get('url')
    if not db_url:
        print("Error: Could not find database URL in config.")
        return

    try:
        conn = psycopg2.connect(db_url)
        cur = conn.cursor()
    except Exception as e:
        print(f"Error connecting to database: {e}")
        return

    def upload_file(file_path, user_name, cal_name):
        cur.execute("""
            SELECT c.id, c.path 
            FROM calendars c 
            JOIN users u ON c.owner_id = u.id 
            WHERE u.name = %s AND c.name = %s
        """, (user_name, cal_name))
        res = cur.fetchone()
        if not res:
            print(f"Warning: Calendar '{cal_name}' for user '{user_name}' not found in DB. Skipping {file_path}")
            return False
        
        cal_id, cal_base_path = res
        with open(file_path, 'r') as f:
            data = f.read()
        
        filename = os.path.basename(file_path)
        obj_path = os.path.join(cal_base_path, filename)
        etag = f'"{count_events(data)}"'
        
        cur.execute("""
            INSERT INTO calendar_objects (calendar_id, path, data, etag) 
            VALUES (%s, %s, %s, %s)
            ON CONFLICT (calendar_id, path) DO UPDATE 
            SET data = EXCLUDED.data, etag = EXCLUDED.etag
        """, (cal_id, obj_path, data, etag))
        return True

    # Gather all files
    files_to_process = []
    if os.path.isfile(args.input):
        files_to_process.append(args.input)
    elif os.path.isdir(args.input):
        for root, _, filenames in os.walk(args.input):
            for f in filenames:
                if f.endswith('.ics'):
                    files_to_process.append(os.path.join(root, f))
    else:
        print(f"Error: Input '{args.input}' not found.")
        return

    success_count = 0
    for f_path in files_to_process:
        # Determine source structure
        # If input is a dir, rel_path is relative to it. If file, it's just the filename.
        if os.path.isdir(args.input):
            rel_path = os.path.relpath(f_path, args.input)
        else:
            rel_path = os.path.basename(f_path)
            
        parts = rel_path.split(os.sep)
        
        source_user = None
        source_cal = None
        
        # Structure: user/calendar/file.ics (len 3)
        if len(parts) >= 3:
            source_user = parts[-3]
            source_cal = parts[-2]
        # Structure: calendar/file.ics (len 2)
        elif len(parts) == 2:
            source_cal = parts[-2]
            
        # 1. Apply Filtering: If flag is set and we have a source value, they must match.
        if args.user and source_user and args.user != source_user:
            continue
        if args.calendar and source_cal and args.calendar != source_cal:
            continue
            
        # 2. Determine Target: Flag overrides source.
        target_user = args.user or source_user
        target_cal = args.calendar or source_cal
        
        if not target_user or not target_cal:
            print(f"Skipping {f_path}: Cannot determine user/calendar. Use --user and --calendar flags.")
            continue
            
        if upload_file(f_path, target_user, target_cal):
            print(f"Imported: {target_user}/{target_cal}/{os.path.basename(f_path)}")
            success_count += 1

    conn.commit()
    cur.close()
    conn.close()
    print(f"Done. Successfully imported {success_count} events.")

if __name__ == "__main__":
    import_events()
