"""
Y Build Container - High-level container operations.
"""

from typing import Optional, List, Dict, Any, Iterator, Callable, TYPE_CHECKING
import base64
import json
import shlex

from .models import ContainerInfo, ContainerStatus, ExecResult, FileInfo, ProcessInfo
from .exceptions import ContainerNotFoundError, ContainerStoppedError

if TYPE_CHECKING:
    from .client import YBuildClient


class Container:
    """
    Container instance for executing commands and managing files.

    Example:
        container = client.create_container(thread_id="my-thread")

        # Execute commands
        result = container.exec("python -c 'print(1+1)'")
        print(result.stdout)  # "2"

        # File operations
        container.write_file("/workspace/script.py", "print('hello')")
        content = container.read_file("/workspace/script.py")

        # Stop when done
        container.stop()
    """

    def __init__(self, client: "YBuildClient", info: ContainerInfo):
        """
        Initialize container wrapper.

        Args:
            client: YBuildClient instance
            info: Container information
        """
        self._client = client
        self._info = info

    @property
    def id(self) -> str:
        """Container ID."""
        return self._info.container_id

    @property
    def thread_id(self) -> str:
        """Thread ID."""
        return self._info.thread_id

    @property
    def status(self) -> ContainerStatus:
        """Current status (may be stale, call refresh() to update)."""
        return self._info.status

    @property
    def info(self) -> ContainerInfo:
        """Full container info."""
        return self._info

    @property
    def is_running(self) -> bool:
        """Check if container is running."""
        return self._info.status == ContainerStatus.RUNNING

    @property
    def ports(self) -> Dict[int, int]:
        """Port mappings: {container_port: host_port}."""
        return self._info.ports

    @property
    def urls(self) -> Dict[int, str]:
        """
        Preview URLs for exposed ports.

        Example:
            container.urls[5173]  # "https://thread-123-5173.zbuild.ai"
        """
        return self._info.urls

    def get_preview_url(self, port: int) -> Optional[str]:
        """
        Get preview URL for a specific port.

        Args:
            port: Container port (e.g., 5173, 3000, 8080)

        Returns:
            Preview URL or None if port not exposed
        """
        return self._info.urls.get(port)

    def _check_running(self):
        """Raise error if container is not running."""
        if self._info.status.is_terminal():
            raise ContainerStoppedError(
                f"Container {self.id} is {self._info.status.value}"
            )

    # ==================== Container Lifecycle ====================

    def refresh(self) -> "Container":
        """
        Refresh container status from API.

        Returns:
            Self for chaining
        """
        data = self._client._request("GET", f"/api/containers/{self.id}/status")
        self._info = ContainerInfo.from_dict(data)
        return self

    def wait_until_running(self, timeout: float = 60.0, poll_interval: float = 1.0) -> "Container":
        """
        Wait until container is running.

        Args:
            timeout: Maximum time to wait in seconds
            poll_interval: Time between status checks

        Returns:
            Self for chaining

        Raises:
            TimeoutError: If container doesn't start within timeout
            ContainerError: If container enters failed state
        """
        import time
        start = time.time()

        while time.time() - start < timeout:
            self.refresh()

            if self._info.status == ContainerStatus.RUNNING:
                return self

            if self._info.status.is_terminal():
                raise ContainerStoppedError(
                    f"Container {self.id} entered terminal state: {self._info.status.value}"
                )

            time.sleep(poll_interval)

        from .exceptions import TimeoutError
        raise TimeoutError(
            f"Container {self.id} did not start within {timeout}s (status: {self._info.status.value})"
        )

    def stop(self) -> "Container":
        """
        Stop the container.

        Returns:
            Self for chaining
        """
        self._client._request("DELETE", f"/api/containers/{self.id}")
        self._info.status = ContainerStatus.STOPPED
        return self

    def start(self) -> "Container":
        """
        Start a stopped container.

        Returns:
            Self for chaining
        """
        data = self._client._request("POST", f"/api/containers/{self.id}/start")
        self._info = ContainerInfo.from_dict(data)
        return self

    def restore(self) -> "Container":
        """
        Restore an archived container.

        Returns:
            Self for chaining
        """
        data = self._client._request("POST", f"/api/containers/{self.id}/restore")
        self._info = ContainerInfo.from_dict(data)
        return self

    def update_activity(self) -> "Container":
        """
        Update last activity timestamp (prevent idle timeout).

        Returns:
            Self for chaining
        """
        self._client._request("POST", f"/api/containers/{self.id}/activity")
        return self

    # ==================== Command Execution ====================

    def exec(
        self,
        command: str,
        *,
        timeout: int = 300,
        workdir: Optional[str] = None,
        user: Optional[str] = None,
        env: Optional[Dict[str, str]] = None,
    ) -> ExecResult:
        """
        Execute a command in the container.

        Args:
            command: Shell command to execute
            timeout: Command timeout in seconds (default: 5 min)
            workdir: Working directory (default: /workspace)
            user: User to run as (default: container default)
            env: Additional environment variables

        Returns:
            ExecResult with exit_code, stdout, stderr

        Example:
            result = container.exec("python script.py")
            if result.success:
                print(result.stdout)
            else:
                print(f"Error: {result.stderr}")
        """
        self._check_running()

        payload: Dict[str, Any] = {
            "command": command,
            "timeout": timeout,
        }
        if workdir:
            payload["workdir"] = workdir
        if user:
            payload["user"] = user
        if env:
            payload["env"] = env

        data = self._client._request(
            "POST",
            f"/api/containers/{self.id}/exec",
            json=payload,
        )

        return ExecResult.from_dict(data)

    def exec_python(self, code: str, timeout: int = 300) -> ExecResult:
        """
        Execute Python code.

        Args:
            code: Python code to execute
            timeout: Timeout in seconds

        Returns:
            ExecResult
        """
        return self.exec(f"python3 -c {shlex.quote(code)}", timeout=timeout)

    def exec_node(self, code: str, timeout: int = 300) -> ExecResult:
        """
        Execute Node.js code.

        Args:
            code: JavaScript code to execute
            timeout: Timeout in seconds

        Returns:
            ExecResult
        """
        return self.exec(f"node -e {shlex.quote(code)}", timeout=timeout)

    def exec_stream(
        self,
        command: str,
        *,
        timeout: int = 300,
        workdir: Optional[str] = None,
        user: Optional[str] = None,
        on_stdout: Optional[Callable[[str], None]] = None,
        on_stderr: Optional[Callable[[str], None]] = None,
    ) -> ExecResult:
        """
        Execute a command with streaming output.

        Output is delivered in real-time via callbacks. This is useful for
        long-running commands like `npm install` where you want to see
        progress as it happens.

        Requires: pip install websockets

        Args:
            command: Shell command to execute
            timeout: Command timeout in seconds
            workdir: Working directory
            user: User to run as
            on_stdout: Callback for stdout chunks
            on_stderr: Callback for stderr chunks

        Returns:
            ExecResult with combined output

        Example:
            result = container.exec_stream(
                "npm install",
                on_stdout=lambda s: print(s, end=""),
                on_stderr=lambda s: print(s, end="", file=sys.stderr),
            )
        """
        self._check_running()

        try:
            import websockets.sync.client as ws_sync
        except ImportError:
            raise ImportError(
                "websockets is required for exec_stream. "
                "Install it with: pip install ybuild[stream]"
            )

        # Build WebSocket URL from HTTP URL
        ws_url = self._client.api_url.replace("https://", "wss://").replace("http://", "ws://")
        ws_url = f"{ws_url}/api/containers/{self.id}/exec/stream"

        # Build headers for auth
        headers = {}
        if self._client.api_key:
            headers["Authorization"] = f"Bearer {self._client.api_key}"

        stdout_parts = []
        stderr_parts = []
        exit_code = -1
        timed_out = False

        with ws_sync.connect(ws_url, additional_headers=headers) as ws:
            # Send command
            ws.send(json.dumps({
                "command": command,
                "timeout": timeout,
                "workdir": workdir,
                "user": user,
            }))

            # Receive streaming output
            for raw_msg in ws:
                msg = json.loads(raw_msg)
                msg_type = msg.get("type")

                if msg_type == "stdout":
                    data = msg.get("data", "")
                    stdout_parts.append(data)
                    if on_stdout:
                        on_stdout(data)

                elif msg_type == "stderr":
                    data = msg.get("data", "")
                    stderr_parts.append(data)
                    if on_stderr:
                        on_stderr(data)

                elif msg_type == "exit":
                    exit_code = msg.get("exit_code", -1)
                    timed_out = msg.get("timed_out", False)
                    break

                elif msg_type == "error":
                    stderr_parts.append(msg.get("message", ""))
                    break

        return ExecResult(
            exit_code=exit_code,
            stdout="".join(stdout_parts),
            stderr="".join(stderr_parts),
            timed_out=timed_out,
        )

    # ==================== Process Management ====================

    def start_process(
        self,
        command: str,
        *,
        workdir: Optional[str] = None,
        user: Optional[str] = None,
    ) -> ProcessInfo:
        """
        Start a long-running background process.

        Unlike exec() which blocks until completion, this starts
        the process and returns immediately with a process_id.

        Args:
            command: Shell command to run
            workdir: Working directory
            user: User to run as

        Returns:
            ProcessInfo with process_id

        Example:
            proc = container.start_process("npm run dev")
            # ... do other work ...
            info = container.get_process(proc.process_id)
            print(info.stdout)  # latest output
        """
        self._check_running()

        payload: Dict[str, Any] = {"command": command}
        if workdir:
            payload["workdir"] = workdir
        if user:
            payload["user"] = user

        data = self._client._request(
            "POST",
            f"/api/containers/{self.id}/processes",
            json=payload,
        )
        return ProcessInfo.from_dict(data)

    def list_processes(self) -> List[ProcessInfo]:
        """
        List all background processes.

        Returns:
            List of ProcessInfo
        """
        self._check_running()

        data = self._client._request(
            "GET",
            f"/api/containers/{self.id}/processes",
        )
        processes = data.get("processes", [])
        return [ProcessInfo.from_dict(p) for p in processes]

    def get_process(self, process_id: str) -> ProcessInfo:
        """
        Get process status and output.

        Args:
            process_id: Process ID from start_process()

        Returns:
            ProcessInfo with current stdout/stderr
        """
        self._check_running()

        data = self._client._request(
            "GET",
            f"/api/containers/{self.id}/processes/{process_id}",
        )
        return ProcessInfo.from_dict(data)

    def stop_process(self, process_id: str) -> None:
        """
        Stop a running background process.

        Args:
            process_id: Process ID to stop
        """
        self._check_running()

        self._client._request(
            "POST",
            f"/api/containers/{self.id}/processes/{process_id}/stop",
        )

    # ==================== File Operations ====================

    def write_file(
        self,
        path: str,
        content: str,
        *,
        create_dirs: bool = True,
    ) -> bool:
        """
        Write content to a file.

        Args:
            path: File path in container
            content: File content
            create_dirs: Create parent directories if needed

        Returns:
            True on success
        """
        self._check_running()

        payload = {
            "path": path,
            "content": content,
            "create_dirs": create_dirs,
        }

        self._client._request(
            "POST",
            f"/api/containers/{self.id}/files/write",
            json=payload,
        )
        return True

    def write_file_bytes(
        self,
        path: str,
        content: bytes,
        *,
        create_dirs: bool = True,
    ) -> bool:
        """
        Write binary content to a file.

        Args:
            path: File path in container
            content: File content as bytes
            create_dirs: Create parent directories if needed

        Returns:
            True on success
        """
        import base64 as _base64

        self._check_running()

        self._client._request(
            "POST",
            f"/api/containers/{self.id}/files/write",
            json={
                "path": path,
                "content": _base64.b64encode(content).decode("ascii"),
                "encoding": "base64",
                "create_dirs": create_dirs,
            },
        )
        return True

    def read_file(self, path: str) -> str:
        """
        Read file content.

        Args:
            path: File path in container

        Returns:
            File content as string
        """
        self._check_running()

        data = self._client._request(
            "GET",
            f"/api/containers/{self.id}/files/read",
            params={"path": path},
        )

        return data.get("content", "")

    def delete_file(self, path: str, *, recursive: bool = False) -> bool:
        """
        Delete a file or directory.

        Args:
            path: Path to delete
            recursive: Delete directories recursively

        Returns:
            True on success
        """
        self._check_running()

        self._client._request(
            "DELETE",
            f"/api/containers/{self.id}/files",
            params={"path": path, "recursive": recursive},
        )
        return True

    def list_files(self, path: str = "/workspace", recursive: bool = False) -> List[FileInfo]:
        """
        List files in a directory.

        Args:
            path: Directory path
            recursive: List files recursively

        Returns:
            List of FileInfo
        """
        self._check_running()

        data = self._client._request(
            "GET",
            f"/api/containers/{self.id}/files/list",
            params={"path": path, "recursive": recursive},
        )

        files = data.get("files", data) if isinstance(data, dict) else data
        return [FileInfo.from_dict(f) for f in files]

    def upload_file(self, path: str, content: bytes) -> bool:
        """
        Upload binary file.

        Args:
            path: Target path in container
            content: Binary content

        Returns:
            True on success
        """
        self._check_running()

        # Encode as base64
        encoded = base64.b64encode(content).decode("utf-8")

        self._client._request(
            "POST",
            f"/api/containers/{self.id}/files/upload",
            json={"path": path, "content": encoded, "encoding": "base64"},
        )
        return True

    def download_file(self, path: str) -> bytes:
        """
        Download file as binary.

        Args:
            path: File path in container

        Returns:
            File content as bytes
        """
        self._check_running()

        data = self._client._request(
            "GET",
            f"/api/containers/{self.id}/files/download",
            params={"path": path},
        )

        content = data.get("content", "")
        encoding = data.get("encoding", "base64")

        if encoding == "base64":
            return base64.b64decode(content)
        return content.encode("utf-8")

    def file_exists(self, path: str) -> bool:
        """
        Check if a file exists.

        Args:
            path: File path

        Returns:
            True if file exists
        """
        result = self.exec(f"test -e {path} && echo exists || echo missing")
        return "exists" in result.stdout

    # ==================== Convenience Methods ====================

    def run_python_file(self, path: str, *, timeout: int = 300) -> ExecResult:
        """
        Run a Python file.

        Args:
            path: Path to Python file
            timeout: Timeout in seconds

        Returns:
            ExecResult
        """
        return self.exec(f"python3 {path}", timeout=timeout)

    def run_node_file(self, path: str, *, timeout: int = 300) -> ExecResult:
        """
        Run a Node.js file.

        Args:
            path: Path to JavaScript file
            timeout: Timeout in seconds

        Returns:
            ExecResult
        """
        return self.exec(f"node {path}", timeout=timeout)

    def install_pip(self, *packages: str, timeout: int = 300) -> ExecResult:
        """
        Install Python packages with pip.

        Args:
            *packages: Package names
            timeout: Timeout in seconds

        Returns:
            ExecResult
        """
        pkg_list = " ".join(packages)
        return self.exec(f"python3 -m pip install {pkg_list}", timeout=timeout)

    def install_npm(self, *packages: str, timeout: int = 300) -> ExecResult:
        """
        Install npm packages.

        Args:
            *packages: Package names
            timeout: Timeout in seconds

        Returns:
            ExecResult
        """
        pkg_list = " ".join(packages)
        return self.exec(f"npm install {pkg_list}", timeout=timeout)

    def __repr__(self) -> str:
        return f"Container(id={self.id!r}, status={self.status.value!r})"
