"""
Metadata storage using DuckDB.

Provides persistent storage for datasets, pipelines, and execution logs.
"""

from datetime import datetime
from pathlib import Path
from typing import Any
from uuid import UUID

import duckdb

from app.core.config import settings
from app.core.exceptions import DatabaseError
from app.core.logging import get_logger
from app.domain.models.dataset import Dataset, DatasetStatus
from app.domain.models.pipeline import ( 
    NodeExecutionLog,
    Pipeline,
    PipelineStatus,
)

logger = get_logger(__name__)


class MetadataStore:
    """DuckDB-based metadata storage."""

    def __init__(self, db_path: str | None = None):
        """Initialize metadata store."""
        self.db_path = db_path or settings.duckdb_path
        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
        self._initialize_schema()

    def _get_connection(self) -> duckdb.DuckDBPyConnection:
        """Get database connection."""
        return duckdb.connect(self.db_path)

    def _initialize_schema(self) -> None:
        """Initialize database schema."""
        try:
            conn = self._get_connection()
            
            # Datasets table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS datasets (
                    id VARCHAR PRIMARY KEY,
                    name VARCHAR NOT NULL,
                    original_filename VARCHAR NOT NULL,
                    file_path VARCHAR NOT NULL,
                    file_format VARCHAR NOT NULL,
                    file_size_bytes BIGINT NOT NULL,
                    status VARCHAR NOT NULL,
                    schema JSON,
                    row_count BIGINT DEFAULT 0,
                    column_count INTEGER DEFAULT 0,
                    quality_metrics JSON,
                    profile JSON,
                    uploaded_at TIMESTAMP NOT NULL,
                    updated_at TIMESTAMP NOT NULL,
                    uploaded_by VARCHAR DEFAULT 'system',
                    tags JSON,
                    description TEXT
                )
            """)

            # Pipelines table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS pipelines (
                    id VARCHAR PRIMARY KEY,
                    name VARCHAR NOT NULL,
                    source_dataset_id VARCHAR NOT NULL,
                    result_dataset_id VARCHAR,
                    ir JSON NOT NULL,
                    natural_language_input TEXT,
                    status VARCHAR NOT NULL,
                    execution_logs JSON,
                    execution_metrics JSON,
                    error_message TEXT,
                    current_version INTEGER DEFAULT 1,
                    versions JSON,
                    created_at TIMESTAMP NOT NULL,
                    updated_at TIMESTAMP NOT NULL,
                    executed_at TIMESTAMP,
                    created_by VARCHAR DEFAULT 'system',
                    tags JSON,
                    FOREIGN KEY (source_dataset_id) REFERENCES datasets(id)
                )
            """)

            # Execution logs table (denormalized for querying)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS execution_logs (
                    id VARCHAR PRIMARY KEY,
                    pipeline_id VARCHAR NOT NULL,
                    node_id VARCHAR NOT NULL,
                    operation VARCHAR NOT NULL,
                    status VARCHAR NOT NULL,
                    error_message TEXT,
                    execution_time_ms BIGINT,
                    rows_before BIGINT,
                    rows_after BIGINT,
                    columns_affected JSON,
                    memory_usage_mb DOUBLE,
                    executed_at TIMESTAMP NOT NULL,
                    FOREIGN KEY (pipeline_id) REFERENCES pipelines(id)
                )
            """)

            # Lineage table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS lineage (
                    id VARCHAR PRIMARY KEY,
                    source_dataset_id VARCHAR NOT NULL,
                    target_dataset_id VARCHAR,
                    pipeline_id VARCHAR NOT NULL,
                    transformation_summary JSON,
                    created_at TIMESTAMP NOT NULL,
                    FOREIGN KEY (source_dataset_id) REFERENCES datasets(id),
                    FOREIGN KEY (pipeline_id) REFERENCES pipelines(id)
                )
            """)

            conn.close()
            logger.info("Database schema initialized", db_path=self.db_path)
        except Exception as e:
            logger.error("Schema initialization failed", error=str(e))
            raise DatabaseError(f"Failed to initialize schema: {e}")

    # Dataset operations
    def save_dataset(self, dataset: Dataset) -> None:
        """Save dataset to storage."""
        try:
            conn = self._get_connection()
            import orjson
            
            conn.execute(
                """
                INSERT OR REPLACE INTO datasets VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                [
                    str(dataset.id),
                    dataset.name,
                    dataset.original_filename,
                    str(dataset.file_path),
                    dataset.file_format,
                    dataset.file_size_bytes,
                    dataset.status.value,
                    orjson.dumps(dataset.schema.model_dump() if dataset.schema else None).decode(),
                    dataset.row_count,
                    dataset.column_count,
                    orjson.dumps(dataset.quality_metrics.model_dump() if dataset.quality_metrics else None).decode(),
                    orjson.dumps(dataset.profile.model_dump() if dataset.profile else None).decode(),
                    dataset.uploaded_at,
                    dataset.updated_at,
                    dataset.uploaded_by,
                    orjson.dumps(dataset.tags).decode(),
                    dataset.description,
                ],
            )
            conn.close()
            logger.info("Dataset saved", dataset_id=str(dataset.id))
        except Exception as e:
            logger.error("Dataset save failed", dataset_id=str(dataset.id), error=str(e))
            raise DatabaseError(f"Failed to save dataset: {e}")

    def get_dataset(self, dataset_id: UUID) -> Dataset | None:
        """Get dataset by ID."""
        try:
            conn = self._get_connection()
            result = conn.execute(
                "SELECT * FROM datasets WHERE id = ?", [str(dataset_id)]
            ).fetchone()
            conn.close()

            if not result:
                return None

            import orjson
            from app.domain.models.ir import DatasetSchema
            from app.domain.models.dataset import DataQualityMetrics, DatasetProfile

            return Dataset(
                id=UUID(result[0]),
                name=result[1],
                original_filename=result[2],
                file_path=Path(result[3]),
                file_format=result[4],
                file_size_bytes=result[5],
                status=DatasetStatus(result[6]),
                schema=DatasetSchema(**orjson.loads(result[7])) if result[7] else None,
                row_count=result[8],
                column_count=result[9],
                quality_metrics=DataQualityMetrics(**orjson.loads(result[10])) if result[10] else None,
                profile=DatasetProfile(**orjson.loads(result[11])) if result[11] else None,
                uploaded_at=result[12],
                updated_at=result[13],
                uploaded_by=result[14],
                tags=orjson.loads(result[15]) if result[15] else [],
                description=result[16],
            )
        except Exception as e:
            logger.error("Dataset retrieval failed", dataset_id=str(dataset_id), error=str(e))
            raise DatabaseError(f"Failed to get dataset: {e}")

    def list_datasets(self, limit: int = 100, offset: int = 0) -> list[Dataset]:
        """List all datasets."""
        try:
            conn = self._get_connection()
            results = conn.execute(
                "SELECT * FROM datasets ORDER BY uploaded_at DESC LIMIT ? OFFSET ?",
                [limit, offset],
            ).fetchall()
            conn.close()

            datasets = []
            import orjson
            from app.domain.models.ir import DatasetSchema
            from app.domain.models.dataset import DataQualityMetrics, DatasetProfile

            for row in results:
                datasets.append(
                    Dataset(
                        id=UUID(row[0]),
                        name=row[1],
                        original_filename=row[2],
                        file_path=Path(row[3]),
                        file_format=row[4],
                        file_size_bytes=row[5],
                        status=DatasetStatus(row[6]),
                        schema=DatasetSchema(**orjson.loads(row[7])) if row[7] else None,
                        row_count=row[8],
                        column_count=row[9],
                        quality_metrics=DataQualityMetrics(**orjson.loads(row[10])) if row[10] else None,
                        profile=DatasetProfile(**orjson.loads(row[11])) if row[11] else None,
                        uploaded_at=row[12],
                        updated_at=row[13],
                        uploaded_by=row[14],
                        tags=orjson.loads(row[15]) if row[15] else [],
                        description=row[16],
                    )
                )
            return datasets
        except Exception as e:
            logger.error("Dataset listing failed", error=str(e))
            raise DatabaseError(f"Failed to list datasets: {e}")

    def delete_dataset(self, dataset_id: UUID) -> bool:
        """Delete dataset."""
        try:
            conn = self._get_connection()
            conn.execute("DELETE FROM datasets WHERE id = ?", [str(dataset_id)])
            affected = conn.execute("SELECT changes()").fetchone()[0]
            conn.close()
            logger.info("Dataset deleted", dataset_id=str(dataset_id))
            return affected > 0
        except Exception as e:
            logger.error("Dataset deletion failed", dataset_id=str(dataset_id), error=str(e))
            raise DatabaseError(f"Failed to delete dataset: {e}")

    # Pipeline operations
    def save_pipeline(self, pipeline: Pipeline) -> None:
        """Save pipeline to storage."""
        try:
            conn = self._get_connection()
            import orjson

            conn.execute(
                """
                INSERT OR REPLACE INTO pipelines VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                [
                    str(pipeline.id),
                    pipeline.name,
                    str(pipeline.source_dataset_id),
                    str(pipeline.result_dataset_id) if pipeline.result_dataset_id else None,
                    orjson.dumps(pipeline.ir.to_json()).decode(),
                    pipeline.natural_language_input,
                    pipeline.status.value,
                    orjson.dumps([log.model_dump() for log in pipeline.execution_logs]).decode(),
                    orjson.dumps(pipeline.execution_metrics.model_dump() if pipeline.execution_metrics else None).decode(),
                    pipeline.error_message,
                    pipeline.current_version,
                    orjson.dumps([v.model_dump() for v in pipeline.versions]).decode(),
                    pipeline.created_at,
                    pipeline.updated_at,
                    pipeline.executed_at,
                    pipeline.created_by,
                    orjson.dumps(pipeline.tags).decode(),
                ],
            )
            conn.close()
            logger.info("Pipeline saved", pipeline_id=str(pipeline.id))
        except Exception as e:
            logger.error("Pipeline save failed", pipeline_id=str(pipeline.id), error=str(e))
            raise DatabaseError(f"Failed to save pipeline: {e}")

    def get_pipeline(self, pipeline_id: UUID) -> Pipeline | None:
        """Get pipeline by ID."""
        try:
            conn = self._get_connection()
            result = conn.execute(
                "SELECT * FROM pipelines WHERE id = ?", [str(pipeline_id)]
            ).fetchone()
            conn.close()

            if not result:
                return None

            import orjson
            from app.domain.models.ir import IRSchema
            from app.domain.models.pipeline import ExecutionMetrics, PipelineVersion

            return Pipeline(
                id=UUID(result[0]),
                name=result[1],
                source_dataset_id=UUID(result[2]),
                result_dataset_id=UUID(result[3]) if result[3] else None,
                ir=IRSchema.from_json(orjson.loads(result[4])),
                natural_language_input=result[5] or "",
                status=PipelineStatus(result[6]),
                execution_logs=[NodeExecutionLog(**log) for log in orjson.loads(result[7])],
                execution_metrics=ExecutionMetrics(**orjson.loads(result[8])) if result[8] else None,
                error_message=result[9],
                current_version=result[10],
                versions=[PipelineVersion(**v) for v in orjson.loads(result[11])],
                created_at=result[12],
                updated_at=result[13],
                executed_at=result[14],
                created_by=result[15],
                tags=orjson.loads(result[16]) if result[16] else [],
            )
        except Exception as e:
            logger.error("Pipeline retrieval failed", pipeline_id=str(pipeline_id), error=str(e))
            raise DatabaseError(f"Failed to get pipeline: {e}")

    def list_pipelines(
        self, dataset_id: UUID | None = None, limit: int = 100, offset: int = 0
    ) -> list[Pipeline]:
        """List pipelines, optionally filtered by dataset."""
        try:
            conn = self._get_connection()
            
            if dataset_id:
                query = "SELECT * FROM pipelines WHERE source_dataset_id = ? ORDER BY created_at DESC LIMIT ? OFFSET ?"
                params = [str(dataset_id), limit, offset]
            else:
                query = "SELECT * FROM pipelines ORDER BY created_at DESC LIMIT ? OFFSET ?"
                params = [limit, offset]
            
            results = conn.execute(query, params).fetchall()
            conn.close()

            pipelines = []
            import orjson
            from app.domain.models.ir import IRSchema
            from app.domain.models.pipeline import ExecutionMetrics, PipelineVersion

            for row in results:
                pipelines.append(
                    Pipeline(
                        id=UUID(row[0]),
                        name=row[1],
                        source_dataset_id=UUID(row[2]),
                        result_dataset_id=UUID(row[3]) if row[3] else None,
                        ir=IRSchema.from_json(orjson.loads(row[4])),
                        natural_language_input=row[5] or "",
                        status=PipelineStatus(row[6]),
                        execution_logs=[NodeExecutionLog(**log) for log in orjson.loads(row[7])],
                        execution_metrics=ExecutionMetrics(**orjson.loads(row[8])) if row[8] else None,
                        error_message=row[9],
                        current_version=row[10],
                        versions=[PipelineVersion(**v) for v in orjson.loads(row[11])],
                        created_at=row[12],
                        updated_at=row[13],
                        executed_at=row[14],
                        created_by=row[15],
                        tags=orjson.loads(row[16]) if row[16] else [],
                    )
                )
            return pipelines
        except Exception as e:
            logger.error("Pipeline listing failed", error=str(e))
            raise DatabaseError(f"Failed to list pipelines: {e}")
