"""
HTTP handler for managing user's DatasetManager.

Provides REST endpoints for dataset operations:
- GET    /api/v1/agents/datasets/{agent_id}              - List datasets
- PATCH  /api/v1/agents/datasets/{agent_id}              - Activate/deactivate dataset
- PUT    /api/v1/agents/datasets/{agent_id}              - Upload Excel/CSV file
- POST   /api/v1/agents/datasets/{agent_id}              - Add SQL query or query slug
- DELETE /api/v1/agents/datasets/{agent_id}              - Delete dataset
- GET    /api/v1/agents/datasets/{agent_id}/{dataset_id} - Describe a single dataset
"""
from __future__ import annotations

from io import BytesIO
from typing import TYPE_CHECKING

import pandas as pd
from aiohttp import web
from navigator.views import BaseView
from navigator_auth.decorators import is_authenticated, user_session
from navigator_session import get_session
from pydantic import ValidationError

from ..models.datasets import (
    DatasetAction,
    DatasetDeleteResponse,
    DatasetListResponse,
    DatasetPatchRequest,
    DatasetQueryRequest,
    DatasetUploadResponse,
)
from ..tools.dataset_manager import DatasetEntry, DatasetManager
from .user_objects import UserObjectsHandler

if TYPE_CHECKING:
    pass


# Maximum file size: 50 MB
MAX_FILE_SIZE = 50 * 1024 * 1024


# ---------------------------------------------------------------------------
# Module-level helpers shared by both handlers
# ---------------------------------------------------------------------------

def _clone_agent_dm(agent: object, logger: object) -> DatasetManager | None:
    """Clone an agent's DatasetManager into a fresh user-scoped instance.

    Copies both loaded DataFrames and query-slug-based datasets so the
    user starts with the full catalog the agent was configured with.

    Args:
        agent: Agent instance that exposes ``_dataset_manager``.
        logger: Logger used for per-entry warnings.

    Returns:
        A new DatasetManager seeded with the agent's datasets, or
        ``None`` if the agent has no DatasetManager or it is empty.
    """
    agent_dm: DatasetManager | None = getattr(agent, "_dataset_manager", None)
    if not agent_dm or not agent_dm._datasets:
        return None

    user_dm = DatasetManager()
    for name, entry in agent_dm._datasets.items():
        try:
            cloned_entry = DatasetEntry(
                name=name,
                source=entry.source,
                metadata=dict(entry.metadata or {}),
                is_active=entry.is_active,
                auto_detect_types=entry.auto_detect_types,
                cache_ttl=entry.cache_ttl,
            )
            if entry.loaded and entry.df is not None:
                cloned_entry._df = entry.df.copy(deep=True)
            if entry.column_types:
                cloned_entry._column_types = dict(entry.column_types)
            user_dm._datasets[name] = cloned_entry
        except Exception as exc:
            logger.warning("Failed to clone dataset '%s' from agent: %s", name, exc)

    logger.debug(
        "Cloned %d datasets from agent into user DatasetManager",
        len(user_dm._datasets),
    )
    return user_dm


async def _resolve_dataset_manager(
    request: web.Request,
    agent_id: str,
    session_key: str,
    logger: object,
) -> DatasetManager:
    """Get or create a session-scoped DatasetManager, seeding from the agent if empty.

    Returns the existing user DM when it already has datasets (preserves any
    customisations the user made).  When absent or empty, looks up the agent in
    BotManager and clones its catalog.

    Args:
        request: The current aiohttp request.
        agent_id: Agent identifier to look up in BotManager.
        session_key: Session key under which the DM is stored.
        logger: Logger for warnings/debug output.

    Returns:
        The user's DatasetManager (may be empty for a fresh agent with no datasets).

    Raises:
        KeyError: If the agent is not registered in BotManager and there is no
            existing session DatasetManager to fall back to.
    """
    request_session = await get_session(request)

    # Return existing non-empty user DM — user may have uploaded their own data.
    dm = request_session.get(session_key)
    if dm is not None and isinstance(dm, DatasetManager) and dm._datasets:
        return dm

    # No usable session DM — the agent must exist to seed from.
    bot_manager = request.app.get("bot_manager")
    agent = bot_manager._bots.get(agent_id) if bot_manager else None
    if agent is None:
        raise KeyError(agent_id)

    # Seed from agent's catalog (may be empty for a freshly-configured agent).
    dm = _clone_agent_dm(agent, logger) or DatasetManager()
    request_session[session_key] = dm
    return dm


