import os
import sys
import subprocess
import requests
import typer
import yaml

app = typer.Typer(help="Cifer CLI", rich_markup_mode="rich")

securetrain_app = typer.Typer(help="Secure training commands")
dataset_app = typer.Typer(help="Dataset commands")

app.add_typer(securetrain_app, name="securetrain")
app.add_typer(dataset_app, name="dataset")


# ================= CONFIG =================
def load_config():
    config_path = os.path.expanduser("~/.cifer/config.yaml")

    if os.path.exists(config_path):
        with open(config_path) as f:
            return yaml.safe_load(f)

    return {}

# ================= VERSION =================
@app.command()
def version():
    """Show Cifer version"""
    from importlib.metadata import version
    typer.echo(version("cifer"))


# ================= PROJECT INIT =================
@app.command("init")
def init_project(name: str):
    """Create new Cifer project"""

    os.makedirs(name, exist_ok=True)
    os.makedirs(f"{name}/dataset", exist_ok=True)
    os.makedirs(f"{name}/models", exist_ok=True)

    config = {
        "project": name,
        "server": "ws://localhost:8080"
    }

    os.makedirs(os.path.expanduser("~/.cifer"), exist_ok=True)

    with open(f"{name}/config.yaml", "w") as f:
        yaml.dump(config, f)

    typer.echo(f"✅ Cifer project '{name}' created")


# ================= SERVER =================
@app.command("server")
def run_server(
    host: str = typer.Option("0.0.0.0"),
    port: int = typer.Option(8080),
):
    """Run Cifer Federated Server"""

    typer.echo(f"🚀 Starting Cifer server at {host}:{port}")

    from cifer.server import CiferServer

    server = CiferServer(host=host, port=port)
    server.start()


# ================= CLIENT =================
@app.command("client")
def run_client(
    dataset: str = typer.Option(...),
):
    """Run federated client"""

    config = load_config()
    server = config.get("server", "ws://localhost:8080")

    typer.echo(f"🔗 Connecting to {server}")

    from cifer.client import CiferClient

    client = CiferClient(server_address=server)
    client.train(dataset)


# ================= DATASET =================
@dataset_app.command("download")
def download_dataset(name: str):
    """Download dataset from HuggingFace"""

    typer.echo(f"📦 Downloading dataset: {name}")

    from datasets import load_dataset

    dataset = load_dataset(name)

    output_dir = f"./dataset/{name}"
    os.makedirs("./dataset", exist_ok=True)

    dataset.save_to_disk(output_dir)

    typer.echo(f"✅ Dataset saved to {output_dir}")


# ================= STATUS =================
@app.command("status")
def status(server: str = typer.Option("http://localhost:8080")):
    """Check server status"""

    try:
        r = requests.get(f"{server}/metrics", timeout=3)
        typer.echo(r.json())

    except Exception:
        typer.echo("❌ Server not reachable")


# ================= Securetrain =================
@securetrain_app.command("encrypt-dataset")
def securetrain_encrypt(
    dataset: str = typer.Option(...),
    output: str = typer.Option(...),
):
    """Encrypt dataset"""

    from cifer.securetrain import encrypt_dataset

    encrypt_dataset(dataset, output)


@securetrain_app.command("train")
def securetrain_train(
    encrypted_data: str = typer.Option(...),
    output_model: str = typer.Option(...),
):
    """Train model"""

    from cifer.securetrain import train_model

    train_model(encrypted_data, output_model)


@securetrain_app.command("decrypt-model")
def securetrain_decrypt(
    input_model: str = typer.Option(...),
    output_model: str = typer.Option(...),
):
    """Decrypt trained model"""

    from cifer.securetrain import decrypt_model

    decrypt_model(input_model, output_model)


# ================= JUPYTER KERNEL =================
@app.command("register-kernel")
def register_kernel():
    """Register Jupyter Kernel"""

    kernel_name = "cifer-kernel"
    display_name = "Cifer Kernel"

    typer.echo("📓 Registering Jupyter kernel...")

    subprocess.run(
        [
            sys.executable,
            "-m",
            "ipykernel",
            "install",
            "--user",
            "--name",
            kernel_name,
            "--display-name",
            display_name,
        ],
        check=True,
    )

    typer.echo("✅ Kernel registered")


# ================= AGENT =================
@app.command("agent-ace")
def agent_ace(
    port: int = typer.Option(9999),
):
    """Launch ACE agent"""

    from cifer.agent_ace import run_agent_ace

    run_agent_ace(port=port)


# ================= NOTEBOOK =================
@app.command("download-notebook")
def download_notebook(
    url: str = typer.Option(...),
):
    """Download notebook"""

    filename = url.split("/")[-1]

    try:

        typer.echo(f"⬇️ Downloading {filename}")

        response = requests.get(url, stream=True)

        if response.status_code == 200:

            with open(filename, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)

            typer.echo(f"✅ Notebook downloaded: {filename}")

        else:
            typer.echo("❌ Download failed")

    except Exception as e:
        typer.echo(f"❌ Error: {e}")


# ================= TRAINING SIM =================
@app.command("train")
def simulate_training(
    epochs: int = typer.Option(5),
    lr: float = typer.Option(0.001),
):
    """Simulate training"""

    typer.echo(f"⚙️ Training for {epochs} epochs")

    for i in range(1, epochs + 1):
        typer.echo(f"Epoch {i}/{epochs}")

    typer.echo("✅ Training completed")


# ================= ENTRYPOINT =================
def main():
    app()


if __name__ == "__main__":
    main()