import asyncio
import aiosqlite
import ast
import json
import re
import threading
import dataclasses
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, TypeVar, Generic, Type, get_type_hints, Union, Tuple
from dataclasses import dataclass, asdict, fields, is_dataclass, field
from functools import lru_cache
import uuid
import contextlib
import enum
from loguru import logger

_VALID_IDENTIFIER = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$')

def _validate_identifier(name: str) -> str:
    if not _VALID_IDENTIFIER.match(name):
        raise ValueError(f"Invalid SQL identifier: {name!r}")
    return name

T = TypeVar('T')
SchemaType = TypeVar('SchemaType', bound='BaseModel')

@dataclass
class BaseModel:
    id: str = field(default_factory=lambda: str(uuid.uuid4()), metadata={"primary_key": True})
    created_at: datetime = field(default_factory=datetime.now)
    updated_at: datetime = field(default_factory=datetime.now)


class AsyncDB(Generic[SchemaType]):
    """High-performance async SQLite with dataclass schema and reliable connection handling."""

    OPERATOR_MAP = {
        'gt': '>', 'lt': '<', 'gte': '>=', 'lte': '<=',
        'neq': '!=', 'like': 'LIKE', 'in': 'IN', 'eq': '='
    }

    # Shared write locks per database file (class-level)
    _db_locks: dict[str, asyncio.Lock] = {}
    # Lock for schema initialization (class-level)
    _schema_init_lock: asyncio.Lock = None
    # Threading lock to guard class-level dict mutations
    _db_locks_guard = threading.Lock()

    def __init__(self, db_path: Union[str, Path], table_name: str, schema_class: Type[SchemaType]):
        """Initialize AsyncDB with a path and schema dataclass."""
        if not is_dataclass(schema_class):
            raise TypeError(f"Schema must be a dataclass, got {schema_class}")

        self.db_path = Path(db_path).resolve()
        self.schema_class = schema_class
        self.table_name = _validate_identifier(table_name)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        # Validate all field names upfront
        for f in fields(schema_class):
            _validate_identifier(f.name)

        # Make schema initialization unique per instance
        self._db_key = f"{str(self.db_path)}:{self.table_name}:{self.schema_class.__name__}"

        # Use shared lock per database file (not per instance)
        db_path_str = str(self.db_path)
        with AsyncDB._db_locks_guard:
            if db_path_str not in AsyncDB._db_locks:
                AsyncDB._db_locks[db_path_str] = asyncio.Lock()
            self._write_lock = AsyncDB._db_locks[db_path_str]

        self._type_hints = get_type_hints(schema_class)

        # Persistent connection (lazy init)
        self._connection: Optional[aiosqlite.Connection] = None

        # Use a class-level set to track initialized schemas
        if not hasattr(AsyncDB, '_initialized_schemas'):
            AsyncDB._initialized_schemas = set()
    
    async def _ensure_connection(self, max_retries: int = 5) -> aiosqlite.Connection:
        """Return the persistent connection, creating it on first call with retry logic."""
        if self._connection is not None:
            return self._connection

        # Ensure schema init lock exists (lazy init for asyncio compatibility)
        with AsyncDB._db_locks_guard:
            if AsyncDB._schema_init_lock is None:
                AsyncDB._schema_init_lock = asyncio.Lock()

        last_error = None
        for attempt in range(max_retries):
            try:
                db = aiosqlite.connect(self.db_path, timeout=30.0)
                # Mark aiosqlite's thread as daemon so it won't block process exit
                db.daemon = True
                db = await db
                # Fast WAL mode with minimal sync
                await db.execute("PRAGMA journal_mode=WAL")
                await db.execute("PRAGMA synchronous=NORMAL")
                await db.execute("PRAGMA cache_size=10000")
                await db.execute("PRAGMA busy_timeout=30000")  # 30s busy timeout

                # Initialize schema if needed (with lock to prevent race condition)
                if self._db_key not in AsyncDB._initialized_schemas:
                    async with AsyncDB._schema_init_lock:
                        # Double-check after acquiring lock
                        if self._db_key not in AsyncDB._initialized_schemas:
                            await self._init_schema(db)
                            AsyncDB._initialized_schemas.add(self._db_key)

                self._connection = db
                return db
            except Exception as e:
                last_error = e
                if attempt < max_retries - 1:
                    # Exponential backoff: 0.1s, 0.2s, 0.4s, 0.8s, 1.6s
                    wait_time = 0.1 * (2 ** attempt)
                    await asyncio.sleep(wait_time)
                    continue
                raise
        raise last_error

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    async def close(self) -> None:
        """Explicitly close the persistent connection."""
        if self._connection is not None:
            try:
                await self._connection.close()
            except Exception:
                pass
            self._connection = None

    async def _init_schema(self, db: aiosqlite.Connection) -> None:
        """Generate schema from dataclass structure with support for field additions."""
        logger.debug(f"Initializing schema for {self.schema_class.__name__} in table {self.table_name}")
        
        field_defs = []
        indexes = []
        
        # First check if table exists
        cursor = await db.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
            (self.table_name,)
        )
        table_exists = await cursor.fetchone() is not None
        
        existing_columns = set()
        if table_exists:
            # Get existing columns if table exists
            cursor = await db.execute(f"PRAGMA table_info({self.table_name})")
            columns = await cursor.fetchall()
            existing_columns = {col[1] for col in columns}  # col[1] is the column name
        
        # Process all fields in the dataclass - ONLY THIS SCHEMA CLASS
        schema_fields = fields(self.schema_class)
        logger.debug(f"Processing {len(schema_fields)} fields for {self.schema_class.__name__}")
        
        for f in schema_fields:
            field_name = f.name
            field_type = self._type_hints.get(field_name)
            logger.debug(f"  Field: {field_name} -> {field_type}")
            
            # Map Python types to SQLite types
            if field_type in (int, bool):
                sql_type = "INTEGER"
            elif field_type in (float,):
                sql_type = "REAL"
            elif field_type == bytes:
                sql_type = "BLOB"
            elif field_type in (str, enum.EnumType):
                sql_type = "TEXT"
            elif field_type in (datetime,):
                sql_type = "TIMESTAMP"
            elif field_type == List[str]:
                sql_type = "TEXT"  # Stored as JSON
            else:
                sql_type = "TEXT"  # Default to TEXT/JSON for complex types
                
            # Handle special field metadata
            constraints = []
            if f.metadata.get('primary_key'):
                constraints.append("PRIMARY KEY")
            if f.metadata.get('unique'):
                constraints.append("UNIQUE")
            if f.default is dataclasses.MISSING and f.default_factory is dataclasses.MISSING and f.metadata.get('required', True):
                constraints.append("NOT NULL")
                
            field_def = f"{field_name} {sql_type} {' '.join(constraints)}"
            
            if not table_exists:
                # Add field definition for new table creation
                field_defs.append(field_def)
            elif field_name not in existing_columns:
                # Alter table to add the new column without NOT NULL constraint
                alter_sql = f"ALTER TABLE {self.table_name} ADD COLUMN {field_name} {sql_type}"
                logger.debug(f"  Adding new column: {alter_sql}")
                await db.execute(alter_sql)
                await db.commit()
                
            # Handle indexes
            if f.metadata.get('index'):
                index_name = f"idx_{self.table_name}_{field_name}"
                index_sql = f"CREATE INDEX IF NOT EXISTS {index_name} ON {self.table_name}({field_name})"
                indexes.append(index_sql)
        
        # Create table if it doesn't exist
        if not table_exists:
            # Check for table constraints
            table_constraints = getattr(self.schema_class, '__table_constraints__', [])

            constraints_sql = ""
            if table_constraints:
                constraints_sql = ", " + ", ".join(table_constraints)

            create_sql = f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    {', '.join(field_defs)}{constraints_sql}
                )
            """
            logger.debug(f"Creating table: {create_sql}")
            await db.execute(create_sql)
        
        # Create indexes
        for idx_stmt in indexes:
            await db.execute(idx_stmt)
            
        await db.commit()
        logger.debug(f"Schema initialization complete for {self.schema_class.__name__}")
    
    @contextlib.asynccontextmanager
    async def transaction(self):
        """Run operations in a transaction with reliable cleanup and auto-reconnect."""
        db = await self._ensure_connection()
        try:
            yield db
            await db.commit()
        except Exception as e:
            try:
                await db.rollback()
            except Exception:
                pass
            if "closed" in str(e).lower() or "no active connection" in str(e).lower():
                self._connection = None
                db = await self._ensure_connection()
                try:
                    yield db
                    await db.commit()
                except Exception:
                    await db.rollback()
                    raise
            else:
                raise
    
    # @lru_cache(maxsize=128)
    def _serialize_value(self, value: Any) -> Any:
        """Fast value serialization with type-based optimization."""
        if value is None or isinstance(value, (int, float, bool, str)):
            return value
        if isinstance(value, bytes):
            return value 
        if isinstance(value, datetime):
            return value.isoformat()
        if isinstance(value, enum.Enum):
            return value.value
        if isinstance(value, (list, dict, tuple)):
            return json.dumps(value)
        return str(value)
    
    def _deserialize_value(self, field_name: str, value: Any) -> Any:
        """Deserialize values based on field type."""
        if value is None:
            return value

        field_type = self._type_hints.get(field_name)

        # Handle bytes fields - keep as bytes
        if field_type == bytes:
            if isinstance(value, bytes):
                return value
            # If somehow stored as string, convert back
            if isinstance(value, str):
                try:
                    return ast.literal_eval(value)
                except (ValueError, SyntaxError):
                    return value.encode('utf-8')

        # Handle bool fields - SQLite stores as INTEGER, need to convert back
        if field_type is bool:
            return bool(value)

        # Handle string fields - ensure phone numbers are strings
        if field_type is str or (hasattr(field_type, '__origin__') and field_type.__origin__ is Union and str in getattr(field_type, '__args__', ())):
            return str(value)

        if field_type is datetime and isinstance(value, str):
            return datetime.fromisoformat(value)

        # Handle enum types
        if hasattr(field_type, '__origin__') and field_type.__origin__ is Union:
            # Handle Optional[EnumType] case
            args = getattr(field_type, '__args__', ())
            for arg in args:
                if arg is not type(None) and isinstance(arg, type) and issubclass(arg, enum.Enum):
                    try:
                        return arg(value)
                    except (ValueError, TypeError):
                        pass
        elif isinstance(field_type, type) and issubclass(field_type, enum.Enum):
            # Handle direct enum types
            try:
                return field_type(value)
            except (ValueError, TypeError):
                pass

        if isinstance(value, str):
            try:
                return json.loads(value)
            except (json.JSONDecodeError, TypeError):
                pass

        return value
    
    @lru_cache(maxsize=64)
    def _generate_save_sql(self, field_names: Tuple[str, ...]) -> str:
        """Generate efficient SQL for upsert with proper conflict handling."""
        columns = ','.join(field_names)
        placeholders = ','.join('?' for _ in field_names)

        set_clause = ','.join(f'{col}=excluded.{col}' for col in field_names if col != 'created_at')
        return f"""
            INSERT INTO {self.table_name} ({columns},id)
            VALUES ({placeholders},?)
            ON CONFLICT(id) DO UPDATE SET {set_clause}
        """
    
    def _prepare_item(self, item: SchemaType) -> Tuple[str, List[Any]]:
        """Prepare an item for saving. Returns (sql, values)."""
        data = asdict(item)
        item_id = data.pop('id', None) or str(uuid.uuid4())
        now = datetime.now()
        if not data.get('created_at'):
            data['created_at'] = now
        data['updated_at'] = now
        field_names = tuple(sorted(data.keys()))
        sql = self._generate_save_sql(field_names)
        values = [self._serialize_value(data[name]) for name in field_names]
        values.append(item_id)
        return sql, values

    async def save_batch(self, items: List[SchemaType], skip_errors: bool = True) -> int:
        """Save multiple items in a single transaction for better performance.

        Args:
            items: List of schema objects to save
            skip_errors: If True, skip items that cause errors

        Returns:
            Number of items successfully saved
        """
        if not items:
            return 0

        saved_count = 0

        max_retries = 3
        for attempt in range(max_retries):
            try:
                saved_count = 0
                async with self._write_lock:
                    async with self.transaction() as db:
                        for item in items:
                            try:
                                if not isinstance(item, self.schema_class):
                                    if not skip_errors:
                                        raise TypeError(f"Expected {self.schema_class.__name__}, got {type(item).__name__}")
                                    continue

                                sql, values = self._prepare_item(item)
                                await db.execute(sql, values)
                                saved_count += 1

                            except Exception as e:
                                if skip_errors:
                                    logger.warning(f"Save error (skipped): {e}")
                                    continue
                                raise
                break
            except Exception as e:
                if "database is locked" in str(e) and attempt < max_retries - 1:
                    wait_time = 0.2 * (2 ** attempt)
                    logger.debug(f"DB locked, retry {attempt + 1}/{max_retries} in {wait_time}s")
                    await asyncio.sleep(wait_time)
                    continue
                raise

        return saved_count

    async def save(self, item: SchemaType, skip_errors: bool = True) -> bool:
        """Store a schema object with upsert functionality and error handling.

        Args:
            item: The schema object to save
            skip_errors: If True, silently skip errors and return False. If False, raise errors.

        Returns:
            True if save was successful, False if error occurred and skip_errors=True
        """
        try:
            if not isinstance(item, self.schema_class):
                if skip_errors:
                    return False
                raise TypeError(f"Expected {self.schema_class.__name__}, got {type(item).__name__}")

            sql, values = self._prepare_item(item)

            # Perform save with reliable transaction (retry on "database is locked")
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    async with self._write_lock:
                        async with self.transaction() as db:
                            await db.execute(sql, values)
                    break
                except Exception as e:
                    if "database is locked" in str(e) and attempt < max_retries - 1:
                        wait_time = 0.2 * (2 ** attempt)
                        logger.debug(f"DB locked, retry {attempt + 1}/{max_retries} in {wait_time}s")
                        await asyncio.sleep(wait_time)
                        continue
                    raise

            return True

        except Exception as e:
            if skip_errors:
                logger.warning(f"Save error (skipped): {e}")
                return False
            raise
    
    async def get_by_id(self, id: str) -> Optional[SchemaType]:
        """Fetch an item by ID with reliable connection handling."""
        async with self.transaction() as db:
            cursor = await db.execute(f"SELECT * FROM {self.table_name} WHERE id = ?", (id,))
            row = await cursor.fetchone()
            
            if not row:
                return None
                
            # Get column names and build data dictionary
            columns = [desc[0] for desc in cursor.description]
            return self.schema_class(**{
                col: self._deserialize_value(col, row[i]) 
                for i, col in enumerate(columns)
            })
    
    def _build_where_clause(self, filters: Dict[str, Any]) -> Tuple[str, List[Any]]:
        """Build optimized WHERE clause for queries."""
        if not filters:
            return "", []
            
        conditions = []
        values = []
        
        for key, value in filters.items():
            # Parse field and operator
            parts = key.split('__', 1)
            field = parts[0]
            
            if len(parts) > 1 and parts[1] in self.OPERATOR_MAP:
                op_str = self.OPERATOR_MAP[parts[1]]
                
                # Handle IN operator specially
                if op_str == 'IN' and isinstance(value, (list, tuple)):
                    placeholders = ','.join(['?'] * len(value))
                    conditions.append(f"{field} IN ({placeholders})")
                    values.extend(value)
                else:
                    conditions.append(f"{field} {op_str} ?")
                    values.append(value)
            else:
                # Default to equality
                conditions.append(f"{field} = ?")
                values.append(value)
                
        return f"WHERE {' AND '.join(conditions)}", values
    
    async def find(self, order_by=None, limit: int = None, offset: int = None, **filters) -> List[SchemaType]:
        """Query items with reliable connection handling."""
        where_clause, values = self._build_where_clause(filters)

        # Build query
        query = f"SELECT * FROM {self.table_name} {where_clause}"

        # Add ORDER BY clause if specified
        if order_by:
            order_fields = [order_by] if isinstance(order_by, str) else order_by
            order_clauses = [
                f"{field[1:]} DESC" if field.startswith('-') else f"{field} ASC"
                for field in order_fields
            ]
            query += f" ORDER BY {', '.join(order_clauses)}"

        # Add LIMIT/OFFSET (SQLite requires LIMIT before OFFSET)
        if limit is not None:
            query += " LIMIT ?"
            values.append(limit)
        elif offset is not None:
            query += " LIMIT -1"
        if offset is not None:
            query += " OFFSET ?"
            values.append(offset)

        # Execute query with reliable transaction
        async with self.transaction() as db:
            cursor = await db.execute(query, values)
            rows = await cursor.fetchall()
            
            if not rows:
                return []
                
            # Process results
            columns = [desc[0] for desc in cursor.description]
            return [
                self.schema_class(**{
                    col: self._deserialize_value(col, row[i])
                    for i, col in enumerate(columns)
                })
                for row in rows
            ]
    
    async def count(self, **filters) -> int:
        """Count items matching filters with reliable connection handling."""
        where_clause, values = self._build_where_clause(filters)
        query = f"SELECT COUNT(*) FROM {self.table_name} {where_clause}"
        
        async with self.transaction() as db:
            cursor = await db.execute(query, values)
            result = await cursor.fetchone()
            return result[0] if result else 0
    
    async def fetch_all(self) -> List[SchemaType]:
        """Retrieve all items."""
        return await self.find()
    
    async def delete(self, id: str) -> bool:
        """Delete an item by ID with reliable transaction handling."""
        async with self._write_lock:
            async with self.transaction() as db:
                cursor = await db.execute(f"DELETE FROM {self.table_name} WHERE id = ?", (id,))
                return cursor.rowcount > 0

    async def exists(self, **filters) -> bool:
        """Check if any record matches the filters without fetching data."""
        return await self.count(**filters) > 0

    async def delete_many(self, **filters) -> int:
        """Delete all items matching filters. Returns count of deleted rows."""
        if not filters:
            raise ValueError("delete_many() requires at least one filter to prevent accidental full table delete")
        where_clause, values = self._build_where_clause(filters)
        async with self._write_lock:
            async with self.transaction() as db:
                cursor = await db.execute(f"DELETE FROM {self.table_name} {where_clause}", values)
                return cursor.rowcount

    async def update_fields(self, id: str, **fields) -> bool:
        """Update specific fields on a record by ID without fetching the full record."""
        if not fields:
            return False
        fields['updated_at'] = datetime.now()
        set_clause = ', '.join(f"{_validate_identifier(k)} = ?" for k in fields)
        values = [self._serialize_value(v) for v in fields.values()]
        values.append(id)
        async with self._write_lock:
            async with self.transaction() as db:
                cursor = await db.execute(
                    f"UPDATE {self.table_name} SET {set_clause} WHERE id = ?", values
                )
                return cursor.rowcount > 0