"""
CLI commands for Lindr personality classifier training.

Usage:
    lindr-train export-data -o data.json
    lindr-train train -d data.json -o model.onnx
    lindr-train push -m model.onnx -v v1.0.0
"""

from __future__ import annotations

import json
import logging
import os
import sys
from pathlib import Path

import click


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


@click.group()
@click.option("--verbose", "-v", is_flag=True, help="Enable verbose output")
def cli(verbose: bool) -> None:
    """Lindr personality classifier training CLI."""
    if verbose:
        logging.getLogger().setLevel(logging.DEBUG)


@cli.command("bootstrap")
@click.option(
    "--output", "-o",
    type=click.Path(),
    default="./bootstrap_data.json",
    help="Output file path",
)
@click.option(
    "--samples-per-trait",
    type=int,
    default=50,
    help="Samples per personality trait",
)
@click.option(
    "--include-embeddings/--no-embeddings",
    default=True,
    help="Pre-compute embeddings",
)
def bootstrap(output: str, samples_per_trait: int, include_embeddings: bool) -> None:
    """Generate bootstrap training data (no API needed)."""
    from .bootstrap import create_bootstrap_dataset

    click.echo(f"Generating bootstrap dataset with {samples_per_trait} samples per trait...")
    path = create_bootstrap_dataset(
        output_path=output,
        samples_per_trait=samples_per_trait,
        compute_embeddings=include_embeddings,
    )
    click.echo(f"Bootstrap dataset created: {path}")


@cli.command("export-data")
@click.option(
    "--api-key",
    envvar="LINDR_API_KEY",
    help="Lindr API key (or set LINDR_API_KEY env var)",
)
@click.option(
    "--output", "-o",
    type=click.Path(),
    default="./training_data.json",
    help="Output file path",
)
@click.option(
    "--base-url",
    default="https://api.lindr.io",
    help="Lindr API base URL",
)
@click.option(
    "--min-confidence",
    type=click.Choice(["low", "medium", "high"]),
    default="medium",
    help="Minimum confidence level for samples",
)
@click.option(
    "--include-embeddings/--no-embeddings",
    default=True,
    help="Include pre-computed embeddings",
)
@click.option(
    "--limit",
    type=int,
    default=10000,
    help="Maximum number of samples to export",
)
def export_data(
    api_key: str | None,
    output: str,
    base_url: str,
    min_confidence: str,
    include_embeddings: bool,
    limit: int,
) -> None:
    """Export training data from Lindr API."""
    if not api_key:
        click.echo("Error: API key required. Set LINDR_API_KEY or use --api-key", err=True)
        sys.exit(1)

    try:
        import httpx
    except ImportError:
        click.echo("Error: httpx is required. Install with: pip install httpx", err=True)
        sys.exit(1)

    click.echo(f"Exporting training data from {base_url}...")

    # Build request URL
    params = {
        "minConfidence": min_confidence,
        "includeEmbeddings": str(include_embeddings).lower(),
        "limit": str(limit),
        "format": "json",
    }

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    try:
        with httpx.Client(timeout=300) as client:
            response = client.get(
                f"{base_url}/api/v1/training-data/export",
                params=params,
                headers=headers,
            )
            response.raise_for_status()
            data = response.json()
    except httpx.HTTPStatusError as e:
        click.echo(f"Error: API request failed: {e.response.status_code}", err=True)
        click.echo(e.response.text, err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

    # Save to file
    output_path = Path(output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)

    click.echo(f"Exported {data.get('totalSamples', 0)} samples to {output_path}")


@cli.command("train")
@click.option(
    "--data", "-d",
    type=click.Path(exists=True),
    required=True,
    help="Path to training data JSON file",
)
@click.option(
    "--external",
    multiple=True,
    type=click.Choice(["big5-chat", "pandora"]),
    help="Include external datasets",
)
@click.option(
    "--output", "-o",
    type=click.Path(),
    default="./model.onnx",
    help="Output ONNX model path",
)
@click.option("--epochs", default=50, type=int, help="Number of training epochs")
@click.option("--batch-size", default=64, type=int, help="Training batch size")
@click.option("--learning-rate", default=1e-3, type=float, help="Learning rate")
@click.option("--val-split", default=0.1, type=float, help="Validation split ratio")
@click.option(
    "--checkpoint-dir",
    default="./checkpoints",
    help="Directory for saving checkpoints",
)
@click.option("--seed", default=42, type=int, help="Random seed")
def train(
    data: str,
    external: tuple[str, ...],
    output: str,
    epochs: int,
    batch_size: int,
    learning_rate: float,
    val_split: float,
    checkpoint_dir: str,
    seed: int,
) -> None:
    """Train personality classifier."""
    try:
        import torch
    except ImportError:
        click.echo(
            "Error: PyTorch is required for training. "
            "Install with: pip install torch",
            err=True,
        )
        sys.exit(1)

    from .model import PersonalityClassifier
    from .dataset import PersonalityDataset, CombinedDataset
    from .trainer import PersonalityTrainer, TrainingConfig
    from .exporter import export_to_onnx

    click.echo(f"Loading training data from {data}...")
    dataset = PersonalityDataset(data, precomputed_embeddings=True)
    click.echo(f"  Loaded {len(dataset)} samples")

    # Add external datasets if specified
    datasets = [dataset]
    if external:
        from .external import load_external_dataset

        for ext_name in external:
            click.echo(f"Loading external dataset: {ext_name}...")
            try:
                ext_dataset = load_external_dataset(ext_name, max_samples=10000)
                datasets.append(ext_dataset)
                click.echo(f"  Loaded {len(ext_dataset)} samples from {ext_name}")
            except Exception as e:
                click.echo(f"  Warning: Failed to load {ext_name}: {e}", err=True)

    # Combine datasets
    if len(datasets) > 1:
        combined = CombinedDataset(datasets)
        click.echo(f"Total training samples: {len(combined)}")
    else:
        combined = dataset

    # Create model
    click.echo("Creating model...")
    model = PersonalityClassifier()

    # Configure training
    config = TrainingConfig(
        epochs=epochs,
        batch_size=batch_size,
        learning_rate=learning_rate,
        val_split=val_split,
        checkpoint_dir=checkpoint_dir,
        seed=seed,
    )

    # Train
    click.echo(f"Training for {epochs} epochs...")
    trainer = PersonalityTrainer(model, combined, config=config)

    def progress_callback(epoch: int, total: int, loss: float) -> None:
        pass  # Logging handled by trainer

    result = trainer.train(progress_callback=progress_callback)

    click.echo(f"\nTraining complete!")
    click.echo(f"  Best validation loss: {result.best_val_loss:.4f}")
    click.echo(f"  Best epoch: {result.best_epoch}")
    click.echo(f"  Final epoch: {result.final_epoch}")

    # Export to ONNX
    click.echo(f"\nExporting model to {output}...")
    export_to_onnx(model, output, verify=True)

    click.echo(f"\nModel saved to {output}")


@cli.command("push")
@click.option(
    "--model", "-m",
    type=click.Path(exists=True),
    required=True,
    help="Path to ONNX model file",
)
@click.option(
    "--version", "-v",
    required=True,
    help="Model version (e.g., v1.0.0)",
)
@click.option(
    "--hf-token",
    envvar="HF_TOKEN",
    help="HuggingFace API token (or set HF_TOKEN env var)",
)
@click.option(
    "--repo-id",
    default="lindr-ai/personality-classifier",
    help="HuggingFace repository ID",
)
@click.option(
    "--private/--public",
    default=True,
    help="Repository visibility",
)
def push(
    model: str,
    version: str,
    hf_token: str | None,
    repo_id: str,
    private: bool,
) -> None:
    """Push trained model to HuggingFace Hub."""
    if not hf_token:
        click.echo(
            "Error: HuggingFace token required. Set HF_TOKEN or use --hf-token",
            err=True,
        )
        sys.exit(1)

    from .hub import push_to_hub

    click.echo(f"Pushing model to {repo_id}...")

    try:
        url = push_to_hub(
            onnx_path=model,
            version=version,
            repo_id=repo_id,
            token=hf_token,
            private=private,
        )
        click.echo(f"\nModel pushed successfully!")
        click.echo(f"URL: {url}")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)


