"""
vr pull command - Pull agent code from the server.
"""

import base64
import io
import zipfile

import typer
import yaml

from pathlib import Path
from typing import Optional

from vr_cli.utils.utils import (
    console,
    make_request,
    print_success,
    print_error,
    print_info,
    confirm,
    calculate_project_checksum,
)
from vr_cli.utils.agent_lock import load_agent_lock, save_agent_lock
from vr_cli.entities.agent import AgentRepository


def ensure_agent_yaml(voicerun_dir: Path, agent_name: str, agent_description: str = ""):
    """Create agent.yaml in .voicerun dir if it doesn't exist."""
    agent_yaml_path = voicerun_dir / "agent.yaml"
    if agent_yaml_path.exists():
        return

    data = {
        "apiVersion": "voicerun/v1",
        "name": agent_name,
        "description": agent_description or "",
    }
    with open(agent_yaml_path, "w") as f:
        yaml.dump(data, f, default_flow_style=False, sort_keys=False)


def extract_code(code: str, target_dir: str) -> list[str]:
    """Extract code into target_dir.

    If the code is a base64-encoded ZIP, decode and extract it.
    Otherwise, treat it as plain text and write it to handler.py.

    Returns list of extracted file paths (relative to target_dir).
    """
    # Try to detect if code is a base64-encoded ZIP
    try:
        padded = code
        padding = 4 - len(padded) % 4
        if padding != 4:
            padded += "=" * padding
        zip_bytes = base64.b64decode(padded, validate=True)
        # Check for ZIP magic bytes
        if zip_bytes[:4] == b"PK\x03\x04":
            extracted = []
            with zipfile.ZipFile(io.BytesIO(zip_bytes), "r") as zf:
                for info in zf.infolist():
                    if info.is_dir():
                        continue
                    zf.extract(info, target_dir)
                    extracted.append(info.filename)
            return extracted
    except Exception:
        pass

    # Plain text code — write to handler.py
    handler_path = Path(target_dir) / "handler.py"
    handler_path.write_text(code)
    return ["handler.py"]


def pull(
    agent_id: Optional[str] = typer.Argument(
        None,
        help="Agent ID to pull. Not required if inside a voicerun project.",
    ),
    output: Optional[str] = typer.Option(
        None,
        "--output",
        "-o",
        help="Output directory (defaults to agent name).",
    ),
    yes: bool = typer.Option(
        False,
        "--yes",
        "-y",
        help="Skip confirmation prompt.",
    ),
):
    """Pull agent code from the server."""
    project_dir = Path.cwd()
    voicerun_dir = project_dir / ".voicerun"
    is_existing_project = voicerun_dir.exists()

    # Flow A: Inside a voicerun project
    if is_existing_project and agent_id is None:
        agent_lock = load_agent_lock(voicerun_dir)
        if not agent_lock or not agent_lock.get("id"):
            print_error("agent.lock is missing or invalid. Cannot determine which agent to pull.")
            print_info("Provide an agent ID: vr pull <agent-id>")
            raise typer.Exit(1)

        lock_agent_id = agent_lock["id"]
        lock_function_id = agent_lock.get("functionId")

        if not yes:
            console.print()
            console.print("[bold]Pull Summary:[/bold]")
            console.print(f"  Agent ID: [cyan]{lock_agent_id[:8]}...[/cyan]")
            if lock_function_id:
                console.print(f"  Function ID: [cyan]{lock_function_id[:8]}...[/cyan]")
            console.print(f"  Target: [cyan]{project_dir}[/cyan]")
            console.print()
            if not confirm("Pull and overwrite local files?", default=True):
                print_info("Aborted.")
                raise typer.Exit(0)

        print_info("Pulling code from server...")
        payload = {"agentId": lock_agent_id}
        if lock_function_id:
            payload["functionId"] = lock_function_id

        response = make_request("/v1/agents/pull", method="POST", json=payload)
        if not response:
            print_error("Failed to pull code from server.")
            raise typer.Exit(1)

        code = response.get("code")
        if not code:
            print_error("Server returned no code.")
            raise typer.Exit(1)

        extracted = extract_code(code, str(project_dir))

        # Ensure agent.yaml exists
        agent_repo = AgentRepository()
        agent = agent_repo.get_by_id(lock_agent_id)
        agent_name = agent.name if agent and agent.name else lock_agent_id[:8]
        agent_desc = agent.description if agent and agent.description else ""
        ensure_agent_yaml(voicerun_dir, agent_name, agent_desc)

        # Update agent.lock with returned functionId + new checksum
        response_function_id = response.get("functionId", lock_function_id)
        checksum = calculate_project_checksum(str(project_dir))
        save_agent_lock(voicerun_dir, {
            "id": lock_agent_id,
            "functionId": response_function_id,
            "checksum": checksum,
        })

        console.print()
        print_success(f"Pulled {len(extracted)} file(s).")
        for f in extracted:
            console.print(f"  [dim]{f}[/dim]")

    # Flow B: No voicerun project — bootstrap a new directory
    else:
        if agent_id is None:
            print_error("Not inside a voicerun project and no agent ID provided.")
            print_info("Usage: vr pull <agent-id>")
            print_info("Or run from inside a voicerun project directory.")
            raise typer.Exit(1)

        print_info("Pulling code from server...")
        payload = {"agentId": agent_id}
        response = make_request("/v1/agents/pull", method="POST", json=payload)
        if not response:
            print_error("Failed to pull code from server.")
            raise typer.Exit(1)

        code = response.get("code")
        if not code:
            print_error("Server returned no code.")
            raise typer.Exit(1)

        response_agent_id = response.get("agentId", agent_id)
        response_function_id = response.get("functionId")

        # Fetch agent metadata
        agent_repo = AgentRepository()
        agent = agent_repo.get_by_id(agent_id)

        # Determine output directory
        if output:
            target_dir = Path(output)
        else:
            dir_name = agent.name if agent and agent.name else agent_id[:8]
            target_dir = project_dir / dir_name

        if target_dir.exists() and any(target_dir.iterdir()):
            if not yes:
                if not confirm(f"Directory '{target_dir}' already exists. Overwrite?", default=False):
                    print_info("Aborted.")
                    raise typer.Exit(0)

        target_dir.mkdir(parents=True, exist_ok=True)

        extracted = extract_code(code, str(target_dir))

        # Create .voicerun directory, agent.yaml, and agent.lock
        new_voicerun_dir = target_dir / ".voicerun"
        new_voicerun_dir.mkdir(exist_ok=True)

        agent_name = agent.name if agent and agent.name else agent_id[:8]
        agent_desc = agent.description if agent and agent.description else ""
        ensure_agent_yaml(new_voicerun_dir, agent_name, agent_desc)

        checksum = calculate_project_checksum(str(target_dir))
        save_agent_lock(new_voicerun_dir, {
            "id": response_agent_id,
            "functionId": response_function_id,
            "checksum": checksum,
        })

        console.print()
        print_success(f"Pulled {len(extracted)} file(s) into '{target_dir}'.")
        for f in extracted:
            console.print(f"  [dim]{f}[/dim]")
        console.print()
        print_info(f"Next steps:")
        console.print(f"  cd {target_dir}")
        console.print(f"  # edit your code")
        console.print(f"  vr push")
