"""AIR client - main entry point for the SDK."""

from __future__ import annotations

import asyncio
import json
import os
import time
import uuid
from typing import Any, Callable

import httpx

from .exceptions import AIRError, AuthError
from .project import Project
from .ws_session import OneShotResult, WebSocketSession

_DEFAULT_BASE_URL = "http://localhost:8000"
_DEFAULT_TIMEOUT = 120.0  # seconds


class AIR:
    """
    Client for the AIR Backend REST API.

    Usage::

        import air

        client = air.AIR(api_key="air_k1_...")
        keywords = client.keywords("dark matter and lensing", n=5)

        # Standalone paper review with Skepthical
        result = client.review("/path/to/paper.pdf", thoroughness="High")

        # Project pipeline
        project = client.create_project("my-research", data_description="We study...")
        idea = project.idea()
        project.literature()
        project.methods()
        project.paper(journal="AAS")
        review = project.review(review_engine="skepthical")

    Args:
        api_key: AIR API key (``air_k1_...``). Falls back to ``AIR_API_KEY`` env var.
        base_url: Base URL of the AIR backend. Defaults to ``http://localhost:8000``.
        timeout: Default HTTP request timeout in seconds.
    """

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: float = _DEFAULT_TIMEOUT,
        local_dir: str | None = None,
    ):
        self._api_key = api_key or os.environ.get("AIR_API_KEY", "")
        if not self._api_key:
            raise AuthError(
                "API key required. Pass api_key= or set the AIR_API_KEY environment variable."
            )

        self._base_url = (base_url or os.environ.get("AIR_BASE_URL", _DEFAULT_BASE_URL)).rstrip("/")
        self._local_dir = local_dir or os.environ.get("AIR_LOCAL_DIR", "~/ai-scientist")

        self._client = httpx.Client(
            base_url=self._base_url,
            headers={"Authorization": f"Bearer {self._api_key}"},
            timeout=timeout,
        )

    def close(self) -> None:
        """Close the underlying HTTP client."""
        self._client.close()

    def __enter__(self) -> "AIR":
        return self

    def __exit__(self, *exc: Any) -> None:
        self.close()

    # ------------------------------------------------------------------
    # Health
    # ------------------------------------------------------------------

    def health(self) -> dict:
        """Check if the backend is healthy."""
        resp = self._client.get("/api/health")
        resp.raise_for_status()
        return resp.json()

    # ------------------------------------------------------------------
    # Standalone tools
    # ------------------------------------------------------------------

    def keywords(
        self,
        text: str,
        n: int = 5,
        kw_type: str = "unesco",
    ) -> Any:
        """
        Extract keywords from text.

        Args:
            text: Input text.
            n: Number of keywords to extract.
            kw_type: Keyword type ("unesco", "aas", etc.).

        Returns:
            Keywords result (list or dict depending on kw_type).
        """
        resp = self._client.post(
            "/api/v1/keywords",
            json={"text": text, "n_keywords": n, "kw_type": kw_type},
        )
        self._check(resp)
        return resp.json().get("keywords")

    def arxiv(self, text: str) -> dict:
        """
        Extract arXiv URLs from text and download papers.

        Args:
            text: Input text containing arXiv URLs.

        Returns:
            Dict with download results.
        """
        resp = self._client.post("/api/v1/arxiv", json={"text": text})
        self._check(resp)
        return resp.json().get("result", {})

    def enhance(
        self,
        text: str,
        max_workers: int = 2,
        max_depth: int = 10,
    ) -> str:
        """
        Enhance input text with contextual information from arXiv papers.

        Args:
            text: Input text (may contain arXiv URLs).
            max_workers: Parallel workers for downloading.
            max_depth: Max summarisation depth.

        Returns:
            Enhanced text.
        """
        resp = self._client.post(
            "/api/v1/enhance",
            json={"text": text, "max_workers": max_workers, "max_depth": max_depth},
        )
        self._check(resp)
        return resp.json().get("enhanced_text", "")

    def ocr(self, file_path: str) -> dict:
        """
        Process a PDF with OCR (server-side path).

        Args:
            file_path: Path to the PDF on the server.

        Returns:
            OCR result dict.
        """
        resp = self._client.post("/api/v1/ocr", json={"file_path": file_path})
        self._check(resp)
        return resp.json().get("result", {})

    def review(
        self,
        file_path: str,
        *,
        output_dir: str | None = None,
        thoroughness: str = "Standard",
        figures_review: bool = False,
        verify_statements: bool = False,
        review_maths: bool = False,
        review_numerics: bool = False,
        emails: list[str] | None = None,
        timeout: float = 3600,
        poll_interval: float = 5.0,
    ) -> dict:
        """
        Review a paper with Skepthical (standalone, not tied to a project).

        Args:
            file_path: Path to the paper PDF on the server.
            output_dir: Local directory where the review files (MD and PDF)
                will be saved. Defaults to ``~/ai-scientist/reviews/``.
            thoroughness: "Standard" (single reviewer) or "High" (two reviewers merged).
            figures_review: Enable detailed figure/diagram analysis.
            verify_statements: Enable key statement verification.
            review_maths: Enable mathematical derivation audit.
            review_numerics: Enable numerical computation audit.
            emails: Optional email recipients for the review report.
            timeout: Max seconds to wait for completion.
            poll_interval: Seconds between status polls.

        Returns:
            Dict with keys:

            - ``review`` (str): The markdown review text.
            - ``review_md_path`` (str): Local path to the saved markdown file.
            - ``review_pdf_path`` (str): Local path to the saved PDF file
              (empty string if PDF generation failed on the server).
        """
        body: dict[str, Any] = {
            "file_path": file_path,
            "thoroughness": thoroughness,
            "figures_review": figures_review,
            "verify_statements": verify_statements,
            "review_maths": review_maths,
            "review_numerics": review_numerics,
        }
        if emails is not None:
            body["emails"] = emails

        resp = self._client.post("/api/v1/review", json=body)
        self._check(resp)
        task_id = resp.json()["task_id"]

        deadline = time.monotonic() + timeout
        result_data: dict = {}
        while time.monotonic() < deadline:
            status_resp = self._client.get(f"/api/v1/tasks/{task_id}")
            status_resp.raise_for_status()
            data = status_resp.json()

            if data["status"] == "completed":
                result_resp = self._client.get(f"/api/v1/tasks/{task_id}/result")
                result_resp.raise_for_status()
                result_data = result_resp.json()
                break

            if data["status"] == "failed":
                result_resp = self._client.get(f"/api/v1/tasks/{task_id}/result")
                result_resp.raise_for_status()
                err_data = result_resp.json()
                raise AIRError(
                    f"Review task {task_id} failed: {err_data.get('error', 'unknown error')}"
                )

            time.sleep(poll_interval)
        else:
            raise AIRError(f"Review task {task_id} timed out after {timeout}s")

        # Extract review content from the backend result
        inner = result_data.get("result", {})
        if isinstance(inner, dict) and "result" in inner:
            inner = inner["result"]
        md_text = inner.get("report", "") if isinstance(inner, dict) else ""
        pdf_b64 = inner.get("report_pdf_base64", "") if isinstance(inner, dict) else ""

        # Save locally
        import base64 as _b64
        import pathlib

        local_dir = output_dir or os.path.join(
            os.path.expanduser(self._local_dir), "reviews"
        )
        review_path = pathlib.Path(local_dir) / task_id
        review_path.mkdir(parents=True, exist_ok=True)

        md_path = review_path / "review.md"
        md_path.write_text(md_text, encoding="utf-8")

        pdf_path_str = ""
        if pdf_b64:
            pdf_path = review_path / "review.pdf"
            pdf_path.write_bytes(_b64.b64decode(pdf_b64))
            pdf_path_str = str(pdf_path)

        return {
            "review": md_text,
            "review_md_path": str(md_path),
            "review_pdf_path": pdf_path_str,
        }

    # ------------------------------------------------------------------
    # Project management
    # ------------------------------------------------------------------

    def create_project(
        self,
        name: str,
        data_description: str = "",
        iteration: int = 0,
        local_dir: str | None = None,
        auto_pull: bool = True,
    ) -> Project:
        """
        Create a new AIR research project.

        Args:
            name: Project name (used as directory name on server).
            data_description: Description of the research data/topic.
            iteration: Starting iteration number.

        Returns:
            A Project instance for further pipeline steps.
        """
        resp = self._client.post(
            "/api/v1/projects",
            json={
                "name": name,
                "data_description": data_description,
                "iteration": iteration,
            },
        )
        self._check(resp)
        if local_dir is not None:
            self._local_dir = local_dir
        return Project(name, self._client, base_url=self._base_url, api_key=self._api_key, local_dir=self._local_dir, auto_pull=auto_pull, data_description=data_description)

    def get_project(self, name: str, auto_pull: bool = True) -> Project:
        """
        Get an existing project by name.

        Args:
            name: Project name.

        Returns:
            A Project instance.

        Raises:
            AIRError: If the project does not exist.
        """
        resp = self._client.get(f"/api/v1/projects/{name}")
        self._check(resp)
        return Project(name, self._client, base_url=self._base_url, api_key=self._api_key, local_dir=self._local_dir, auto_pull=auto_pull)

    def list_projects(self) -> list[str]:
        """
        List all projects for this API key.

        Returns:
            List of project names.
        """
        resp = self._client.get("/api/v1/projects")
        self._check(resp)
        return resp.json().get("projects", [])

    def delete_project(self, name: str) -> dict:
        """Delete a project."""
        resp = self._client.delete(f"/api/v1/projects/{name}")
        self._check(resp)
        return resp.json()

    # ------------------------------------------------------------------
    # One-shot / deep-research (local execution via WebSocket)
    # ------------------------------------------------------------------

    def one_shot(
        self,
        task: str,
        *,
        model: str = "gpt-4.1-2025-04-14",
        max_rounds: int = 25,
        max_attempts: int = 1,
        agent: str = "engineer",
        work_dir: str | None = None,
        timeout: int = 86400,
        on_output: Callable[[str], None] | None = lambda s: print(s, end="\n", flush=True),
        python_path: str | None = None,
        venv_path: str | None = None,
        **config_overrides: Any,
    ) -> OneShotResult:
        """
        Execute a one-shot task with local code execution.

        The backend orchestrates the AI agent loop while all generated code
        runs locally on your machine in an isolated virtual environment.

        Args:
            task: Natural-language description of the task.
            model: LLM model for the backend agent.
            max_rounds: Maximum agent conversation rounds.
            max_attempts: Retry attempts on failure.
            agent: Agent type (``"engineer"``, etc.).
            work_dir: Local directory for code and outputs.
                Defaults to ``~/ai-scientist``.
            timeout: Maximum wall-clock seconds for the entire session.
            on_output: Callback invoked with streaming text.
            python_path: Explicit Python binary for the task venv.
            venv_path: Explicit path to an existing virtual environment.
                When *None*, the executor walks up from *work_dir*
                looking for a ``.venv`` directory; if none is found
                it creates one inside *work_dir*.
            **config_overrides: Extra keys merged into the config sent
                to the backend.

        Returns:
            An :class:`OneShotResult` with output text, result data,
            created files, and the local work directory.

        Raises:
            WebSocketError: If the WebSocket connection fails.
            AuthError: If the API key is invalid.
        """
        task_id = f"task_{uuid.uuid4().hex[:12]}"

        base_work_dir = work_dir or os.environ.get("AIR_WORK_DIR", "~/ai-scientist")
        base_work_dir = os.path.expanduser(base_work_dir)
        task_work_dir = os.path.join(base_work_dir, task_id)

        config = {
            "mode": "one-shot",
            "agent": agent,
            "model": model,
            "maxRounds": max_rounds,
            "maxAttempts": max_attempts,
            "remoteExecution": True,
            **config_overrides,
        }

        session = WebSocketSession(
            base_url=self._base_url,
            api_key=self._api_key,
            task_id=task_id,
            work_dir=task_work_dir,
            on_output=on_output,
            timeout=timeout,
            python_path=python_path,
            venv_path=venv_path,
        )

        coro = session.run(task, config)

        # Sync-to-async bridge: detect whether an event loop is already
        # running (e.g. Jupyter) and handle accordingly.
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None

        if loop is None:
            # Normal script — no loop running yet.
            return asyncio.run(coro)
        else:
            # Already inside an event loop (Jupyter, etc.) — run in a
            # background thread to avoid blocking the existing loop.
            import concurrent.futures

            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                future = pool.submit(asyncio.run, coro)
                return future.result()

    def deep_research(
        self,
        task: str,
        *,
        engineer_model: str = "gpt-4.1-2025-04-14",
        researcher_model: str = "gpt-4.1-2025-04-14",
        planner_model: str = "gpt-4.1-2025-04-14",
        plan_reviewer_model: str = "o3-mini-2025-01-31",
        max_plan_steps: int = 3,
        n_plan_reviews: int = 1,
        max_rounds_planning: int = 50,
        max_rounds: int = 100,
        max_attempts: int = 3,
        plan_instructions: str = "",
        engineer_instructions: str = "",
        researcher_instructions: str = "",
        hardware_constraints: str = "",
        enable_vlm_review: bool = False,
        vlm_model: str | None = None,
        max_vlm_review_attempts: int = 2,
        code_execution_timeout: int | None = None,
        work_dir: str | None = None,
        timeout: int = 86400,
        on_output: Callable[[str], None] | None = lambda s: print(s, end="\n", flush=True),
        python_path: str | None = None,
        venv_path: str | None = None,
        **config_overrides: Any,
    ) -> OneShotResult:
        """
        Execute a deep-research task with planning and multi-step control.

        The backend orchestrates a planning phase (planner + plan reviewer)
        followed by step-by-step control execution with specialised agents.
        All generated code runs locally on your machine.

        Args:
            task: Natural-language description of the research task.
            engineer_model: LLM model for the engineer agent.
            researcher_model: LLM model for the researcher agent.
            planner_model: LLM model for the planner agent.
            plan_reviewer_model: LLM model for the plan reviewer agent.
            max_plan_steps: Maximum number of plan steps.
            n_plan_reviews: Number of plan review rounds.
            max_rounds_planning: Maximum conversation rounds during planning.
            max_rounds: Maximum conversation rounds per control step.
            max_attempts: Retry attempts on failure per step.
            plan_instructions: Extra instructions for the planner.
            engineer_instructions: Extra instructions for the engineer.
            researcher_instructions: Extra instructions for the researcher.
            hardware_constraints: Hardware constraint description for agents.
            enable_vlm_review: Enable VLM-based review of generated plots.
            vlm_model: Model to use for VLM review (required if enabled).
            max_vlm_review_attempts: Max VLM review iterations.
            code_execution_timeout: Timeout in seconds for code execution.
            work_dir: Local directory for code and outputs.
                Defaults to ``~/ai-scientist``.
            timeout: Maximum wall-clock seconds for the entire session.
            on_output: Callback invoked with streaming text.
            python_path: Explicit Python binary for the task venv.
            venv_path: Explicit path to an existing virtual environment.
            **config_overrides: Extra keys merged into the config sent
                to the backend.

        Returns:
            An :class:`OneShotResult` with output text, result data,
            created files, and the local work directory.

        Raises:
            WebSocketError: If the WebSocket connection fails.
            AuthError: If the API key is invalid.
        """
        task_id = f"task_{uuid.uuid4().hex[:12]}"

        base_work_dir = work_dir or os.environ.get("AIR_WORK_DIR", "~/ai-scientist")
        base_work_dir = os.path.expanduser(base_work_dir)
        task_work_dir = os.path.join(base_work_dir, task_id)

        config = {
            "mode": "deep_research",
            "model": engineer_model,
            "researcherModel": researcher_model,
            "plannerModel": planner_model,
            "planReviewerModel": plan_reviewer_model,
            "maxPlanSteps": max_plan_steps,
            "nPlanReviews": n_plan_reviews,
            "maxRoundsPlanning": max_rounds_planning,
            "maxRounds": max_rounds,
            "maxAttempts": max_attempts,
            "planInstructions": plan_instructions,
            "engineerInstructions": engineer_instructions,
            "researcherInstructions": researcher_instructions,
            "hardwareConstraints": hardware_constraints,
            "enableVlmReview": enable_vlm_review,
            "vlmModel": vlm_model,
            "maxVlmReviewAttempts": max_vlm_review_attempts,
            "codeExecutionTimeout": code_execution_timeout,
            "remoteExecution": True,
            **config_overrides,
        }

        session = WebSocketSession(
            base_url=self._base_url,
            api_key=self._api_key,
            task_id=task_id,
            work_dir=task_work_dir,
            on_output=on_output,
            timeout=timeout,
            python_path=python_path,
            venv_path=venv_path,
        )

        coro = session.run(task, config)

        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None

        if loop is None:
            return asyncio.run(coro)
        else:
            import concurrent.futures

            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                future = pool.submit(asyncio.run, coro)
                return future.result()

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _check(self, resp: httpx.Response) -> None:
        """Raise appropriate exceptions for error responses."""
        if resp.status_code == 401:
            raise AuthError(resp.json().get("detail", "Authentication failed"))
        if resp.status_code >= 400:
            detail = ""
            try:
                body = resp.json()
                # Try common error keys; fall back to the full JSON body.
                detail = (
                    body.get("detail")
                    or body.get("error")
                    or body.get("message")
                    or json.dumps(body)
                )
                # If detail is a dict/list itself, serialize it for readability.
                if not isinstance(detail, str):
                    detail = json.dumps(detail)
            except Exception:
                detail = resp.text
            raise AIRError(
                f"[HTTP {resp.status_code}] {resp.request.method} {resp.request.url.path}: {detail}",
                status_code=resp.status_code,
            )

    def __repr__(self) -> str:
        return f"AIR(base_url='{self._base_url}')"
