services/file_manager.py

"""
File manager — upload, CSV→Parquet conversion, download, cleanup.
"""
from __future__ import annotations

import os
import uuid
import shutil
import polars as pl

from config import DATA_DIR, UPLOAD_DIR, LAZY_THRESHOLD_BYTES, MAX_UPLOAD_BYTES
from services.session_manager import session_manager
from core.column_registry import column_registry


def _human_size(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} TB"


def _detect_separator(path: str) -> str:
    """Sniff the separator from the first few KB."""
    with open(path, "r", errors="replace") as f:
        head = f.read(8192)
    for sep in (",", ";", "\t", "|"):
        if sep in head:
            return sep
    return ","


def handle_upload(file_content: bytes, file_name: str) -> dict:
    """Save uploaded file, convert to Parquet, create session."""
    # Size check
    if len(file_content) > MAX_UPLOAD_BYTES:
        raise ValueError(f"File too large. Max {MAX_UPLOAD_BYTES // (1024**2)} MB allowed.")

    # Generate session
    session_id = uuid.uuid4().hex[:12]

    # Save raw upload
    ext = os.path.splitext(file_name)[1].lower()
    raw_path = os.path.join(UPLOAD_DIR, f"{session_id}{ext}")
    with open(raw_path, "wb") as f:
        f.write(file_content)

    # Determine file type and read
    if ext in (".csv", ".tsv", ".txt"):
        sep = _detect_separator(raw_path)
        lf = pl.scan_csv(raw_path, separator=sep, try_parse_dates=True)
    elif ext in (".xlsx", ".xls"):
        # Polars can read Excel but needs the feature flag.
        # Fall back to eager read for Excel.
        df = pl.read_excel(raw_path)  # type: ignore[attr-defined]
        parquet_path = os.path.join(DATA_DIR, f"{session_id}.parquet")
        df.write_parquet(parquet_path)
        os.remove(raw_path)
    elif ext == ".parquet":
        shutil.copy2(raw_path, os.path.join(DATA_DIR, f"{session_id}.parquet"))
        os.remove(raw_path)
        lf = pl.scan_parquet(os.path.join(DATA_DIR, f"{session_id}.parquet"))
    else:
        os.remove(raw_path)
        raise ValueError(f"Unsupported file format: {ext}")

    # For CSV: stream-convert to Parquet
    if ext in (".csv", ".tsv", ".txt"):
        parquet_path = os.path.join(DATA_DIR, f"{session_id}.parquet")
        lf.sink_parquet(parquet_path)
        os.remove(raw_path)
        lf = pl.scan_parquet(parquet_path)

    # Read metadata (no full data load)
    schema = lf.collect_schema()
    row_count = lf.select(pl.len()).collect().item()
    columns = [{"name": name, "dtype": str(dtype)} for name, dtype in schema.items()]

    # Register in session manager
    session_manager.create(
        session_id=session_id,
        file_name=file_name,
        file_size_bytes=len(file_content),
        columns=columns,
        row_count=row_count,
    )

    # Register columns for fuzzy resolution
    column_registry.register(session_id, schema.names())

    return {
        "session_id": session_id,
        "file_name": file_name,
        "rows": row_count,
        "columns": schema.names(),
        "size_human": _human_size(len(file_content)),
    }


def get_download_path(session_id: str) -> Optional[str]:
    """Return path to the Parquet file, or None if session doesn't exist."""
    path = session_manager.get_filepath(session_id)
    if os.path.exists(path) and session_manager.get(session_id):
        return path
    return None


def export_to_csv(session_id: str) -> Optional[str]:
    """Export Parquet to a temporary CSV file and return its path."""
    pq_path = session_manager.get_filepath(session_id)
    if not os.path.exists(pq_path):
        return None

    csv_path = os.path.join(UPLOAD_DIR, f"{session_id}.csv")
    lf = pl.scan_parquet(pq_path)
    if os.path.getsize(pq_path) > LAZY_THRESHOLD_BYTES:
        lf.sink_csv(csv_path)
    else:
        lf.collect().write_csv(csv_path)
    return csv_path


def delete_session(session_id: str) -> bool:
    """Remove session metadata, Parquet file, and column registry."""
    import glob
    pattern = os.path.join(DATA_DIR, f"{session_id}*.parquet")
    removed = False
    for pq_path in glob.glob(pattern):
        try:
            os.remove(pq_path)
            removed = True
        except Exception:
            pass
    session_manager.remove(session_id)
    column_registry.remove(session_id)
    return removed