"""
NexMind AI Python SDK — clean ML training data with AI agents.
"""
import io
import time
import warnings
from typing import Any, Dict, List, Optional, Tuple

import httpx
import pandas as pd


class CleaningJob:
    """Represents a cleaning job submitted to NexMind AI."""

    def __init__(self, job_data: Dict, client: "NexMindClient"):
        self._data = job_data
        self._client = client

    @property
    def id(self) -> str:
        return self._data["id"]

    @property
    def status(self) -> str:
        return self._data["status"]

    @property
    def report(self) -> Optional[Dict]:
        return self._data.get("_report")

    @property
    def quality_before(self) -> Optional[float]:
        return self._data.get("quality_score_before")

    @property
    def quality_after(self) -> Optional[float]:
        return self._data.get("quality_score_after")

    def wait(self, poll_interval: float = 2.0, timeout: float = 3600.0) -> "CleaningJob":
        """Block until job completes or timeout."""
        start = time.time()
        while True:
            self.refresh()
            if self.status in ("completed", "failed", "cancelled"):
                if self.status == "completed":
                    try:
                        self._data["_report"] = self._client.get_report(self.id)
                    except Exception:
                        pass
                return self
            if time.time() - start > timeout:
                raise TimeoutError(f"Job {self.id} did not complete within {timeout}s")
            time.sleep(poll_interval)

    def refresh(self) -> "CleaningJob":
        self._data = self._client._get(f"/jobs/{self.id}")
        return self

    def __repr__(self) -> str:
        return f"<CleaningJob id={self.id} status={self.status}>"


class NexMindClient:
    """
    NexMind AI Python SDK — clean ML training data with AI-powered agents.

    Usage:
        client = NexMindClient(api_key="dcs_your_key")

        # Quick clean
        clean_df, report = client.clean(my_dataframe)

        # Or with options
        clean_df, report = client.clean(
            my_dataframe,
            agents=["schema", "anomaly", "dedup", "repair"],
        )
    """

    DEFAULT_BASE_URL = "https://api.getnexmind.com/api/v1"

    def __init__(
        self,
        api_key: str,
        base_url: str = DEFAULT_BASE_URL,
        timeout: float = 30.0,
    ):
        """
        Initialize the NexMind client.

        Args:
            api_key: Your API key (starts with 'dcs_'). Get one at app.getnexmind.com/api-keys
            base_url: API base URL (default: https://api.getnexmind.com/api/v1)
            timeout: Request timeout in seconds
        """
        if not api_key:
            raise ValueError("api_key is required. Get one at https://app.getnexmind.com/api-keys")
        self._base_url = base_url.rstrip("/")
        self._headers = {
            "X-API-Key": api_key,
            "Content-Type": "application/json",
        }
        self._timeout = timeout

    def clean(
        self,
        df: pd.DataFrame,
        agents: Optional[List[str]] = None,
        agent_config: Optional[Dict] = None,
        wait: bool = True,
        poll_interval: float = 2.0,
    ) -> Tuple[pd.DataFrame, Dict]:
        """
        Clean a DataFrame using AI agents.

        Args:
            df: The dirty DataFrame to clean
            agents: List of agents to run (default: all 5)
            agent_config: Per-agent configuration overrides
            wait: Wait for job to complete (default: True)
            poll_interval: Seconds between status checks

        Returns:
            Tuple of (cleaned_DataFrame, report_dict)

        Example:
            clean_df, report = client.clean(my_df)
            print(f"Quality: {report['quality_score_before']} → {report['quality_score_after']}")
        """
        # Upload CSV
        buf = io.BytesIO()
        df.to_csv(buf, index=False)
        buf.seek(0)

        with httpx.Client(timeout=60) as client:
            resp = client.post(
                f"{self._base_url}/datasets/upload",
                headers={k: v for k, v in self._headers.items() if k != "Content-Type"},
                files={"file": ("upload.csv", buf, "text/csv")},
            )
            resp.raise_for_status()
            upload = resp.json()

        dataset_id = upload["dataset_id"]
        source = {
            "type": "minio",
            "bucket": "raw-data",
            "key": f"uploads/{dataset_id}/upload.csv",
        }
        output = {
            "type": "minio",
            "bucket": "clean-data",
            "key": f"cleaned/{dataset_id}/output.parquet",
            "format": "parquet",
        }

        job = self.submit_job(
            name=f"SDK Job - {dataset_id[:8]}",
            source=source,
            output=output,
            agents=agents,
            agent_config=agent_config,
        )

        if wait:
            job.wait(poll_interval=poll_interval)
            if job.status == "failed":
                raise RuntimeError(f"Cleaning job failed: {job._data.get('error_message')}")
            # Try downloading cleaned data
            try:
                clean_df = self._download_output(job)
                return clean_df, job.report or {}
            except Exception:
                warnings.warn("Could not download cleaned data. Check output_path in the report.", stacklevel=2)
                return df, job.report or {}

        return df, {}

    def submit_job(
        self,
        name: str,
        source: Dict,
        output: Dict,
        agents: Optional[List[str]] = None,
        agent_config: Optional[Dict] = None,
    ) -> CleaningJob:
        """Submit an async cleaning job and return immediately."""
        payload = {
            "name": name,
            "source": source,
            "output": output,
            "agents": agents or ["schema", "anomaly", "dedup", "repair", "labeling"],
            "agent_config": agent_config or {},
        }
        data = self._post("/jobs", payload)
        return CleaningJob(data, self)

    def list_jobs(self, status: Optional[str] = None, limit: int = 20) -> List[Dict]:
        """List cleaning jobs."""
        params = {"limit": limit}
        if status:
            params["status"] = status
        return self._get("/jobs", params=params)

    def get_report(self, job_id: str) -> Dict:
        """Get the cleaning report for a completed job."""
        return self._get(f"/jobs/{job_id}/report")

    def _download_output(self, job: CleaningJob) -> pd.DataFrame:
        output_path = job._data.get("output_path")
        if not output_path:
            raise RuntimeError("No output_path in job data")
        with httpx.Client(timeout=120) as client:
            resp = client.get(
                f"{self._base_url}/jobs/{job.id}/download",
                headers=self._headers,
            )
            resp.raise_for_status()
            if output_path.endswith(".parquet"):
                return pd.read_parquet(io.BytesIO(resp.content))
            return pd.read_csv(io.BytesIO(resp.content))

    def _get(self, path: str, params: Optional[Dict] = None) -> Any:
        with httpx.Client(timeout=self._timeout) as client:
            resp = client.get(
                f"{self._base_url}{path}",
                headers=self._headers,
                params=params,
            )
            resp.raise_for_status()
            return resp.json()

    def _post(self, path: str, body: Dict) -> Any:
        with httpx.Client(timeout=self._timeout) as client:
            resp = client.post(
                f"{self._base_url}{path}",
                headers=self._headers,
                json=body,
            )
            resp.raise_for_status()
            return resp.json()
