"""Save query results as CSV files and optionally as PG tables."""

from __future__ import annotations

import csv
import json
from pathlib import Path
from typing import Any

import pandas as pd

from config import AppSettings
from database.connection_manager import ConnectionManager
from database.local_pg import LocalPG


class ResultStorage:
    """Store query results for later retrieval."""

    def __init__(self, settings: AppSettings, cm: ConnectionManager, local_pg: LocalPG) -> None:
        self._csv_path = Path(settings.results_csv_path)
        self._csv_path.mkdir(parents=True, exist_ok=True)
        self._cm = cm
        self._local = local_pg

    async def store_results(
        self,
        message_id: str,
        rows: list[dict],
        columns: list[str],
    ) -> dict:
        """Store query results as CSV and return storage info.

        Returns: {"csv_path": str, "result_table": str | None, "row_count": int}
        """
        # Save as CSV
        csv_filename = f"{message_id}.csv"
        csv_path = self._csv_path / csv_filename

        df = pd.DataFrame(rows, columns=columns)
        df.to_csv(csv_path, index=False, quoting=csv.QUOTE_NONNUMERIC)

        # For larger result sets, also store in PG for fast querying
        result_table = None
        if len(rows) > 0:
            result_table = f"result_{message_id.replace('-', '_')}"
            try:
                pg_columns = [{"name": c, "type": "TEXT"} for c in columns]
                await self._local.create_table_from_columns(
                    result_table, pg_columns, schema="v3q_results"
                )
                await self._local.insert_dataframe(
                    result_table, df, schema="v3q_results"
                )
            except Exception:
                result_table = None

        return {
            "csv_path": str(csv_path),
            "result_table": result_table,
            "row_count": len(rows),
        }

    async def get_results(self, message_id: str) -> dict | None:
        """Retrieve stored results for a message."""
        csv_path = self._csv_path / f"{message_id}.csv"
        if not csv_path.exists():
            return None

        import numpy as np
        df = pd.read_csv(csv_path)
        # Replace NaN/NaT with None so JSON serialization works (nan is not valid JSON)
        df = df.replace({np.nan: None})
        return {
            "columns": list(df.columns),
            "rows": df.to_dict("records"),
            "row_count": len(df),
        }

    def get_csv_path(self, message_id: str) -> Path | None:
        csv_path = self._csv_path / f"{message_id}.csv"
        return csv_path if csv_path.exists() else None
