"""SQLite dataset and schema implementations.

Provides SQLite-specific classes for dataset management:
- SqliteColumn: Column definition with SQL constraints
- SqliteSchema: Schema with SQL generation and display methods
- SQLiteDataset: SQLite dataset with JSON data loading
"""

import json
import sqlite3
from typing import Any

from pydantic import Field

from agentbyte.dataset.base import (
    BaseDataset,
    Column,
    DataRecord,
    DataStore,
    DatasetConfig,
    Schema,
    SchemaFormat,
)


class SqliteColumn(Column):
    """SQLite-specific column definition with SQL constraints.
    
    Extends the generic Column with SQLite-specific constraint definitions
    and SQL generation methods.
    
    Attributes:
        primary_key: Is primary key
        not_null: Is NOT NULL
        unique: Is UNIQUE
    """

    primary_key: bool = Field(default=False, description="Is primary key")
    not_null: bool = Field(default=False, description="Is NOT NULL")
    unique: bool = Field(default=False, description="Is UNIQUE")

    def to_sql_def(self) -> str:
        """Generate SQL column definition string."""
        parts = [self.name, self.type]
        if self.primary_key:
            parts.append("PRIMARY KEY")
        if self.not_null:
            parts.append("NOT NULL")
        if self.unique:
            parts.append("UNIQUE")
        return " ".join(parts)


class SqliteSchema(Schema):
    """SQLite-specific schema with SQL generation and display options.
    
    Extends the generic Schema with SQLite-specific columns, SQL generation,
    and formatted display of schema information.
    
    Attributes:
        columns: List of SqliteColumn definitions with SQLite-specific constraints
    """

    columns: list[SqliteColumn] = Field(...)

    def to_create_table_sql(self, table_name: str) -> str:
        """Generate CREATE TABLE SQL."""
        col_defs = ",\n    ".join(col.to_sql_def() for col in self.columns)
        return f"CREATE TABLE {table_name} (\n    {col_defs}\n)"

    def get_column_names(self) -> list[str]:
        """Get all column names in order."""
        return [col.name for col in self.columns]

    def show(self, format: SchemaFormat = SchemaFormat.PRETTY) -> None:
        """Display schema in specified format.
        
        Args:
            format: Display format - "raw" (list of dicts) or "pretty" (formatted table)
        """
        if format == SchemaFormat.RAW:
            self._show_raw()
        elif format == SchemaFormat.PRETTY:
            self._show_pretty()
        else:
            raise ValueError(f"Unknown format: {format}")

    def _show_raw(self) -> None:
        """Display schema as raw list of dicts."""
        for col in self.columns:
            col_dict = col.model_dump()
            print(col_dict)

    def _show_pretty(self) -> None:
        """Display schema as formatted table."""
        print(f"Total columns: {len(self.columns)}\n")
        
        # Print header
        print(f"{'Column Name':<25} {'Type':<15} {'Constraints'}")
        print("-" * 70)
        
        # Print each column
        for col in self.columns:
            constraints = []
            if col.primary_key:
                constraints.append("PRIMARY KEY")
            if col.not_null:
                constraints.append("NOT NULL")
            if col.unique:
                constraints.append("UNIQUE")
            
            constraint_str = f"[{', '.join(constraints)}]" if constraints else ""
            print(f"{col.name:<25} {col.type:<15} {constraint_str}")


class SQLiteDataset(BaseDataset):
    """SQLite-specific dataset implementation.
    
    Handles loading JSON data into in-memory or persistent SQLite databases.
    
    Attributes:
        dataset_id: Identifier (e.g., "contracts/asmd")
        config: Dataset configuration
        data_text: Raw data.json content
        schema_text: Raw schema.json content
    """

    def __init__(
        self,
        dataset_id: str,
        config: DatasetConfig,
        data_text: str,
        schema_text: str | None,
        source_uri: str,
    ) -> None:
        """Initialize SQLite dataset.
        
        Args:
            dataset_id: Dataset identifier
            config: DatasetConfig object
            data_text: Raw ``data.json`` content
            schema_text: Raw ``schema.json`` content when present
            source_uri: Human-readable dataset source location
        """
        self.dataset_id = dataset_id
        self.config = config
        self.data_text = data_text
        self.schema_text = schema_text
        self.source_uri = source_uri
        self._connection: sqlite3.Connection | None = None
        self._schema: SqliteSchema | None = None

    @property
    def engine(self) -> str:
        """Get the database engine name."""
        return self.config.engine

    @property
    def in_memory(self) -> bool:
        """Whether this is an in-memory database."""
        return self.config.in_memory

    @property
    def table_name(self) -> str:
        """Get the table name."""
        return self.config.table_name

    @property
    def schema(self) -> SqliteSchema | None:
        """Get the table schema (loaded on first connect() call).
        
        Returns a SqliteSchema object with methods to display the schema
        in different formats: schema.show() or schema.show(format="raw")
        """
        return self._schema

    def connect(self) -> sqlite3.Connection:
        """Get a connection to the SQLite dataset.
        
        Creates/opens in-memory or file-based connection.
        Loads schema from schema.json and data from data.json on first call.
        
        Returns:
            sqlite3.Connection ready to query
        """
        if self._connection is not None:
            return self._connection

        if self.config.engine != "sqlite":
            raise NotImplementedError(
                f"Engine '{self.config.engine}' not yet supported. "
                "Currently only 'sqlite' is implemented."
            )

        if self.schema_text is None:
            raise FileNotFoundError(
                f"Schema file not found for dataset: {self.source_uri}/schema.json"
            )

        schema_data = json.loads(self.schema_text)
        
        # Create SqliteSchema object for display and SQL generation
        schema_obj = SqliteSchema(columns=[SqliteColumn(**col) for col in schema_data])
        self._schema = schema_obj

        data = json.loads(self.data_text)

        # Parse records (data.json should only have "records")
        if "records" not in data:
            raise ValueError(
                f"Data file must have 'records' key for dataset: {self.source_uri}/data.json"
            )

        records = data["records"]

        # Create SQLite connection
        if self.config.in_memory:
            conn = sqlite3.connect(":memory:")
        else:
            # Future: support persistent SQLite
            conn = sqlite3.connect(":memory:")

        conn.row_factory = sqlite3.Row

        # Create table
        create_sql = schema_obj.to_create_table_sql(self.config.table_name)
        conn.execute(create_sql)

        # Insert data
        if records:
            col_names = schema_obj.get_column_names()
            placeholders = ", ".join(["?"] * len(col_names))
            insert_sql = (
                f"INSERT INTO {self.config.table_name} "
                f"({', '.join(col_names)}) VALUES ({placeholders})"
            )

            for record in records:
                values = tuple(record.get(col) for col in col_names)
                conn.execute(insert_sql, values)

        conn.commit()
        self._connection = conn
        return conn

    def __repr__(self) -> str:
        """String representation of dataset."""
        return (
            f"SQLiteDataset(id='{self.dataset_id}', engine='{self.engine}', "
            f"table='{self.table_name}', in_memory={self.in_memory})"
        )


