"""
Database schema definition for KuzuMemory.

Defines the Kuzu graph database schema including node tables, relationship tables,
and indices for optimal performance. Includes version management for migrations.
"""

from __future__ import annotations

from pathlib import Path

# Database schema version for migration support
SCHEMA_VERSION = "1.0"

# Core schema definition
SCHEMA_DDL = """
CREATE NODE TABLE IF NOT EXISTS SchemaVersion (
    version STRING PRIMARY KEY,
    created_at TIMESTAMP,
    description STRING
);

CREATE NODE TABLE IF NOT EXISTS Memory (
    id STRING PRIMARY KEY,
    content STRING,
    content_hash STRING,
    created_at TIMESTAMP,
    valid_from TIMESTAMP,
    valid_to TIMESTAMP,
    accessed_at TIMESTAMP,
    access_count INT32 DEFAULT 0,
    memory_type STRING,
    knowledge_type STRING DEFAULT 'note',
    project_tag STRING DEFAULT '',
    importance FLOAT DEFAULT 0.5,
    confidence FLOAT DEFAULT 1.0,
    source_type STRING DEFAULT 'conversation',
    agent_id STRING DEFAULT 'default',
    user_id STRING,
    session_id STRING,
    metadata STRING DEFAULT '{}'
);

CREATE NODE TABLE IF NOT EXISTS Entity (
    id STRING PRIMARY KEY,
    name STRING,
    entity_type STRING,
    normalized_name STRING,
    first_seen TIMESTAMP,
    last_seen TIMESTAMP,
    mention_count INT32 DEFAULT 1,
    confidence FLOAT DEFAULT 1.0
);

CREATE NODE TABLE IF NOT EXISTS Session (
    id STRING PRIMARY KEY,
    user_id STRING,
    agent_id STRING DEFAULT 'default',
    created_at TIMESTAMP,
    last_activity TIMESTAMP,
    memory_count INT32 DEFAULT 0,
    metadata STRING DEFAULT '{}'
);

CREATE REL TABLE IF NOT EXISTS MENTIONS (
    FROM Memory TO Entity,
    confidence FLOAT DEFAULT 1.0,
    position_start INT32,
    position_end INT32,
    extraction_method STRING DEFAULT 'pattern'
);

CREATE REL TABLE IF NOT EXISTS RELATES_TO (
    FROM Memory TO Memory,
    relationship_type STRING,
    strength FLOAT DEFAULT 1.0,
    created_at TIMESTAMP
);

CREATE REL TABLE IF NOT EXISTS BELONGS_TO_SESSION (
    FROM Memory TO Session,
    created_at TIMESTAMP
);

CREATE REL TABLE IF NOT EXISTS CO_OCCURS_WITH (
    FROM Entity TO Entity,
    co_occurrence_count INT32 DEFAULT 1,
    last_co_occurrence TIMESTAMP
);

CREATE REL TABLE IF NOT EXISTS CONSOLIDATED_INTO (
    FROM Memory TO Memory,
    consolidation_date TIMESTAMP,
    cluster_id STRING,
    similarity_score FLOAT
);

CREATE NODE TABLE IF NOT EXISTS ArchivedMemory (
    id STRING PRIMARY KEY,
    original_id STRING,
    content STRING,
    memory_type STRING,
    knowledge_type STRING DEFAULT 'note',
    project_tag STRING DEFAULT '',
    source_type STRING,
    importance FLOAT,
    created_at TIMESTAMP,
    archived_at TIMESTAMP,
    prune_score FLOAT,
    prune_reason STRING,
    expires_at TIMESTAMP
);
"""

# Performance indices
# NOTE: Kuzu does not support traditional secondary indexes (CREATE INDEX).
# Instead, it uses automatic optimizations:
# - Hash indexes on primary keys (automatic)
# - CSR-based adjacency list indices for edges (automatic)
# - Columnar storage with vectorized execution (automatic)
# - Specialized indexes via function calls: CREATE_FTS_INDEX, CREATE_VECTOR_INDEX
INDICES_DDL = ""

# Initial data insertion
INITIAL_DATA_DDL = f"""
CREATE (sv:SchemaVersion {{version: '{SCHEMA_VERSION}', description: 'Initial KuzuMemory schema with memory, entity, and session tracking'}});
"""

# Full schema (combines all DDL)
FULL_SCHEMA_DDL = SCHEMA_DDL + "\n" + INDICES_DDL + "\n" + INITIAL_DATA_DDL

