"""
TacoDataset - metadata container with lazy SQL interface.

Provides STAC-like metadata with DuckDB connection for lazy SQL queries.
Queries are not executed until .data or .collect() is called.
"""

import uuid
from typing import Any, Literal

import polars as pl
from pydantic import BaseModel, ConfigDict, PrivateAttr

from tacoreader.schema import PITSchema


# ============================================================================
# TACODATASET
# ============================================================================


class TacoDataset(BaseModel):
    """
    TACO dataset with lazy SQL interface.
    
    Container for dataset metadata with lazy query capabilities via DuckDB.
    Queries create views without materializing data until explicitly requested.
    
    Public attributes (STAC-like metadata):
        id: Dataset identifier
        version: Dataset version string
        description: Dataset description
        tasks: Machine learning task types
        extent: Spatial and temporal boundaries
        providers: Dataset creators/contributors
        licenses: License identifiers
        title: Human-friendly title (optional)
        curators: Dataset curators (optional)
        keywords: Searchable tags (optional)
        schema: PIT schema for hierarchy structure
    
    Private attributes:
        _path: Original dataset path
        _format: Backend format type ("zip", "folder", "tacocat")
        _collection: Full COLLECTION.json content
        _consolidated_files: Cached metadata file paths
        _duckdb: DuckDB in-memory connection
        _view_name: Current view name for queries
        _root_path: VSI root path for file access
    
    Example:
        >>> dataset = load("data.tacozip")
        >>> print(dataset.id)
        'sentinel2-l2a'
        >>> 
        >>> # Lazy queries (not executed)
        >>> peru = dataset.sql("SELECT * FROM data WHERE country = 'Peru'")
        >>> low_cloud = peru.sql("SELECT * FROM data WHERE cloud_cover < 10")
        >>> 
        >>> # Materialization (executed)
        >>> df = low_cloud.data
        >>> tdf = low_cloud.collect()
    """
    
    # Public metadata fields (STAC-like)
    id: str
    version: str
    description: str
    tasks: list[str]
    extent: dict[str, Any]
    providers: list[dict[str, Any]]
    licenses: list[str]
    title: str | None = None
    curators: list[dict[str, Any]] | None = None
    keywords: list[str] | None = None
    schema: PITSchema

    # Private attributes (internal state)
    _path: str = PrivateAttr()
    _format: Literal["zip", "folder", "tacocat"] = PrivateAttr()
    _collection: dict[str, Any] = PrivateAttr()
    _consolidated_files: dict[str, str] = PrivateAttr()
    _duckdb: Any = PrivateAttr(default=None)
    _view_name: str = PrivateAttr(default="data")
    _root_path: str = PrivateAttr(default="")

    model_config = ConfigDict(arbitrary_types_allowed=True)

    # ========================================================================
    # PROPERTIES
    # ========================================================================

    @property
    def data(self) -> pl.DataFrame:
        """
        Materialize current view to Polars DataFrame.
        
        Executes DuckDB query and returns materialized data in memory.
        This is the point where lazy evaluation ends and data is loaded.
        
        Returns:
            Polars DataFrame with query results
        
        Example:
            >>> df = dataset.data  # Materializes all data
            >>> filtered = dataset.sql("SELECT * FROM data LIMIT 10").data
        """
        return self._duckdb.execute(
            f"SELECT * FROM {self._view_name}"
        ).pl()
    
    @property
    def field_schema(self) -> dict[str, Any]:
        """
        Field schema from collection metadata.
        
        Returns:
            Dictionary mapping level names to field definitions
            Example: {"level0": [["id", "string"], ["type", "string"], ...]}
        """
        return self._collection.get("taco:field_schema", {})

    # ========================================================================
    # LAZY SQL INTERFACE
    # ========================================================================

    def sql(self, query: str) -> "TacoDataset":
        """
        Execute SQL query and return new TacoDataset with lazy view.
        
        Query is NOT executed immediately. Creates a temporary view in DuckDB
        that references the current view. Multiple queries can be chained.
        
        Always use 'data' as the table name in queries, even when chaining.
        The method automatically replaces 'data' with the current view name.
        
        Args:
            query: SQL query string (use table name 'data')
        
        Returns:
            New TacoDataset with filtered view (still lazy)
        
        Example:
            >>> # Single query
            >>> peru = dataset.sql("SELECT * FROM data WHERE country = 'Peru'")
            >>> 
            >>> # Chained queries (both lazy)
            >>> low_cloud = peru.sql("SELECT * FROM data WHERE cloud_cover < 10")
            >>> 
            >>> # Materialization happens here
            >>> df = low_cloud.data
        """
        new_view_name = f"view_{uuid.uuid4().hex[:8]}"
        
        # Replace 'data' with current view for chaining
        modified_query = query.replace("FROM data", f"FROM {self._view_name}")
        
        # Create temporary view (lazy)
        self._duckdb.execute(f"""
            CREATE TEMP VIEW {new_view_name} AS
            {modified_query}
        """)
        
        # Count rows for schema update
        new_n = self._duckdb.execute(
            f"SELECT COUNT(*) FROM {new_view_name}"
        ).fetchone()[0]
        
        # Update schema with new count
        new_schema = self.schema.with_n(new_n)
        
        # Return new TacoDataset with same connection, different view
        return TacoDataset.model_construct(
            id=self.id,
            version=self.version,
            description=self.description,
            tasks=self.tasks,
            extent=self.extent,
            providers=self.providers,
            licenses=self.licenses,
            title=self.title,
            curators=self.curators,
            keywords=self.keywords,
            schema=new_schema,
            _path=self._path,
            _format=self._format,
            _collection=self._collection,
            _consolidated_files=self._consolidated_files,
            _duckdb=self._duckdb,  # Shared connection
            _view_name=new_view_name,  # New view
            _root_path=self._root_path,
        )

    # ========================================================================
    # MATERIALIZATION
    # ========================================================================

    def collect(self) -> "TacoDataFrame":
        """
        Materialize to TacoDataFrame with navigation methods.
        
        Convenience wrapper that materializes data and wraps it in
        TacoDataFrame for hierarchical navigation and stats methods.
        
        Returns:
            TacoDataFrame with read() and stats_*() methods
        
        Example:
            >>> tdf = dataset.collect()
            >>> tdf.stats_mean()
            >>> child = tdf.read("sample_001")
        """
        return self.to_dataframe()
    
    def to_dataframe(self) -> "TacoDataFrame":
        """
        Materialize to TacoDataFrame (explicit name).
        
        Alias for .collect() with more explicit naming.
        
        Returns:
            TacoDataFrame with TACO navigation methods
        """
        from tacoreader.dataframe import TacoDataFrame
        
        pdf = self.data
        
        return TacoDataFrame(
            data=pdf,
            format_type=self._format,
            consolidated_files=self._consolidated_files,
        )

    # ========================================================================
    # REPRESENTATION
    # ========================================================================

    def __repr__(self) -> str:
        """
        Rich text representation of dataset metadata.
        
        Returns:
            Multi-line string with tree-like structure showing key metadata
        """
        lines = [f"<TacoDataset '{self.id}'>"]
        lines.append(f"├── Version: {self.version}")
        
        desc_short = self.description[:80] + "..." if len(self.description) > 80 else self.description
        lines.append(f"├── Description: {desc_short}")
        lines.append(f"├── Tasks: {', '.join(self.tasks)}")

        if "spatial" in self.extent:
            spatial = self.extent["spatial"]
            lines.append(
                f"├── Spatial Extent: [{spatial[0]:.1f}, {spatial[1]:.1f}, "
                f"{spatial[2]:.1f}, {spatial[3]:.1f}]"
            )

        if "temporal" in self.extent and self.extent["temporal"]:
            temporal = self.extent["temporal"]
            lines.append(f"├── Temporal Extent: {temporal[0]} -> {temporal[1]}")

        lines.append("│")
        lines.append(f"└── Level 0: {self.schema.root['n']} rows")

        return "\n".join(lines)

    def _repr_html_(self):
        """
        Rich HTML representation for Jupyter notebooks.
        
        Returns:
            HTML string with interactive display
        """
        from tacoreader._html import build_html_repr
        return build_html_repr(self)