from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, Optional, Union
from urllib.parse import urlparse

import pandas as pd

from ..context import Context


class DataModule:
    def __init__(self, context: Context) -> None:
        self._ctx = context

    # ---------------------------------------------------------
    # Load Data
    # ---------------------------------------------------------
    def load(self, source: Union[str, Path, pd.DataFrame]) -> pd.DataFrame:
        # 1️⃣ DataFrame input
        if isinstance(source, pd.DataFrame):
            df = source.copy()

        # 2️⃣ String or Path input
        elif isinstance(source, (str, Path)):
            source_str = str(source)

            # Detect URL
            parsed = urlparse(source_str)
            is_url = parsed.scheme in ("http", "https")

            if is_url:
                suffix = Path(parsed.path).suffix.lower()

                if suffix == ".csv":
                    df = pd.read_csv(source_str)

                elif suffix == ".json":
                    df = pd.read_json(source_str)

                elif suffix in (".xlsx", ".xls"):
                    df = pd.read_excel(source_str)

                else:
                    raise ValueError(f"Unsupported remote format: {suffix}")

            else:
                path = Path(source_str)

                if not path.exists():
                    raise FileNotFoundError(str(path))

                suffix = path.suffix.lower()

                if suffix == ".csv":
                    df = pd.read_csv(path)

                elif suffix in (".xlsx", ".xls"):
                    df = pd.read_excel(path)

                elif suffix == ".json":
                    df = pd.read_json(path)

                else:
                    raise ValueError(f"Unsupported format: {suffix}")

        else:
            raise TypeError("Source must be DataFrame, file path, or URL.")

        # 3️⃣ Validate non-empty
        if df.empty:
            raise ValueError("Dataset is empty.")

        self._ctx.set_dataframe(df)
        return df

    # ---------------------------------------------------------
    # From Dictionary
    # ---------------------------------------------------------
    def from_dict(self, data: Dict[str, Any]) -> pd.DataFrame:
        return self.load(pd.DataFrame(data))

    # ---------------------------------------------------------
    # Target Handling
    # ---------------------------------------------------------
    def set_target(self, target: Optional[str]) -> None:
        if target:
            self._ctx.set_target(target)

    # ---------------------------------------------------------
    # Preview
    # ---------------------------------------------------------
    def preview(self, rows: int = 5) -> pd.DataFrame:
        return self._ctx.ensure_dataframe().head(rows)