# ---------------------------------------------------------------------------
# Collection handler  — /api/v1/agents/datasets/{agent_id}
# ---------------------------------------------------------------------------

@is_authenticated()
@user_session()
class DatasetManagerHandler(BaseView):
    """HTTP handler for managing a user's DatasetManager via REST API.

    Endpoints:
        GET    /api/v1/agents/datasets/{agent_id}              - List datasets
        GET    /api/v1/agents/datasets/{agent_id}/{dataset_id} - Describe a single dataset
        PATCH  /api/v1/agents/datasets/{agent_id}              - Activate/deactivate dataset
        PUT    /api/v1/agents/datasets/{agent_id}              - Upload file as dataset
        POST   /api/v1/agents/datasets/{agent_id}              - Add query as dataset
        DELETE /api/v1/agents/datasets/{agent_id}              - Delete dataset
    """

    _user_objects_handler: UserObjectsHandler = None

    @property
    def user_objects_handler(self) -> UserObjectsHandler:
        """Lazy-initialized UserObjectsHandler instance."""
        if self._user_objects_handler is None:
            self._user_objects_handler = UserObjectsHandler(logger=self.logger)
        return self._user_objects_handler

    async def _get_dataset_manager(self, agent_id: str) -> DatasetManager:
        """Get or create user's DatasetManager from session.

        If the session DM is absent or empty, the agent's catalog is cloned
        into a new user-scoped DatasetManager.

        Raises:
            KeyError: If the agent is not found in BotManager and there is no
                existing session DatasetManager for this agent_id.
        """
        session_key = self.user_objects_handler.get_session_key(
            agent_id, "dataset_manager"
        )
        return await _resolve_dataset_manager(
            self.request, agent_id, session_key, self.logger
        )

    async def get(self) -> web.Response:
        """List all datasets or describe a single dataset.

        Endpoint:
            GET /api/v1/agents/datasets/{agent_id}
            GET /api/v1/agents/datasets/{agent_id}/{dataset_id}
            
        Query params for detail:
            eda:           bool - Include EDA summary (default: true)
            samples:       bool - Include sample rows (default: false)
            column_stats:  bool - Include per-column statistics (default: false)
            metrics_guide: bool - Include metrics guide (default: false)
            column:        str  - Describe a single column only (optional)

        Returns:
            DatasetListResponse with dataset information, or single Dataset metadata.
        """
        agent_id = self.request.match_info.get("agent_id")
        dataset_id = self.request.match_info.get("dataset_id")
        if not agent_id:
            return self.json_response({"error": "agent_id is required"}, status=400)

        try:
            dm = await self._get_dataset_manager(agent_id)
        except KeyError:
            return self.json_response(
                {"error": f"Agent '{agent_id}' not found"}, status=404
            )

        if dataset_id:
            q = self.request.query
            include_eda = q.get("eda", "true").lower() == "true"
            include_samples = q.get("samples", "").lower() == "true"
            include_column_stats = q.get("column_stats", "").lower() == "true"
            include_metrics_guide = q.get("metrics_guide", "").lower() == "true"
            column = q.get("column") or None

            try:
                metadata = await dm.get_metadata(
                    name=dataset_id,
                    include_eda=include_eda,
                    include_samples=include_samples,
                    include_column_stats=include_column_stats,
                    include_metrics_guide=include_metrics_guide,
                    column=column,
                )
                return self.json_response(metadata)
            except KeyError:
                return self.json_response(
                    {"error": f"Dataset '{dataset_id}' not found"}, status=404
                )
            except Exception as e:
                self.logger.error("Error describing dataset '%s': %s", dataset_id, e)
                return self.json_response({"error": str(e)}, status=500)
        else:
            try:
                datasets_info = await dm.list_available()

                datasets = []
                for info in datasets_info:
                    dataset_dict = {
                        "name": info.get("name", ""),
                        "type": info.get("source_type", "dataframe"),
                        "description": info.get("description", ""),
                        "shape": info.get("shape", (0, 0)),
                        "is_active": info.get("is_active", True),
                        "loaded": info.get("loaded", False),
                    }
                    datasets.append(dataset_dict)

                response = DatasetListResponse(
                    datasets=datasets,
                    total=len(datasets),
                    active_count=sum(1 for d in datasets if d.get("is_active", True)),
                )
                return self.json_response(response.model_dump())
            except Exception as e:
                self.logger.error("Error listing datasets: %s", e)
                return self.json_response({"error": str(e)}, status=500)

    async def patch(self) -> web.Response:
        """Activate or deactivate a dataset.

        Body: DatasetPatchRequest
        """
        agent_id = self.request.match_info.get("agent_id")
        if not agent_id:
            return self.json_response({"error": "agent_id is required"}, status=400)

        try:
            data = await self.request.json()
            request_body = DatasetPatchRequest(**data)
        except ValidationError as e:
            return self.json_response(
                {"error": "Invalid request", "detail": str(e)}, status=400
            )
        except Exception as e:
            return self.json_response({"error": f"Invalid JSON: {e}"}, status=400)

        try:
            dm = await self._get_dataset_manager(agent_id)

            if request_body.action == DatasetAction.ACTIVATE:
                result = dm.activate(request_body.dataset_name)
                if not result:
                    return self.json_response(
                        {"error": f"Dataset '{request_body.dataset_name}' not found"},
                        status=404,
                    )
                message = f"Dataset '{request_body.dataset_name}' activated"
            else:
                result = dm.deactivate(request_body.dataset_name)
                if not result:
                    return self.json_response(
                        {"error": f"Dataset '{request_body.dataset_name}' not found"},
                        status=404,
                    )
                message = f"Dataset '{request_body.dataset_name}' deactivated"

            return self.json_response({
                "name": request_body.dataset_name,
                "action": request_body.action,
                "message": message,
            })
        except KeyError:
            return self.json_response(
                {"error": f"Agent '{agent_id}' not found"}, status=404
            )
        except Exception as e:
            self.logger.error("Error patching dataset: %s", e)
            return self.json_response({"error": str(e)}, status=500)

    async def put(self) -> web.Response:
        """Upload an Excel/CSV file as a new dataset.

        Accepts multipart/form-data with:
            file: The file to upload
            name: Optional dataset name (defaults to filename without extension)
        """
        agent_id = self.request.match_info.get("agent_id")
        if not agent_id:
            return self.json_response({"error": "agent_id is required"}, status=400)

        try:
            reader = await self.request.multipart()
            dataset_name = None
            df = None
            filename = None

            async for field in reader:
                if field.name == "name":
                    dataset_name = (await field.read()).decode("utf-8")
                elif field.name == "file":
                    filename = field.filename
                    data = await field.read()

                    if len(data) > MAX_FILE_SIZE:
                        return self.json_response(
                            {
                                "error": (
                                    f"File too large. "
                                    f"Maximum size is {MAX_FILE_SIZE // (1024 * 1024)}MB"
                                )
                            },
                            status=400,
                        )

                    if filename.endswith((".xlsx", ".xls")):
                        df = pd.read_excel(BytesIO(data))
                    elif filename.endswith(".csv"):
                        df = pd.read_csv(BytesIO(data))
                    else:
                        return self.json_response(
                            {"error": "Unsupported file format. Use .xlsx, .xls, or .csv"},
                            status=400,
                        )

            if df is None:
                return self.json_response({"error": "No file provided"}, status=400)

            if not dataset_name:
                dataset_name = (
                    filename.rsplit(".", 1)[0] if filename else "uploaded_dataset"
                )

            dm = await self._get_dataset_manager(agent_id)
            dm.add_dataframe(dataset_name, df)

            response = DatasetUploadResponse(
                name=dataset_name,
                rows=len(df),
                columns=len(df.columns),
                columns_list=list(df.columns),
            )
            return self.json_response(response.model_dump(), status=201)

        except KeyError:
            return self.json_response(
                {"error": f"Agent '{agent_id}' not found"}, status=404
            )
        except Exception as e:
            self.logger.error("Error uploading dataset: %s", e)
            return self.json_response({"error": str(e)}, status=500)

    async def post(self) -> web.Response:
        """Add a dataset from query/sql/datasource configuration.

        Body: DatasetQueryRequest

        """
        agent_id = self.request.match_info.get("agent_id")
        if not agent_id:
            return self.json_response({"error": "agent_id is required"}, status=400)

        try:
            data = await self.request.json()
            request_body = DatasetQueryRequest(**data)
            request_body.validate_query_source()
        except ValidationError as e:
            return self.json_response(
                {"error": "Invalid request", "detail": str(e)}, status=400
            )
        except ValueError as e:
            return self.json_response({"error": str(e)}, status=400)
        except Exception as e:
            return self.json_response({"error": f"Invalid JSON: {e}"}, status=400)

        try:
            dm = await self._get_dataset_manager(agent_id)

            if request_body.datasource:
                ds = request_body.datasource
                ds_type = str(ds.get("type", "")).lower().strip()
                metadata = {"description": request_body.description or ""}
                if ds.get("metadata") and isinstance(ds.get("metadata"), dict):
                    metadata.update(ds.get("metadata"))

                raw_cache_ttl = ds.get("cache_ttl", 3600)
                try:
                    cache_ttl = int(raw_cache_ttl)
                except (TypeError, ValueError):
                    return self.json_response(
                        {"error": "datasource.cache_ttl must be an integer (seconds)"},
                        status=400,
                    )
                if cache_ttl < 0:
                    return self.json_response(
                        {"error": "datasource.cache_ttl must be a non-negative integer (seconds)"},
                        status=400,
                    )

                if ds_type == "query_slug":
                    slug = ds.get("slug") or ds.get("query_slug")
                    if not slug:
                        return self.json_response({"error": "datasource.slug is required for query_slug"}, status=400)
                    dm.add_query(name=request_body.name, query_slug=slug, metadata=metadata)
                elif ds_type == "sql":
                    sql = ds.get("sql")
                    driver = ds.get("driver")
                    if not sql or not driver:
                        return self.json_response({"error": "datasource.sql and datasource.driver are required for sql"}, status=400)
                    dm.add_sql_source(
                        name=request_body.name,
                        sql=sql,
                        driver=driver,
                        dsn=ds.get("dsn"),
                        metadata=metadata,
                        cache_ttl=cache_ttl,
                    )
                elif ds_type == "table":
                    table = ds.get("table")
                    driver = ds.get("driver")
                    if not table or not driver:
                        return self.json_response({"error": "datasource.table and datasource.driver are required for table"}, status=400)
                    await dm.add_table_source(
                        name=request_body.name,
                        table=table,
                        driver=driver,
                        dsn=ds.get("dsn"),
                        credentials=ds.get("credentials"),
                        metadata=metadata,
                        cache_ttl=cache_ttl,
                        strict_schema=bool(ds.get("strict_schema", True)),
                    )
                elif ds_type == "airtable":
                    base_id = ds.get("base_id")
                    table = ds.get("table")
                    api_key = ds.get("api_key")
                    if not base_id or not table:
                        return self.json_response({"error": "datasource.base_id and datasource.table are required for airtable"}, status=400)
                    await dm.add_airtable_source(
                        name=request_body.name,
                        base_id=base_id,
                        table=table,
                        api_key=api_key,
                        view=ds.get("view"),
                        metadata=metadata,
                        cache_ttl=cache_ttl,
                        fetch_on_create=True,
                    )
                elif ds_type == "smartsheet":
                    sheet_id = ds.get("sheet_id")
                    access_token = ds.get("access_token")
                    if not sheet_id:
                        return self.json_response({"error": "datasource.sheet_id is required for smartsheet"}, status=400)
                    await dm.add_smartsheet_source(
                        name=request_body.name,
                        sheet_id=str(sheet_id),
                        access_token=access_token,
                        metadata=metadata,
                        cache_ttl=cache_ttl,
                        fetch_on_create=True,
                    )
                elif ds_type == "dataframe":
                    data = ds.get("data")
                    if data is None:
                        return self.json_response({"error": "datasource.data is required for dataframe"}, status=400)
                    df = pd.DataFrame(data)
                    dm.add_dataframe(request_body.name, df, metadata=metadata)
                else:
                    return self.json_response({"error": f"Unsupported datasource type '{ds_type}'"}, status=400)

                if ds.get("refresh"):
                    # Only support automatic refresh for datasource types that can
                    # be materialized without additional parameters. For other
                    # types, return a controlled 400 instead of risking a 500.
                    if ds_type in ("airtable", "smartsheet", "query_slug", "dataframe"):
                        await dm.materialize(request_body.name, force_refresh=True)
                    else:
                        return self.json_response(
                            {
                                "error": f"Refresh is not supported for datasource type '{ds_type}' via this endpoint"
                            },
                            status=400,
                        )

                return self.json_response(
                    {
                        "name": request_body.name,
                        "type": ds_type,
                        "message": f"Datasource dataset '{request_body.name}' added successfully",
                    },
                    status=201,
                )

            if request_body.query_slug:
                dm.add_query(
                    name=request_body.name,
                    query_slug=request_body.query_slug,
                    metadata={"description": request_body.description or ""},
                )
                query_type = "query_slug"
            elif request_body.query:
                dm.add_query(
                    name=request_body.name,
                    query_slug=request_body.name,
                    metadata={
                        "description": request_body.description or "",
                        "raw_sql": request_body.query,
                    },
                )
                query_type = "query"
            else:
                return self.json_response(
                    {"error": "Either 'query', 'query_slug', or 'datasource' must be provided"},
                    status=400,
                )

            return self.json_response(
                {
                    "name": request_body.name,
                    "type": query_type,
                    "message": f"Query dataset '{request_body.name}' added successfully",
                },
                status=201,
            )
        except KeyError:
            return self.json_response(
                {"error": f"Agent '{agent_id}' not found"}, status=404
            )
        except Exception as e:
            self.logger.error("Error adding query dataset: %s", e)
            return self.json_response({"error": str(e)}, status=500)

    async def delete(self) -> web.Response:
        """Delete a dataset from the DatasetManager.

        Query params:
            name: str - Dataset name to delete
        """
        agent_id = self.request.match_info.get("agent_id")
        if not agent_id:
            return self.json_response({"error": "agent_id is required"}, status=400)

        dataset_name = self.request.query.get("name")
        if not dataset_name:
            return self.json_response(
                {"error": "Query parameter 'name' is required"}, status=400
            )

        try:
            dm = await self._get_dataset_manager(agent_id)
            dm.remove(dataset_name)

            response = DatasetDeleteResponse(name=dataset_name)
            return self.json_response(response.model_dump())

        except KeyError:
            return self.json_response(
                {"error": f"Agent '{agent_id}' not found"}, status=404
            )
        except ValueError:
            # remove() raises ValueError when dataset not found
            return self.json_response(
                {"error": f"Dataset '{dataset_name}' not found"}, status=404
            )
        except Exception as e:
            self.logger.error("Error deleting dataset: %s", e)
            return self.json_response({"error": str(e)}, status=500)


