"""Cloud platform knowledge base.

GPU instance types, pricing, and command templates for deploying
training workloads on major cloud providers.
"""

from __future__ import annotations

CLOUD_KNOWLEDGE: dict[str, dict] = {
    "gcp": {
        "name": "Google Cloud Platform",
        "gpu_instances": {
            "a100-40gb": {"type": "a2-highgpu-1g", "cost": "$3.67/hr"},
            "a100-80gb": {"type": "a2-ultragpu-1g", "cost": "$5.00/hr"},
            "l4": {"type": "g2-standard-4", "cost": "$0.70/hr"},
            "t4": {"type": "n1-standard-4 + T4", "cost": "$0.35/hr"},
        },
        "command_template": (
            "gcloud compute instances create {instance_name} \\\n"
            "  --machine-type={machine_type} \\\n"
            "  --accelerator=count=1,type={gpu_type} \\\n"
            "  --image-family=pytorch-latest-gpu \\\n"
            "  --image-project=deeplearning-platform-release \\\n"
            "  --boot-disk-size=200GB"
        ),
        "tips": [
            "Use preemptible/spot instances for 60-91% savings",
            "Cloud Run is for serving, not training — use Compute Engine or Vertex AI",
        ],
    },
    "aws": {
        "name": "Amazon Web Services",
        "gpu_instances": {
            "a100-40gb": {"type": "p4d.24xlarge (8x A100)", "cost": "$32.77/hr"},
            "a10g": {"type": "g5.xlarge", "cost": "$1.01/hr"},
            "t4": {"type": "g4dn.xlarge", "cost": "$0.53/hr"},
        },
        "command_template": (
            "aws ec2 run-instances \\\n"
            "  --instance-type {instance_type} \\\n"
            "  --image-id {ami_id} \\\n"
            "  --key-name {key_name}"
        ),
        "tips": [
            "g5.xlarge (A10G) is the best value for single-GPU training",
            "Use spot instances for non-critical training runs",
        ],
    },
    "runpod": {
        "name": "RunPod",
        "gpu_instances": {
            "a100-80gb": {"type": "A100 80GB", "cost": "$1.64/hr"},
            "a100-40gb": {"type": "A100 40GB", "cost": "$1.24/hr"},
            "a6000": {"type": "A6000 48GB", "cost": "$0.76/hr"},
            "4090": {"type": "RTX 4090 24GB", "cost": "$0.44/hr"},
        },
        "command_template": (
            "runpodctl create pod \\\n"
            '  --name {instance_name} \\\n'
            '  --gpuType "{gpu_type}" \\\n'
            "  --imageName pytorch/pytorch:latest \\\n"
            "  --volumeSize 50"
        ),
        "tips": [
            "Best value for single-GPU training runs",
            "Serverless GPUs available for burst workloads",
            "Pre-built templates for common ML frameworks",
        ],
    },
    "lambda": {
        "name": "Lambda Cloud",
        "gpu_instances": {
            "a100-80gb": {"type": "1x A100 80GB", "cost": "$1.29/hr"},
            "h100-80gb": {"type": "1x H100 80GB", "cost": "$2.49/hr"},
        },
        "command_template": (
            "lambda cloud create-instance \\\n"
            "  --instance-type {instance_type} \\\n"
            "  --ssh-key {ssh_key}"
        ),
        "tips": [
            "Cheapest A100s among major providers",
            "Good availability for 1x GPU instances",
        ],
    },
    "osmo": {
        "name": "NVIDIA OSMO",
        "description": "Orchestrates across heterogeneous clusters — not a cloud provider",
        "gpu_instances": {},
        "command_template": "osmo workflow submit {workflow_yaml} --pool {pool_name}",
        "tips": [
            "Requires Kubernetes cluster with OSMO deployed",
            "Platforms: gb200, rtx-pro-6000, jetson-agx-thor",
            "Best for multi-stage pipelines (gate -> train -> verify)",
        ],
    },
}


def recommend_cloud_gpu(
    vram_needed_gb: int,
    budget_priority: bool = True,
) -> list[dict]:
    """Recommend cloud GPU instances based on VRAM requirements.

    Returns list of {provider, gpu, instance_type, cost, notes} sorted by cost.
    """
    recommendations: list[dict] = []

    for provider_key, provider in CLOUD_KNOWLEDGE.items():
        for gpu_key, gpu_info in provider.get("gpu_instances", {}).items():
            # Parse VRAM from gpu key
            vram = _parse_vram(gpu_key)
            if vram and vram >= vram_needed_gb:
                cost_str = gpu_info.get("cost", "unknown")
                cost_val = _parse_cost(cost_str)
                recommendations.append({
                    "provider": provider["name"],
                    "provider_key": provider_key,
                    "gpu": gpu_key,
                    "instance_type": gpu_info["type"],
                    "cost": cost_str,
                    "cost_value": cost_val,
                })

    if budget_priority:
        recommendations.sort(key=lambda x: x["cost_value"])

    return recommendations


def _parse_vram(gpu_key: str) -> int | None:
    """Extract VRAM in GB from a GPU key like 'a100-80gb' or '4090'."""
    gpu_key = gpu_key.lower()
    # Known mappings
    known = {
        "4090": 24,
        "a6000": 48,
        "a10g": 24,
        "t4": 16,
        "l4": 24,
    }
    for k, v in known.items():
        if k in gpu_key:
            return v
    # Try to parse XBgb pattern
    import re

    m = re.search(r"(\d+)gb", gpu_key)
    if m:
        return int(m.group(1))
    return None


def _parse_cost(cost_str: str) -> float:
    """Parse cost string like '$1.29/hr' to float."""
    import re

    m = re.search(r"\$([\d.]+)", cost_str)
    if m:
        return float(m.group(1))
    return 999.0  # unknown = expensive


def get_cloud_summary() -> str:
    """Generate concise cloud provider reference for the system prompt."""
    parts: list[str] = []
    for key, provider in CLOUD_KNOWLEDGE.items():
        if key == "osmo":
            parts.append(f"### {provider['name']} (Orchestration)")
            parts.append("- Not a cloud provider — orchestrates across clusters")
            parts.append(f"- Command: `{provider['command_template']}`")
        else:
            parts.append(f"### {provider['name']}")
            for gpu_key, gpu_info in provider.get("gpu_instances", {}).items():
                parts.append(f"  - {gpu_key}: {gpu_info['type']} ({gpu_info['cost']})")
        for tip in provider.get("tips", []):
            parts.append(f"  - {tip}")
        parts.append("")
    return "\n".join(parts)
