import base64
import io
import json
import csv as csv_mod

import pandas as pd
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import Optional

router = APIRouter(prefix="/dataset")


# ── Request models ─────────────────────────────────────────────────────────────

class UploadRequest(BaseModel):
    file_name: str
    file_size: int
    file_data: str          # text for JSONL/CSV, base64 for Parquet
    is_base64: bool = False

class HFLoadRequest(BaseModel):
    hf_id: str
    hf_config: str = "default"
    hf_split: str = "train"
    sample_size: int = 1000

    @property
    def capped_sample_size(self) -> int:
        return min(self.sample_size, 1000)


# ── Shared: dataframe → summary ───────────────────────────────────────────────

def _df_summary(df: pd.DataFrame, source_name: str) -> dict:
    """Turn any DataFrame into the standard DatasetSource metadata response."""
    row_count = len(df)
    columns = df.columns.tolist()

    # Sample: up to 20 rows as list of dicts, NaN → None
    sample = df.head(300).where(pd.notnull(df), None).to_dict(orient="records")

    # Basic column type info
    col_types = {c: str(df[c].dtype) for c in columns}

    # Null counts per column
    null_counts = df.isnull().sum().to_dict()

    return {
        "name": source_name,
        "row_count": row_count,
        "columns": columns,
        "col_types": col_types,
        "null_counts": null_counts,
        "sample": sample,
    }


# ── /dataset/upload ────────────────────────────────────────────────────────────

@router.post("/upload")
async def dataset_upload(req: UploadRequest):
    name = req.file_name
    ext = name.rsplit(".", 1)[-1].lower()

    try:
        if req.is_base64 or ext == "parquet":
            raw = base64.b64decode(req.file_data)
            df = pd.read_parquet(io.BytesIO(raw))

        elif ext in ("jsonl", "json"):
            lines = req.file_data.strip().splitlines()
            # Try JSONL first (one JSON object per line)
            try:
                records = [json.loads(l) for l in lines if l.strip()]
                df = pd.DataFrame(records)
            except json.JSONDecodeError:
                # Fall back to a single JSON array
                records = json.loads(req.file_data)
                df = pd.DataFrame(records if isinstance(records, list) else [records])

        elif ext == "csv":
            df = pd.read_csv(
                io.StringIO(req.file_data),
                encoding_errors="replace",
                on_bad_lines="skip",
            )

        elif ext == "tsv":
            df = pd.read_csv(
                io.StringIO(req.file_data),
                sep="\t",
                encoding_errors="replace",
                on_bad_lines="skip",
            )

        else:
            raise HTTPException(status_code=400, detail=f"Unsupported file type: .{ext}")

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=422, detail=f"Failed to parse file: {e}")

    return _df_summary(df, name)


# ── /dataset/load-hf ──────────────────────────────────────────────────────────

@router.post("/load-hf")
async def dataset_load_hf(req: HFLoadRequest):
    try:
        from datasets import load_dataset, get_dataset_config_names
    except ImportError:
        raise HTTPException(status_code=503, detail="datasets library not installed. Run: pip install datasets")

    try:
        config = req.hf_config if req.hf_config != "default" else None

        # Try to list configs — if the dataset has no configs, pass None
        try:
            configs = get_dataset_config_names(req.hf_id)
            if config not in configs:
                config = configs[0] if configs else None
        except Exception:
            config = None

        load_kwargs: dict = dict(
            path=req.hf_id,
            split=req.hf_split,
            trust_remote_code=False,
        )
        if config:
            load_kwargs["name"] = config

        load_kwargs["streaming"] = True
        raw = load_dataset(**load_kwargs)
        records = list(raw.take(req.capped_sample_size))

        df = pd.DataFrame(records)

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=422, detail=f"Failed to load HuggingFace dataset: {e}")

    summary = _df_summary(df, req.hf_id)
    summary["hf_id"] = req.hf_id
    summary["hf_split"] = req.hf_split
    summary["sample_size"] = req.sample_size
    return summary