# Common queries for schema operations
SCHEMA_QUERIES = {
    "get_schema_version": """
        MATCH (sv:SchemaVersion)
        RETURN sv.version, sv.created_at, sv.description
        ORDER BY sv.created_at DESC
        LIMIT 1
    """,
    "check_table_exists": """
        CALL SHOW_TABLES()
        RETURN name WHERE name = $table_name
    """,
    "get_memory_count": """
        MATCH (m:Memory)
        RETURN COUNT(m) as count
    """,
    "get_entity_count": """
        MATCH (e:Entity)
        RETURN COUNT(e) as count
    """,
    "get_session_count": """
        MATCH (s:Session)
        RETURN COUNT(s) as count
    """,
    "get_database_stats": """
        MATCH (m:Memory)
        WITH COUNT(m) as memory_count
        MATCH (e:Entity)
        WITH memory_count, COUNT(e) as entity_count
        MATCH (s:Session)
        WITH memory_count, entity_count, COUNT(s) as session_count
        MATCH ()-[r]->()
        RETURN memory_count, entity_count, session_count, COUNT(r) as relationship_count
    """,
    "cleanup_expired_memories": """
        MATCH (m:Memory)
        WHERE m.valid_to IS NOT NULL AND m.valid_to < $current_time
        DELETE m
    """,
    "get_memory_types_distribution": """
        MATCH (m:Memory)
        WHERE m.valid_to IS NULL OR m.valid_to > $current_time
        RETURN m.memory_type, COUNT(m) as count
        ORDER BY count DESC
    """,
    "get_top_entities": """
        MATCH (e:Entity)
        RETURN e.name, e.entity_type, e.mention_count
        ORDER BY e.mention_count DESC
        LIMIT $limit
    """,
    "find_duplicate_content_hashes": """
        MATCH (m:Memory)
        WITH m.content_hash, COUNT(m) as count, COLLECT(m.id) as memory_ids
        WHERE count > 1
        RETURN m.content_hash, count, memory_ids
    """,
    "get_recent_memories": """
        MATCH (m:Memory)
        WHERE m.created_at > $since_time
        AND (m.valid_to IS NULL OR m.valid_to > $current_time)
        RETURN m
        ORDER BY m.created_at DESC
        LIMIT $limit
    """,
    "get_memories_by_importance": """
        MATCH (m:Memory)
        WHERE (m.valid_to IS NULL OR m.valid_to > $current_time)
        AND m.importance >= $min_importance
        RETURN m
        ORDER BY m.importance DESC, m.created_at DESC
        LIMIT $limit
    """,
}

# Migration queries for schema updates
MIGRATION_QUERIES = {
    "1.0_to_1.1": [
        # Example migration - add new fields or indices
        "ALTER TABLE Memory ADD COLUMN tags STRING DEFAULT '[]'",
        "CREATE INDEX IF NOT EXISTS idx_memory_tags ON Memory(tags)",
    ]
}


def get_schema_ddl() -> str:
    """Get the complete schema DDL."""
    return FULL_SCHEMA_DDL


def get_schema_version() -> str:
    """Get the current schema version."""
    return SCHEMA_VERSION


def get_query(query_name: str) -> str:
    """
    Get a predefined query by name.

    Args:
        query_name: Name of the query to retrieve

    Returns:
        Query string

    Raises:
        KeyError: If query name is not found
    """
    if query_name not in SCHEMA_QUERIES:
        available_queries = ", ".join(SCHEMA_QUERIES.keys())
        raise KeyError(f"Query '{query_name}' not found. Available queries: {available_queries}")

    return SCHEMA_QUERIES[query_name]


def get_migration_queries(from_version: str, to_version: str) -> list[str]:
    """
    Get migration queries between schema versions.

    Args:
        from_version: Source schema version
        to_version: Target schema version

    Returns:
        List of migration queries to execute

    Raises:
        ValueError: If migration path is not supported
    """
    migration_key = f"{from_version}_to_{to_version}"

    if migration_key not in MIGRATION_QUERIES:
        available_migrations = ", ".join(MIGRATION_QUERIES.keys())
        raise ValueError(
            f"Migration from {from_version} to {to_version} not supported. "
            f"Available migrations: {available_migrations}"
        )

    return MIGRATION_QUERIES[migration_key]


def validate_schema_compatibility(current_version: str, required_version: str) -> bool:
    """
    Check if current schema version is compatible with required version.

    Args:
        current_version: Current database schema version
        required_version: Required schema version

    Returns:
        True if compatible, False otherwise
    """
    # For now, only exact version matches are supported
    # In the future, this could support backward compatibility rules
    return current_version == required_version


def ensure_indexes(db_path: Path) -> dict[str, bool]:
    """
    Verify database schema and optimization.

    NOTE: Kuzu does not support traditional secondary indexes (CREATE INDEX)
    on properties. Instead, it uses:
    - Automatic hash indexes on primary keys
    - CSR-based adjacency list indices for edges
    - Columnar storage with vectorized execution for performance
    - Specialized indexes via function calls (FTS, vector)

    This function serves as a validation/verification step rather than
    actually creating indexes. It checks that the database schema is
    properly initialized and returns status information.

    Args:
        db_path: Path to the Kuzu database

    Returns:
        Dictionary with verification results:
        - "schema_valid": True if schema is properly initialized
        - "primary_keys_indexed": True (always, automatic in Kuzu)
        - "columnar_storage": True (always, Kuzu's architecture)

    Raises:
        DatabaseError: If database verification fails

    Example:
        >>> results = ensure_indexes(Path("/tmp/test.db"))
        >>> if results["schema_valid"]:
        ...     print("Database schema verified")
    """
    import kuzu

    from kuzu_memory.utils.exceptions import DatabaseError

    results: dict[str, bool] = {}

    try:
        # Create database connection to verify schema
        db = kuzu.Database(str(db_path))
        conn = kuzu.Connection(db)

        # Verify schema by checking critical tables exist
        try:
            # Check Memory table exists
            conn.execute("MATCH (m:Memory) RETURN COUNT(m) LIMIT 1")
            results["schema_valid"] = True

            # Kuzu automatically provides these optimizations
            results["primary_keys_indexed"] = True  # Hash index on primary keys (automatic)
            results["columnar_storage"] = True  # Kuzu's architecture (automatic)
            results["vectorized_execution"] = True  # Kuzu's query processor (automatic)

        except Exception as e:
            # Schema not initialized or corrupted
            error_msg = str(e).lower()
            if "memory does not exist" in error_msg or "no node table" in error_msg:
                results["schema_valid"] = False
                results["primary_keys_indexed"] = False
                results["columnar_storage"] = False
                results["vectorized_execution"] = False
            else:
                raise DatabaseError(f"Failed to verify schema: {e}") from e

    except Exception as e:
        if isinstance(e, DatabaseError):
            raise
        raise DatabaseError(f"Failed to ensure indexes on {db_path}: {e}") from e

    return results
