import sqlite3
import os
from contextlib import contextmanager
from typing import Optional, List, Dict


DB_PATH = os.path.join(os.path.dirname(__file__), 'publishers.db')


class DatabaseManager:
    def __init__(self, db_path: str = DB_PATH):
        self.db_path = db_path
        self._init_db()
    
    def _init_db(self):
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS publishers (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT UNIQUE NOT NULL,
                    port INTEGER NOT NULL,
                    host TEXT NOT NULL DEFAULT 'localhost',
                    description TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_port_host 
                ON publishers(port, host)
            ''')
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_name 
                ON publishers(name)
            ''')
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS subscriptions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    subscriber_name TEXT NOT NULL,
                    publisher_name TEXT NOT NULL,
                    subscriber_host TEXT NOT NULL DEFAULT 'localhost',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_active TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(subscriber_name, publisher_name)
                )
            ''')
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_subscriber
                ON subscriptions(subscriber_name)
            ''')
            cursor.execute('''
                CREATE INDEX IF NOT EXISTS idx_publisher_sub
                ON subscriptions(publisher_name)
            ''')
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS config (
                    key TEXT PRIMARY KEY,
                    value TEXT NOT NULL,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            cursor.execute('''
                INSERT OR IGNORE INTO config (key, value) VALUES ('port_range_start', '5555')
            ''')
            cursor.execute('''
                INSERT OR IGNORE INTO config (key, value) VALUES ('port_range_count', '100')
            ''')
            conn.commit()
    
    @contextmanager
    def _get_connection(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
        finally:
            conn.close()
    
    def add_publisher(self, name: str, port: int, host: str = 'localhost', description: str = '') -> bool:
        try:
            with self._get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    INSERT INTO publishers (name, port, host, description)
                    VALUES (?, ?, ?, ?)
                ''', (name, port, host, description))
                conn.commit()
                return True
        except sqlite3.IntegrityError:
            return False
    
    def update_publisher(self, name: str, port: Optional[int] = None, 
                        host: Optional[str] = None, description: Optional[str] = None) -> bool:
        updates = []
        params = []
        
        if port is not None:
            updates.append('port = ?')
            params.append(port)
        if host is not None:
            updates.append('host = ?')
            params.append(host)
        if description is not None:
            updates.append('description = ?')
            params.append(description)
        
        if not updates:
            return False
        
        updates.append('updated_at = CURRENT_TIMESTAMP')
        params.append(name)
        
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(f'''
                UPDATE publishers 
                SET {', '.join(updates)}
                WHERE name = ?
            ''', params)
            conn.commit()
            return cursor.rowcount > 0
    
    def remove_publisher(self, name: str) -> bool:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM publishers WHERE name = ?', (name,))
            conn.commit()
            return cursor.rowcount > 0
    
    def get_publisher(self, name: str) -> Optional[Dict]:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT name, port, host, description, created_at, updated_at
                FROM publishers WHERE name = ?
            ''', (name,))
            row = cursor.fetchone()
            if row:
                return dict(row)
            return None
    
    def get_publisher_by_port(self, port: int, host: str = 'localhost') -> Optional[Dict]:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT name, port, host, description, created_at, updated_at
                FROM publishers WHERE port = ? AND host = ?
            ''', (port, host))
            row = cursor.fetchone()
            if row:
                return dict(row)
            return None
    
    def list_all_publishers(self) -> List[Dict]:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT name, port, host, description, created_at, updated_at
                FROM publishers
                ORDER BY name
            ''')
            return [dict(row) for row in cursor.fetchall()]
    
    def publisher_exists(self, name: str) -> bool:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT 1 FROM publishers WHERE name = ?', (name,))
            return cursor.fetchone() is not None
    
    def get_config(self, key: str) -> Optional[str]:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT value FROM config WHERE key = ?', (key,))
            row = cursor.fetchone()
            if row:
                return row['value']
            return None
    
    def set_config(self, key: str, value: str) -> bool:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT OR REPLACE INTO config (key, value, updated_at) 
                VALUES (?, ?, CURRENT_TIMESTAMP)
            ''', (key, value))
            conn.commit()
            return True
    
    def get_port_range(self) -> tuple[int, int]:
        start = int(self.get_config('port_range_start') or '5555')
        count = int(self.get_config('port_range_count') or '100')
        return start, start + count
    
    def set_port_range(self, start: int, count: int) -> bool:
        self.set_config('port_range_start', str(start))
        self.set_config('port_range_count', str(count))
        return True
    
    def get_used_ports(self, host: str = 'localhost') -> List[int]:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT port FROM publishers WHERE host = ? ORDER BY port', (host,))
            return [row['port'] for row in cursor.fetchall()]
    
    def find_available_port(self, host: str = 'localhost') -> Optional[int]:
        start, end = self.get_port_range()
        used_ports = set(self.get_used_ports(host))
        
        for port in range(start, end):
            if port not in used_ports:
                return port
        
        return None
    
    def add_subscription(self, subscriber_name: str, publisher_name: str, subscriber_host: str = 'localhost') -> bool:
        try:
            with self._get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    INSERT INTO subscriptions (subscriber_name, publisher_name, subscriber_host)
                    VALUES (?, ?, ?)
                    ON CONFLICT(subscriber_name, publisher_name) 
                    DO UPDATE SET last_active = CURRENT_TIMESTAMP
                ''', (subscriber_name, publisher_name, subscriber_host))
                conn.commit()
                return True
        except Exception:
            return False
    
    def remove_subscription(self, subscriber_name: str, publisher_name: str) -> bool:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                DELETE FROM subscriptions 
                WHERE subscriber_name = ? AND publisher_name = ?
            ''', (subscriber_name, publisher_name))
            conn.commit()
            return cursor.rowcount > 0
    
    def get_subscribers(self, publisher_name: str) -> List[Dict]:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT subscriber_name, subscriber_host, created_at, last_active
                FROM subscriptions
                WHERE publisher_name = ?
                ORDER BY last_active DESC
            ''', (publisher_name,))
            return [dict(row) for row in cursor.fetchall()]
    
    def get_subscriptions(self, subscriber_name: str) -> List[Dict]:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT publisher_name, created_at, last_active
                FROM subscriptions
                WHERE subscriber_name = ?
                ORDER BY last_active DESC
            ''', (subscriber_name,))
            return [dict(row) for row in cursor.fetchall()]
    
    def get_all_subscriptions(self) -> List[Dict]:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT subscriber_name, publisher_name, subscriber_host, created_at, last_active
                FROM subscriptions
                ORDER BY publisher_name, subscriber_name
            ''')
            return [dict(row) for row in cursor.fetchall()]
    
    def get_graph(self) -> Dict:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            
            publishers = {}
            cursor.execute('SELECT name, port, host FROM publishers ORDER BY name')
            for row in cursor.fetchall():
                publishers[row['name']] = {
                    'port': row['port'],
                    'host': row['host'],
                    'subscribers': []
                }
            
            cursor.execute('''
                SELECT publisher_name, subscriber_name, subscriber_host
                FROM subscriptions
                ORDER BY publisher_name, subscriber_name
            ''')
            for row in cursor.fetchall():
                pub_name = row['publisher_name']
                if pub_name in publishers:
                    publishers[pub_name]['subscribers'].append({
                        'name': row['subscriber_name'],
                        'host': row['subscriber_host']
                    })
            
            return publishers