# -*- coding: utf-8 -*-
import logging
import re
from typing import Any, Dict, List, Optional, Sequence, Tuple

import jmespath
from pymongo.errors import PyMongoError

from .common import CursorIterator
from .error import DatabaseError, ProgrammingError
from .sql.builder import ExecutionPlan

_logger = logging.getLogger(__name__)


class ResultSet(CursorIterator):
    """Result set wrapper for MongoDB command results"""

    def __init__(
        self,
        command_result: Optional[Dict[str, Any]] = None,
        execution_plan: ExecutionPlan = None,
        arraysize: int = None,
        database: Optional[Any] = None,
        **kwargs,
    ) -> None:
        super().__init__(arraysize=arraysize or self.DEFAULT_FETCH_SIZE, **kwargs)

        # Handle command results from db.command
        if command_result is not None:
            self._command_result = command_result
            self._database = database
            # Extract cursor info from command result
            self._result_cursor = command_result.get("cursor", {})
            self._cursor_id = self._result_cursor.get("id", 0)  # 0 means no more results
            self._raw_results = self._result_cursor.get("firstBatch", [])
            self._cached_results: List[Sequence[Any]] = []
        else:
            raise ProgrammingError("command_result must be provided")

        self._execution_plan = execution_plan
        self._is_closed = False
        self._cache_exhausted = False
        self._total_fetched = 0
        self._description: Optional[List[Tuple[str, str, None, None, None, None, None]]] = None
        self._column_names: Optional[List[str]] = None  # Track column order for sequences
        self._errors: List[Dict[str, str]] = []

        # Process firstBatch immediately if available (after all attributes are set)
        if command_result is not None and self._raw_results:
            self._process_and_cache_batch(self._raw_results)

        # Build description from projection
        self._build_description()

    def _process_and_cache_batch(self, batch: List[Dict[str, Any]]) -> None:
        """Process and cache a batch of documents"""
        if not batch:
            return
        # Process results through projection mapping
        processed_batch = [self._process_document(doc) for doc in batch]
        # Convert dictionaries to output format (sequence or dict)
        formatted_batch = [self._format_result(doc) for doc in processed_batch]
        self._cached_results.extend(formatted_batch)
        self._total_fetched += len(batch)

    def _build_description(self) -> None:
        """Build column description from execution plan projection"""
        if not self._execution_plan.projection_stage:
            # No projection specified, description will be built dynamically
            self._description = None
            return

        # Build description from projection (now in MongoDB format {field: 1})
        description = []
        for field_name, include_flag in self._execution_plan.projection_stage.items():
            # SQL cursor description format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
            if include_flag == 1:  # Field is included in projection
                description.append((field_name, "VARCHAR", None, None, None, None, None))

        self._description = description

    def _ensure_results_available(self, count: int = 1) -> None:
        """Ensure we have at least 'count' results available in cache"""
        if self._is_closed:
            raise ProgrammingError("ResultSet is closed")

        if self._cache_exhausted:
            return

        # Fetch more results if needed and cursor has more data
        while len(self._cached_results) < count and self._cursor_id != 0:
            try:
                # Use getMore to fetch next batch
                if self._database is not None and self._execution_plan.collection:
                    getmore_cmd = {
                        "getMore": self._cursor_id,
                        "collection": self._execution_plan.collection,
                    }
                    result = self._database.command(getmore_cmd)

                    # Extract and process next batch
                    cursor_info = result.get("cursor", {})
                    next_batch = cursor_info.get("nextBatch", [])
                    self._process_and_cache_batch(next_batch)

                    # Update cursor ID for next iteration
                    self._cursor_id = cursor_info.get("id", 0)
                else:
                    # No database access, mark as exhausted
                    self._cache_exhausted = True
                    break

            except PyMongoError as e:
                self._errors.append({"error": str(e), "type": type(e).__name__})
                self._cache_exhausted = True
                raise DatabaseError(f"Error fetching more results: {e}")

        # Mark as exhausted if no more results available
        if self._cursor_id == 0:
            self._cache_exhausted = True

    def _process_document(self, doc: Dict[str, Any]) -> Dict[str, Any]:
        """Process a MongoDB document according to projection mapping"""
        if not self._execution_plan.projection_stage:
            # No projection, return document as-is (including _id)
            return dict(doc)

        # Apply projection mapping (now using MongoDB format {field: 1})
        processed = {}
        for field_name, include_flag in self._execution_plan.projection_stage.items():
            if include_flag == 1:  # Field is included in projection
                # Extract value using jmespath-compatible field path (convert numeric dot indexes to bracket form)
                value = self._get_nested_value(doc, field_name)
                # Convert the projection key back to bracket notation for client-facing results
                display_key = self._mongo_to_bracket_key(field_name)
                processed[display_key] = value

        return processed

    def _mongo_to_bracket_key(self, field_path: str) -> str:
        """Convert Mongo dot-index notation to bracket notation.

        Transforms numeric dot segments into bracket indices for both display keys
        and JMESPath-compatible field paths.

        Examples:
            items.0 -> items[0]
            items.1.name -> items[1].name
        """
        if not isinstance(field_path, str):
            return field_path
        # Replace .<number> with [<number>]
        return re.sub(r"\.(\d+)", r"[\1]", field_path)

    def _get_nested_value(self, doc: Dict[str, Any], field_path: str) -> Any:
        """Extract nested field value from document using JMESPath

        Supports:
            - Simple fields: "name" -> doc["name"]
            - Nested fields: "profile.bio" -> doc["profile"]["bio"]
            - Array indexing: "address.coordinates[1]" -> doc["address"]["coordinates"][1]
            - Wildcards: "items[*].name" -> [item["name"] for item in items]
        """
        try:
            # Optimization: for simple field names without dots/brackets, use direct access
            if "." not in field_path and "[" not in field_path:
                return doc.get(field_path)

            # Convert normalized Mongo-style numeric segments to bracket notation
            normalized_field = self._mongo_to_bracket_key(field_path)
            # Use jmespath for complex paths
            return jmespath.search(normalized_field, doc)
        except Exception as e:
            _logger.debug(f"Error extracting field '{field_path}': {e}")
            return None

    def _format_result(self, doc: Dict[str, Any]) -> Tuple[Any, ...]:
        """Format processed document to output format (tuple for DB API 2.0 compliance)"""
        if self._column_names is None:
            # First time - establish column order
            self._column_names = list(doc.keys())

        # Return values in consistent column order
        return tuple(doc.get(col_name) for col_name in self._column_names)

    @property
    def errors(self) -> List[Dict[str, str]]:
        return self._errors.copy()

    @property
    def rowcount(self) -> int:
        """Return number of rows fetched so far (not total available)"""
        return self._total_fetched

    @property
    def description(
        self,
    ) -> Optional[List[Tuple[str, str, None, None, None, None, None]]]:
        """Return column description"""
        if self._description is None and not self._cache_exhausted:
            # Try to fetch one result to build description dynamically
            try:
                self._ensure_results_available(1)
                if self._column_names:
                    # Build description from established column names
                    self._description = [
                        (col_name, "VARCHAR", None, None, None, None, None) for col_name in self._column_names
                    ]
            except Exception as e:
                _logger.warning(f"Could not build dynamic description: {e}")

        return self._description

    def fetchone(self) -> Optional[Sequence[Any]]:
        """Fetch the next row from the result set"""
        if self._is_closed:
            raise ProgrammingError("ResultSet is closed")

        # Ensure we have at least one result
        self._ensure_results_available(1)

        if not self._cached_results:
            return None

        # Return and remove first result
        result = self._cached_results.pop(0)
        self._rownumber = (self._rownumber or 0) + 1
        return result

    def fetchmany(self, size: Optional[int] = None) -> List[Sequence[Any]]:
        """Fetch up to 'size' rows from the result set"""
        if self._is_closed:
            raise ProgrammingError("ResultSet is closed")

        fetch_size = size or self.arraysize

        # Ensure we have enough results
        self._ensure_results_available(fetch_size)

        # Return requested number of results
        results = self._cached_results[:fetch_size]
        self._cached_results = self._cached_results[fetch_size:]

        # Update row number
        self._rownumber = (self._rownumber or 0) + len(results)

        return results

    def fetchall(self) -> List[Sequence[Any]]:
        """Fetch all remaining rows from the result set"""
        if self._is_closed:
            raise ProgrammingError("ResultSet is closed")

        # Fetch all remaining results
        all_results = []

        try:
            # Ensure all results are available in cache by requesting a very large number
            # This will trigger getMore calls until all data is exhausted
            if not self._cache_exhausted and self._cursor_id != 0:
                self._ensure_results_available(float("inf"))

            # Now get everything from cache
            all_results.extend(self._cached_results)
            self._total_fetched += len(self._cached_results)
            self._cached_results.clear()
            self._cache_exhausted = True

        except PyMongoError as e:
            self._errors.append({"error": str(e), "type": type(e).__name__})
            raise DatabaseError(f"Error fetching all results: {e}")

        # Update row number
        self._rownumber = (self._rownumber or 0) + len(all_results)

        return all_results

    @property
    def is_closed(self) -> bool:
        return self._is_closed

    def close(self) -> None:
        """Close the result set and free resources"""
        if not self._is_closed:
            self._is_closed = True
            self._command_result = None
            self._database = None
            self._cached_results.clear()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class DictResultSet(ResultSet):
    """Result set that returns dictionaries instead of sequences"""

    def _format_result(self, doc: Dict[str, Any]) -> Dict[str, Any]:
        """Override to return dictionary directly instead of converting to sequence"""
        return doc


# For backward compatibility
MongoResultSet = ResultSet