@cli.command("download")
@click.option(
    "--version", "-v",
    default="latest",
    help="Model version to download",
)
@click.option(
    "--output", "-o",
    type=click.Path(),
    default="./model.onnx",
    help="Output path for the model",
)
@click.option(
    "--hf-token",
    envvar="HF_TOKEN",
    help="HuggingFace API token (for private repos)",
)
@click.option(
    "--repo-id",
    default="lindr-ai/personality-classifier",
    help="HuggingFace repository ID",
)
def download(
    version: str,
    output: str,
    hf_token: str | None,
    repo_id: str,
) -> None:
    """Download model from HuggingFace Hub."""
    from .hub import download_model
    import shutil

    click.echo(f"Downloading model version {version}...")

    try:
        model_path = download_model(
            version=version,
            repo_id=repo_id,
            token=hf_token,
        )

        # Copy to output path
        output_path = Path(output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(model_path, output_path)

        click.echo(f"Model downloaded to {output_path}")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)


@cli.command("info")
@click.option(
    "--model", "-m",
    type=click.Path(exists=True),
    required=True,
    help="Path to ONNX model file",
)
def info(model: str) -> None:
    """Show information about an ONNX model."""
    from .exporter import get_onnx_model_info

    try:
        model_info = get_onnx_model_info(model)

        click.echo(f"\nModel: {model_info['path']}")
        click.echo(f"Size: {model_info['size_kb']:.1f} KB")
        click.echo(f"ONNX opset: {model_info['opset_version']}")
        click.echo(f"IR version: {model_info['ir_version']}")
        click.echo(f"Nodes: {model_info['num_nodes']}")

        click.echo("\nInputs:")
        for inp in model_info["inputs"]:
            click.echo(f"  {inp['name']}: {inp['shape']} ({inp['dtype']})")

        click.echo("\nOutputs:")
        for out in model_info["outputs"]:
            click.echo(f"  {out['name']}: {out['shape']} ({out['dtype']})")

    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)


@cli.command("verify")
@click.option(
    "--model", "-m",
    type=click.Path(exists=True),
    required=True,
    help="Path to ONNX model file",
)
def verify(model: str) -> None:
    """Verify an ONNX model can be loaded and run."""
    try:
        import onnxruntime as ort
        import numpy as np
    except ImportError:
        click.echo(
            "Error: onnxruntime and numpy required. "
            "Install with: pip install onnxruntime numpy",
            err=True,
        )
        sys.exit(1)

    click.echo(f"Verifying model: {model}")

    try:
        # Load model
        session = ort.InferenceSession(model, providers=["CPUExecutionProvider"])
        click.echo("  Model loaded successfully")

        # Get input info
        input_info = session.get_inputs()[0]
        click.echo(f"  Input: {input_info.name} {input_info.shape}")

        # Run inference
        test_input = np.random.randn(1, 384).astype(np.float32)
        output = session.run(None, {"embedding": test_input})

        click.echo(f"  Output shape: {output[0].shape}")
        click.echo(f"  Output range: [{output[0].min():.3f}, {output[0].max():.3f}]")
        click.echo("\nModel verification passed!")

    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)


def main() -> None:
    """Entry point for the CLI."""
    cli()


if __name__ == "__main__":
    main()
