import json
from datetime import datetime
from typing import Any, Dict, Optional

import typer
from flashenv.config import STATE_FILE


def load_state() -> Dict[str, Any]:
    if not STATE_FILE.exists():
        return {}

    try:
        content = STATE_FILE.read_text()
        if not content.strip():
            return {}
        return json.loads(content)
    except json.JSONDecodeError:
        typer.secho("⚠️ Warning: State file is corrupted. Resetting state.", fg=typer.colors.YELLOW)
        return {}


def save_state(data: Dict[str, Any]):
    STATE_FILE.parent.mkdir(parents=True, exist_ok=True)
    STATE_FILE.write_text(json.dumps(data, indent=4))


def add_env(name: str, instance_id: str, public_ip: str, image: str, spot: bool = False):
    data = load_state()
    if name in data:
        raise ValueError(f"Environment '{name}' already exists locally.")
    data[name] = {
        "id": instance_id,
        "ip": public_ip,
        "image": image,
        "created_at": datetime.now().isoformat(),
        "status": "running",
        "spot": spot,
    }
    save_state(data)


def remove_env(name: str) -> Optional[Dict]:
    data = load_state()
    if name not in data:
        return None
    removed_env = data.pop(name)
    save_state(data)
    return removed_env


def get_env(name: str) -> Optional[Dict]:
    data = load_state()
    return data.get(name)


def list_envs() -> Dict[str, Any]:
    return load_state()
