"""Project class - wraps a single AIR research project."""

from __future__ import annotations

import asyncio
import base64
import os
import pathlib
import time
import uuid
from typing import TYPE_CHECKING, Any, Callable

from .exceptions import AIRError, TaskTimeoutError
from .ws_session import WebSocketSession

if TYPE_CHECKING:
    import httpx


class Project:
    """
    A AIR research project accessed via the AIR API.

    Provides methods for each pipeline step (idea, literature, methods,
    paper, review) and file access.

    Do not instantiate directly -- use ``AIR.create_project()`` or
    ``AIR.get_project()`` instead.
    """

    def __init__(
        self,
        name: str,
        client: "httpx.Client",
        base_url: str | None = None,
        api_key: str | None = None,
        poll_interval: float = 3.0,
        local_dir: str | None = None,
        auto_pull: bool = True,
        data_description: str = "",
    ):
        self.name = name
        self._client = client
        self._base_url = base_url
        self._api_key = api_key
        self._poll_interval = poll_interval
        self._local_dir = local_dir or os.environ.get("AIR_LOCAL_DIR", "~/ai-scientist")
        self._auto_pull_enabled = auto_pull
        self.data_description = data_description

        self.pull_files(key_only=True, iteration=0)

    def _poll_task(self, task_id: str, timeout: float) -> dict:
        """Poll a task until it completes or the timeout expires."""
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            resp = self._client.get(f"/api/v1/tasks/{task_id}")
            resp.raise_for_status()
            data = resp.json()

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

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

            time.sleep(self._poll_interval)

        raise TaskTimeoutError(task_id, timeout)

    def _run_step(self, step: str, timeout: float, **body_overrides: Any) -> dict:
        """Submit a pipeline step and poll until done."""
        resp = self._client.post(
            f"/api/v1/projects/{self.name}/{step}",
            json=body_overrides if body_overrides else {},
        )
        resp.raise_for_status()
        task_id = resp.json()["task_id"]
        return self._poll_task(task_id, timeout)

    # ------------------------------------------------------------------
    # Pipeline steps
    # ------------------------------------------------------------------

    def _auto_pull(self, iteration: int) -> None:
        """Pull key files locally if auto_pull is enabled and a local_dir is configured."""
        if self._auto_pull_enabled and self._local_dir:
            self.pull_files(key_only=True, iteration=iteration)

    def _read_output_file(self, path: str) -> str | None:
        """Try to read a project file, return None if not found."""
        try:
            return self.get_file(path)
        except Exception as e:
            print(f"[debug] Could not read output file '{path}': {e}")
            return None

    def idea(
        self,
        mode: str = "fast",
        iteration: int = 0,
        timeout: float = 600,
        **kwargs: Any,
    ) -> str:
        """
        Run idea generation.

        Args:
            mode: Execution mode ("fast" uses LangGraph, "full" uses P&C).
            iteration: Project iteration number.
            timeout: Max seconds to wait for completion.

        Returns:
            The generated idea text.
        """
        self._run_step("idea", timeout=timeout, mode=mode, iteration=iteration, **kwargs)
        result = self._read_output_file(f"Iteration{iteration}/input_files/idea.md") or ""
        self._auto_pull(iteration)
        return result

    def literature(self, iteration: int = 0, timeout: float = 600, **kwargs: Any) -> str:
        """Run literature search."""
        self._run_step("literature", timeout=timeout, iteration=iteration, **kwargs)
        return self._read_output_file(f"Iteration{iteration}/input_files/literature.md") or ""

    def methods(
        self,
        mode: str = "fast",
        iteration: int = 0,
        timeout: float = 600,
        **kwargs: Any,
    ) -> str:
        """Run methods development."""
        self._run_step("methods", timeout=timeout, mode=mode, iteration=iteration, **kwargs)
        result = self._read_output_file(f"Iteration{iteration}/input_files/methods.md") or ""
        self._auto_pull(iteration)
        return result

    def paper(
        self,
        journal: str = "NONE",
        iteration: int = 0,
        timeout: float = 900,
        **kwargs: Any,
    ) -> str:
        """Run paper writing.

        Args:
            journal: Target journal name.
            iteration: Project iteration number.
            timeout: Max seconds to wait for completion.

        Returns:
            The generated LaTeX paper text.
        """
        self._run_step(
            "paper",
            timeout=timeout,
            journal=journal if journal is not None else "NONE",
            iteration=iteration,
            **kwargs,
        )
        paper_files = {
            f["relative_path"]
            for f in self.list_files()
            if f["relative_path"].startswith(f"Iteration{iteration}/paper_output/paper_v")
            and f["relative_path"].endswith("_final.tex")
        }
        for version in range(4, 0, -1):
            candidate = f"Iteration{iteration}/paper_output/paper_v{version}_final.tex"
            if candidate in paper_files:
                result = self._read_output_file(candidate) or ""
                self._auto_pull(iteration)
                return result
        self._auto_pull(iteration)
        return ""

    def results(
        self,
        iteration: int = 0,
        timeout: int = 86400,
        max_attempts: int = 10,
        max_steps: int = 6,
        work_dir: str | None = None,
        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,
    ) -> str:
        """Run results / analysis via WebSocket with local code execution.

        The backend orchestrates the ``Experiment`` agent loop while all
        generated code runs locally on your machine.  After completion,
        any plots created locally are uploaded to the project on the
        server so that the subsequent ``paper()`` step can embed them.

        Args:
            iteration: Project iteration number.
            timeout: Max wall-clock seconds for the session.
            max_attempts: Retry attempts for the experiment.
            max_steps: Maximum experiment control steps.
            work_dir: Local directory for code and outputs.
                Defaults to ``~/ai-scientist/<task_id>``.
            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 WebSocket config.

        Returns:
            The generated results text.
        """
        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": "analysis",
            "projectName": self.name,
            "projectIteration": iteration,
            "maxAttempts": max_attempts,
            "maxSteps": max_steps,
            "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(self.data_description, config)

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

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

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

        # Upload locally-generated plots to the server so paper() can use them.
        n = self._upload_plots(task_work_dir, iteration)
        if on_output and n:
            on_output(f"Uploaded {n} plot(s) to project.")

        result = self._read_output_file(f"Iteration{iteration}/input_files/results.md") or ""
        self._auto_pull(iteration)
        return result

    def _upload_plots(self, local_work_dir: str, iteration: int) -> int:
        """Upload image files from *local_work_dir* to the project's plots folder.

        Searches common subdirectories where the executor writes plots and
        uploads each image via the REST file API.

        Returns:
            Number of files uploaded.
        """
        _IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".pdf", ".svg"}
        search_dirs = [
            os.path.join(local_work_dir, "control", "data"),
            os.path.join(local_work_dir, "data"),
            os.path.join(local_work_dir, "experiment", "control", "data"),
        ]

        uploaded = 0
        seen: set[str] = set()

        for search_dir in search_dirs:
            if not os.path.isdir(search_dir):
                continue
            for dirpath, _, filenames in os.walk(search_dir):
                for fname in filenames:
                    ext = os.path.splitext(fname)[1].lower()
                    if ext not in _IMAGE_EXTS:
                        continue
                    if fname in seen:
                        continue
                    seen.add(fname)

                    full_path = os.path.join(dirpath, fname)
                    with open(full_path, "rb") as f:
                        data = base64.b64encode(f.read()).decode("ascii")

                    remote_path = f"Iteration{iteration}/input_files/plots/{fname}"
                    self.write_file(remote_path, data, encoding="base64")
                    uploaded += 1

        return uploaded

    def review(self, iteration: int = 0, timeout: float = 600, **kwargs: Any) -> dict:
        """Run review.

        Keyword Args:
            review_engine (str): "denario" (default) or "skepthical".
            review_thoroughness (str): Skepthical only. "Standard" or "High".
            review_figures (bool): Skepthical only. Detailed figure review.
            review_verify_statements (bool): Skepthical only. Statement verification.
            review_maths (bool): Skepthical only. Mathematical audit.
            review_numerics (bool): Skepthical only. Numerical audit.
            pdf_version (str): Paper version to review ("v1"-"v4", default "v2").
            emails (list[str]): Email recipients for the review report.

        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 was not available).
        """
        result = self._run_step("review", timeout=timeout, iteration=iteration, **kwargs)

        # Read markdown from the server
        md_text = self._read_output_file(f"Iteration{iteration}/input_files/referee.md") or ""

        # Pull review files (MD + PDF) into the local review folder
        review_md_path = ""
        review_pdf_path = ""

        if self._local_dir:
            root = pathlib.Path(self._local_dir).expanduser() / self.name
            review_dir = root / f"Iteration{iteration}" / "review"
            review_dir.mkdir(parents=True, exist_ok=True)

            # Save markdown locally
            local_md = review_dir / "referee.md"
            local_md.write_text(md_text, encoding="utf-8")
            review_md_path = str(local_md)

            # Try to get the PDF from the server (saved by the backend into review/)
            try:
                pdf_content = self.get_file(f"Iteration{iteration}/input_files/review/review.pdf")
                if pdf_content:
                    local_pdf = review_dir / "review.pdf"
                    local_pdf.write_bytes(base64.b64decode(pdf_content))
                    review_pdf_path = str(local_pdf)
            except Exception:
                # PDF may not be available (e.g. denario engine doesn't produce one)
                pass

        self._auto_pull(iteration)

        return {
            "review": md_text,
            "review_md_path": review_md_path,
            "review_pdf_path": review_pdf_path,
        }

    # ------------------------------------------------------------------
    # File access
    # ------------------------------------------------------------------

    def get_file(self, path: str) -> str:
        """
        Read a file from the project.

        Args:
            path: Relative path within the project (e.g. "Iteration0/input_files/idea.md").

        Returns:
            File content as a string (base64 for binary files).
        """
        resp = self._client.get(f"/api/v1/projects/{self.name}/files/{path}")
        resp.raise_for_status()
        return resp.json()["content"]

    def list_files(self) -> list[dict]:
        """
        List all files in the project.

        Returns:
            List of dicts with keys: relative_path, size, mtime.
        """
        resp = self._client.get(f"/api/v1/projects/{self.name}")
        resp.raise_for_status()
        return resp.json().get("files", [])

    def write_file(self, path: str, content: str, encoding: str = "utf-8") -> dict:
        """
        Write a file to the project.

        Args:
            path: Relative path within the project.
            content: File content.
            encoding: "utf-8" or "base64".

        Returns:
            Server response dict.
        """
        resp = self._client.post(
            f"/api/v1/projects/{self.name}/files",
            json={"path": path, "content": content, "encoding": encoding},
        )
        resp.raise_for_status()
        return resp.json()

    _BINARY_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg", ".svg", ".gif", ".zip"}

    def pull_files(
        self,
        local_dir: str | None = None,
        iteration: int = 0,
        key_only: bool = False,
    ) -> list[str]:
        """
        Download project files from the backend to a local directory.

        The directory structure mirrors the backend paths under
        ``<local_dir>/<project_name>/``.

        Args:
            local_dir: Root directory to write files into. Falls back to the
                ``local_dir`` passed at client construction, then to the
                ``AIR_LOCAL_DIR`` environment variable.
            iteration: Iteration number used when ``key_only=True``.
            key_only: When ``True``, only download files under
                ``Iteration{iteration}/input_files`` and
                ``Iteration{iteration}/paper_output`` (if present).

        Returns:
            List of local file paths that were written.

        Raises:
            ValueError: If no local directory is configured.
        """
        dest = local_dir or self._local_dir
        if not dest:
            raise ValueError(
                "No local directory specified. Pass local_dir= to pull(), "
                "set it on the AIR client, or set the AIR_LOCAL_DIR environment variable."
            )

        root = pathlib.Path(dest).expanduser() / self.name
        written: list[str] = []

        _key_prefixes = (
            f"Iteration{iteration}/input_files/",
            f"Iteration{iteration}/paper_output/",
        )

        for f in self.list_files():
            rel = f["relative_path"]
            if key_only and not any(rel.startswith(p) for p in _key_prefixes):
                continue
            content = self.get_file(rel)

            local_path = root / rel
            local_path.parent.mkdir(parents=True, exist_ok=True)

            ext = pathlib.Path(rel).suffix.lower()
            if ext in self._BINARY_EXTENSIONS:
                local_path.write_bytes(base64.b64decode(content))
            else:
                local_path.write_text(content, encoding="utf-8")

            written.append(str(local_path))

        return written

    def push_files(self, local_path: str, remote_path: str | None = None) -> list[str]:
        """
        Upload a local file or directory to the backend project.

        The remote path is determined as follows:
        - If ``remote_path`` is given, it is used as-is (for a single file) or
          as the remote prefix (for a directory).
        - Otherwise the path relative to ``<local_dir>/<project_name>/`` is used,
          so the local structure mirrors the backend structure.

        Args:
            local_path: Local file or directory to upload.
            remote_path: Remote path override (relative to the project root).

        Returns:
            List of remote paths that were uploaded.

        Raises:
            ValueError: If ``local_path`` does not exist.
        """
        src = pathlib.Path(local_path).expanduser().resolve()
        if not src.exists():
            raise ValueError(f"Path does not exist: {src}")

        # Base used to compute relative remote paths when no override is given.
        base = (pathlib.Path(self._local_dir).expanduser() / self.name).resolve() if self._local_dir else src.parent

        files = [src] if src.is_file() else sorted(src.rglob("*"))
        uploaded: list[str] = []

        for file in files:
            if not file.is_file():
                continue

            if remote_path is not None:
                if src.is_file():
                    rel = remote_path
                else:
                    rel = str(pathlib.Path(remote_path) / file.relative_to(src))
            else:
                try:
                    rel = str(file.relative_to(base))
                except ValueError:
                    rel = file.name

            ext = file.suffix.lower()
            if ext in self._BINARY_EXTENSIONS:
                content = base64.b64encode(file.read_bytes()).decode("ascii")
                self.write_file(rel, content, encoding="base64")
            else:
                self.write_file(rel, file.read_text(encoding="utf-8"), encoding="utf-8")

            uploaded.append(rel)

        return uploaded

    def delete(self) -> dict:
        """Delete this project from the server."""
        resp = self._client.delete(f"/api/v1/projects/{self.name}")
        resp.raise_for_status()
        return resp.json()

    def __repr__(self) -> str:
        return f"Project('{self.name}')"
