"""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 re
import sqlite3
from typing import Any

from pydantic import Field

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

_SQLITE_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")


def validate_sqlite_identifier(value: str, *, field_name: str = "identifier") -> str:
    """Validate a SQLite identifier used without quoting.

    AgentByte currently interpolates table/column identifiers directly into SQL
    statements, so only simple unquoted SQLite identifiers are supported.
    """
    if not _SQLITE_IDENTIFIER_PATTERN.fullmatch(value):
        raise ValueError(
            f"{field_name} must be a simple SQLite identifier using only letters, "
            f"digits, and underscores, and cannot start with a digit: {value!r}"
        )
    return value


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."""
        validate_sqlite_identifier(table_name, field_name="table_name")
        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
        """
        if config.engine != "sqlite":
            raise ValueError(
                f"SQLiteDataset requires engine='sqlite', got {config.engine!r}."
            )
        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 _load_schema(self) -> SqliteSchema:
        """Parse schema.json into a SqliteSchema object.

        Raises:
            FileNotFoundError: If schema_text is absent (schema.json was not found).
        """
        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)
        return SqliteSchema(columns=[SqliteColumn(**col) for col in schema_data])

    def _open_connection(self) -> sqlite3.Connection:
        """Create and configure a new sqlite3 connection."""
        conn = sqlite3.connect(":memory:", check_same_thread=False)
        conn.row_factory = sqlite3.Row
        return conn

    def _populate_table(
        self, conn: sqlite3.Connection, schema: SqliteSchema
    ) -> None:
        """Create the table and bulk-insert all records.

        Args:
            conn: An open sqlite3 connection.
            schema: The parsed SqliteSchema used to generate DDL and INSERT SQL.

        Raises:
            ValueError: If data.json is missing the 'records' key.
        """
        data = json.loads(self.data_text)
        if "records" not in data:
            raise ValueError(
                f"Data file must have 'records' key for dataset: {self.source_uri}/data.json"
            )

        conn.execute(schema.to_create_table_sql(self.config.table_name))

        records = data["records"]
        if records:
            col_names = schema.get_column_names()
            placeholders = ", ".join(["?"] * len(col_names))
            table = validate_sqlite_identifier(self.config.table_name, field_name="table_name")
            insert_sql = (
                f"INSERT INTO {table} ({', '.join(col_names)}) VALUES ({placeholders})"
            )
            for record in records:
                conn.execute(insert_sql, tuple(record.get(col) for col in col_names))

        conn.commit()

    def connect(self) -> sqlite3.Connection:
        """Get a connection to the SQLite dataset.

        On the first call: loads schema, creates the in-memory table, and
        bulk-inserts all records. Subsequent calls return the cached connection.

        Returns:
            sqlite3.Connection ready to query.
        """
        if self._connection is not None:
            return self._connection

        schema = self._load_schema()
        self._schema = schema
        conn = self._open_connection()
        self._populate_table(conn, schema)
        self._connection = conn
        return conn

    def fetch_all(self) -> list[dict[str, Any]]:
        """Return all rows as a list of plain dicts.

        Calls ``connect()`` if not already connected, then fetches every row
        from the table and converts each ``sqlite3.Row`` to a plain dict.

        Returns:
            List of dicts with string column names as keys.
        """
        conn = self.connect()
        table_name = validate_sqlite_identifier(self.config.table_name, field_name="table_name")
        cursor = conn.execute(f"SELECT * FROM {table_name}")  # noqa: S608
        return [dict(row) for row in cursor.fetchall()]

    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})"
        )