class InMemorySQLiteStore(DataStore[DataRecord]):
    """In-memory SQLite-based data store.
    
    Uses SQLite's in-memory mode (":memory:") for fast, zero-persistence storage.
    Suitable for development, testing, and scenarios where data fits in RAM.
    
    Schema:
        - id (TEXT PRIMARY KEY): Record identifier
        - data (TEXT): JSON-serialized data content
        - metadata (TEXT): JSON-serialized metadata
    """

    def __init__(self, name: str = "default") -> None:
        """Initialize in-memory SQLite store.
        
        Args:
            name: Optional name for this store instance
        """
        super().__init__(name)
        self.connection = sqlite3.connect(":memory:")
        self.cursor = self.connection.cursor()
        self._init_schema()

    def _init_schema(self) -> None:
        """Create the default schema for storing records."""
        self.cursor.execute("""
            CREATE TABLE IF NOT EXISTS records (
                id TEXT PRIMARY KEY,
                data TEXT NOT NULL,
                metadata TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        self.connection.commit()

    async def add(self, record: DataRecord) -> str:
        """Add a single record to the store.
        
        Args:
            record: The record to add
            
        Returns:
            The record ID
        """
        data_json = record.data if isinstance(record.data, str) else str(record.data)
        metadata_json = str(record.metadata)
        
        self.cursor.execute(
            """INSERT INTO records (id, data, metadata) VALUES (?, ?, ?)""",
            (record.id, data_json, metadata_json),
        )
        self.connection.commit()
        self.record_count += 1
        return record.id

    async def add_batch(self, records: list[DataRecord]) -> list[str]:
        """Add multiple records to the store.
        
        Args:
            records: List of records to add
            
        Returns:
            List of record IDs
        """
        record_ids = []
        for record in records:
            record_id = await self.add(record)
            record_ids.append(record_id)
        return record_ids

    async def get(self, record_id: str) -> DataRecord | None:
        """Retrieve a record by ID.
        
        Args:
            record_id: The ID of the record
            
        Returns:
            The record if found, None otherwise
        """
        self.cursor.execute(
            """SELECT id, data, metadata FROM records WHERE id = ?""",
            (record_id,),
        )
        row = self.cursor.fetchone()
        if not row:
            return None
        
        return DataRecord(
            id=row[0],
            data=row[1],
            metadata=eval(row[2]),  # Simple eval for demo; use json.loads in production
        )

    async def query(
        self,
        query_text: str,
        limit: int = 10,
        filters: dict[str, Any] | None = None,
    ) -> list[DataRecord]:
        """Query records by exact match on data field.
        
        Args:
            query_text: Text to search for in data field
            limit: Maximum results to return
            filters: Optional additional filters (not used in basic implementation)
            
        Returns:
            List of matching records
        """
        sql = """SELECT id, data, metadata FROM records 
                 WHERE data LIKE ? LIMIT ?"""
        self.cursor.execute(sql, (f"%{query_text}%", limit))
        rows = self.cursor.fetchall()
        
        results = []
        for row in rows:
            results.append(
                DataRecord(
                    id=row[0],
                    data=row[1],
                    metadata=eval(row[2]),
                )
            )
        return results

    async def delete(self, record_id: str) -> bool:
        """Delete a record from the store.
        
        Args:
            record_id: The ID of the record to delete
            
        Returns:
            True if record was deleted, False if not found
        """
        self.cursor.execute("""DELETE FROM records WHERE id = ?""", (record_id,))
        self.connection.commit()
        deleted = self.cursor.rowcount > 0
        if deleted:
            self.record_count -= 1
        return deleted

    async def clear(self) -> None:
        """Delete all records from the store."""
        self.cursor.execute("""DELETE FROM records""")
        self.connection.commit()
        self.record_count = 0

    async def count(self) -> int:
        """Get total number of records in the store.
        
        Returns:
            Number of records
        """
        self.cursor.execute("""SELECT COUNT(*) FROM records""")
        return self.cursor.fetchone()[0]

    def __del__(self) -> None:
        """Clean up database connection on deletion."""
        if hasattr(self, "connection"):
            self.connection.close()
