db.py

import os
import psycopg2
import psycopg2.extras
from typing import List, Optional, Dict, Any
from datetime import date
from models import Application
from dotenv import load_dotenv
load_dotenv()

DATABASE_URL = os.environ.get("DATABASE_URL")

def get_connection():
    """Returns a connection to the Postgres database."""
    conn = psycopg2.connect(DATABASE_URL)
    return conn

def init_db():
    """Initializes the database and creates the applications table if it doesn't exist."""
    if not DATABASE_URL:
        # Don't fail completely if no DB URL is set yet, but just skip init
        return
        
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS applications (
                    id SERIAL PRIMARY KEY,
                    title TEXT NOT NULL,
                    org TEXT NOT NULL,
                    type TEXT NOT NULL,
                    platform TEXT NOT NULL,
                    url TEXT NOT NULL,
                    deadline DATE,
                    status TEXT NOT NULL,
                    date_applied DATE,
                    notes TEXT NOT NULL
                )
            ''')
            cursor.execute('''
                ALTER TABLE applications ADD COLUMN IF NOT EXISTS contact_used TEXT NOT NULL DEFAULT '';
            ''')
        conn.commit()
    finally:
        conn.close()

def _row_to_app(row) -> Application:
    """Helper to convert psycopg2 dict row to Application dataclass"""
    return Application(
        id=row['id'],
        title=row['title'],
        org=row['org'],
        type=row['type'],
        platform=row['platform'],
        url=row['url'],
        deadline=row['deadline'],
        status=row['status'],
        date_applied=row['date_applied'],
        notes=row['notes'],
        contact_used=row['contact_used']
    )

def add_application(app: Application) -> int:
    try:
        conn = get_connection()
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
            cursor.execute('''
                INSERT INTO applications (title, org, type, platform, url, deadline, status, date_applied, notes, contact_used)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                RETURNING id
            ''', (app.title, app.org, app.type, app.platform, app.url, app.deadline, app.status, app.date_applied, app.notes, app.contact_used))
            row = cursor.fetchone()
        conn.commit()
        return row['id']
    finally:
        conn.close()

def get_all_applications(sort_by_deadline: bool = True) -> List[Application]:
    try:
        conn = get_connection()
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
            query = 'SELECT * FROM applications'
            if sort_by_deadline:
                # Sort by deadline ASC, putting NULLs at the end
                query += ' ORDER BY CASE WHEN deadline IS NULL THEN 1 ELSE 0 END, deadline ASC'
                
            cursor.execute(query)
            return [_row_to_app(row) for row in cursor.fetchall()]
    finally:
        conn.close()

def get_upcoming_deadlines() -> List[Application]:
    """Returns applications with deadlines today or in the future, sorted ascending."""
    try:
        conn = get_connection()
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
            cursor.execute('''
                SELECT * FROM applications 
                WHERE deadline IS NOT NULL AND deadline >= %s
                ORDER BY deadline ASC
            ''', (date.today(),))
            return [_row_to_app(row) for row in cursor.fetchall()]
    finally:
        conn.close()

def get_application(app_id: int) -> Optional[Application]:
    try:
        conn = get_connection()
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
            cursor.execute('SELECT * FROM applications WHERE id = %s', (app_id,))
            row = cursor.fetchone()
            if row:
                return _row_to_app(row)
            return None
    finally:
        conn.close()

def update_application(app_id: int, app: Application):
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            cursor.execute('''
                UPDATE applications 
                SET title = %s, org = %s, type = %s, platform = %s, url = %s, deadline = %s, status = %s, date_applied = %s, notes = %s, contact_used = %s
                WHERE id = %s
            ''', (app.title, app.org, app.type, app.platform, app.url, app.deadline, app.status, app.date_applied, app.notes, app.contact_used, app_id))
        conn.commit()
    finally:
        conn.close()

def update_status(app_id: int, status: str, date_applied: Optional[date] = None):
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            if date_applied:
                cursor.execute('UPDATE applications SET status = %s, date_applied = %s WHERE id = %s', (status, date_applied, app_id))
            else:
                cursor.execute('UPDATE applications SET status = %s WHERE id = %s', (status, app_id))
        conn.commit()
    finally:
        conn.close()

def delete_application(app_id: int):
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            cursor.execute('DELETE FROM applications WHERE id = %s', (app_id,))
        conn.commit()
    finally:
        conn.close()

def search_applications(query: str) -> List[Application]:
    search_term = f'%{query}%'
    try:
        conn = get_connection()
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
            # Using ILIKE for case-insensitive search in Postgres
            cursor.execute('''
                SELECT * FROM applications 
                WHERE title ILIKE %s OR org ILIKE %s OR platform ILIKE %s
            ''', (search_term, search_term, search_term))
            return [_row_to_app(row) for row in cursor.fetchall()]
    finally:
        conn.close()

def get_stats() -> Dict[str, Any]:
    stats = {
        'total': 0,
        'by_type': {},
        'by_status': {},
        'active_count': 0
    }
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            # Total
            cursor.execute('SELECT COUNT(*) FROM applications')
            stats['total'] = cursor.fetchone()[0]
            
            # By type
            cursor.execute('SELECT type, COUNT(*) FROM applications GROUP BY type')
            for row in cursor.fetchall():
                stats['by_type'][row[0]] = row[1]
                
            # By status
            cursor.execute('SELECT status, COUNT(*) FROM applications GROUP BY status')
            for row in cursor.fetchall():
                stats['by_status'][row[0]] = row[1]
                
            # Active count (pending, interview)
            cursor.execute("SELECT COUNT(*) FROM applications WHERE status IN ('pending', 'interview')")
            stats['active_count'] = cursor.fetchone()[0]
            
    finally:
        conn.close()
        
    return stats