"""
CLI Integration Tests
=====================

Tests that verify the AutoTrain CLI actually works end-to-end by running real commands.
"""

import json
import os
import subprocess
import tempfile
from pathlib import Path

import pytest


def run_cli_command(command: str, timeout: int = 120):
    """Run a CLI command and return output."""
    import shlex

    result = subprocess.run(
        shlex.split(command),
        capture_output=True,
        text=True,
        timeout=timeout,
        cwd=Path(__file__).parent.parent,
        env={**os.environ, "PYTHONPATH": "src"},
    )
    return result


@pytest.fixture
def temp_data_dir():
    """Create temporary directory with sample data."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Create sample training data
        data_file = os.path.join(tmpdir, "train.jsonl")
        with open(data_file, "w") as f:
            for i in range(10):
                f.write(json.dumps({"text": f"Sample text {i}"}) + "\n")

        yield tmpdir


@pytest.fixture
def temp_prompts_file():
    """Create temporary prompts file for inference."""
    with tempfile.TemporaryDirectory() as tmpdir:
        prompts_file = os.path.join(tmpdir, "prompts.txt")
        with open(prompts_file, "w") as f:
            f.write("What is AI?\n")
            f.write("Explain machine learning\n")

        yield prompts_file


class TestCLIHelp:
    """Test CLI help commands."""

    def test_main_help(self):
        """Test main help command."""
        result = run_cli_command("python -m autotrain.cli.autotrain --help")
        assert result.returncode == 0 or "autotrain" in result.stdout.lower()

    def test_llm_help(self):
        """Test LLM help command."""
        result = run_cli_command("python -m autotrain.cli.autotrain llm --help")
        assert result.returncode == 0 or "llm" in result.stdout.lower()


class TestCLIParameterValidation:
    """Test CLI parameter parsing."""

    def test_parameters_loaded(self):
        """Test that all new parameters are recognized."""
        # This tests parameter loading, not execution
        cmd = """python -c "
import sys
sys.path.insert(0, 'src')
from autotrain.trainers.clm.params import LLMTrainingParams
p = LLMTrainingParams()
# Check Tinker params exist
assert hasattr(p, 'use_distillation')
assert hasattr(p, 'use_sweep')
assert hasattr(p, 'use_enhanced_eval')
assert hasattr(p, 'rl_gamma')
assert hasattr(p, 'custom_loss')
print('✅ All parameters loaded')
        "
        """
        result = run_cli_command(cmd)
        assert result.returncode == 0
        assert "✅ All parameters loaded" in result.stdout


class TestCLIInference:
    """Test inference mode CLI."""

    def test_inference_with_prompts_file(self, temp_prompts_file):
        """Test inference with prompts from file."""
        with tempfile.TemporaryDirectory() as tmpdir:
            output_file = os.path.join(tmpdir, "results.json")

            # Note: This will fail without a model, but tests CLI parsing
            result = run_cli_command(
                f"python -m autotrain.cli.autotrain llm "
                f"--inference "
                f"--model gpt2 "
                f"--inference-prompts {temp_prompts_file} "
                f"--inference-output {output_file} "
                f"--project-name test-inference"
            )

            # Check if command was parsed correctly (may fail on model load, but that's ok for this test)
            assert "--inference-prompts" not in result.stderr or "error" not in result.stderr.lower()

    def test_inference_with_inline_prompts(self):
        """Test inference with inline prompts."""
        with tempfile.TemporaryDirectory() as tmpdir:
            output_file = os.path.join(tmpdir, "results.json")

            result = run_cli_command(
                f"python -m autotrain.cli.autotrain llm "
                f"--inference "
                f"--model gpt2 "
                f'--inference-prompts "What is AI?,Explain ML" '
                f"--inference-output {output_file} "
                f"--project-name test-inference"
            )

            # Check CLI parsing
            assert "inference" in result.stdout.lower() or result.returncode != 2


class TestCLITraining:
    """Test training CLI commands."""

    def test_sweep_parameters_recognized(self, temp_data_dir):
        """Test that sweep parameters are recognized."""
        # Test parameter recognition (not actual execution)
        result = run_cli_command(
            f"python -m autotrain.cli.autotrain llm "
            f"--train "
            f"--model gpt2 "
            f"--data-path {temp_data_dir} "
            f"--project-name test-sweep "
            f"--use-sweep "
            f"--sweep-backend optuna "
            f"--sweep-n-trials 2 "
            f"--epochs 1 "
            f"--batch-size 1 "
            f"--help"  # Add help to just test parsing
        )

        # Should not error on parameter parsing
        assert "unrecognized arguments" not in result.stderr

    def test_enhanced_eval_parameters_recognized(self, temp_data_dir):
        """Test that enhanced eval parameters are recognized."""
        result = run_cli_command(
            f"python -m autotrain.cli.autotrain llm "
            f"--train "
            f"--model gpt2 "
            f"--data-path {temp_data_dir} "
            f"--project-name test-eval "
            f"--use-enhanced-eval "
            f"--eval-metrics perplexity,bleu "
            f"--eval-save-predictions "
            f"--help"
        )

        assert "unrecognized arguments" not in result.stderr

    def test_distillation_parameters_recognized(self, temp_data_dir):
        """Test distillation CLI parameters."""
        result = run_cli_command(
            f"python -m autotrain.cli.autotrain llm "
            f"--train "
            f"--trainer distillation "
            f"--model gpt2 "
            f"--teacher-model gpt2-medium "
            f"--teacher-prompt-template 'Answer: {{input}}' "
            f"--data-path {temp_data_dir} "
            f"--project-name test-distill "
            f"--help"
        )

        assert "unrecognized arguments" not in result.stderr

    def test_ppo_parameters_recognized(self, temp_data_dir):
        """Test PPO RL parameters."""
        result = run_cli_command(
            f"python -m autotrain.cli.autotrain llm "
            f"--train "
            f"--trainer ppo "
            f"--model gpt2 "
            f"--rl-gamma 0.99 "
            f"--rl-kl-coef 0.1 "
            f"--rl-env-type text_generation "
            f"--data-path {temp_data_dir} "
            f"--project-name test-ppo "
            f"--help"
        )

        assert "unrecognized arguments" not in result.stderr

    def test_custom_loss_parameters_recognized(self, temp_data_dir):
        """Test custom loss parameters."""
        result = run_cli_command(
            f"python -m autotrain.cli.autotrain llm "
            f"--train "
            f"--model gpt2 "
            f"--custom-loss composite "
            f'--custom-loss-weights "[1.0, 0.5]" '
            f"--data-path {temp_data_dir} "
            f"--project-name test-loss "
            f"--help"
        )

        assert "unrecognized arguments" not in result.stderr


class TestCLITrainerOptions:
    """Test all trainer options are available."""

    def test_all_trainers_listed(self):
        """Verify all trainers are recognized."""
        trainers = ["default", "sft", "reward", "dpo", "orpo", "distillation", "ppo"]

        for trainer in trainers:
            # Just test that the trainer option doesn't cause an error
            result = run_cli_command(
                f'python -c "from autotrain.trainers.clm.__main__ import train; '
                f"from autotrain.trainers.clm.params import LLMTrainingParams; "
                f"config = LLMTrainingParams(trainer='{trainer}'); "
                f"print('Trainer {trainer} recognized')\""
            )

            if result.returncode == 0:
                assert f"Trainer {trainer} recognized" in result.stdout


class TestCLIEndToEnd:
    """End-to-end CLI tests (these may take longer)."""

    @pytest.mark.slow
    def test_inference_end_to_end(self):
        """Test actual inference execution."""
        with tempfile.TemporaryDirectory() as tmpdir:
            prompts_file = os.path.join(tmpdir, "prompts.txt")
            with open(prompts_file, "w") as f:
                f.write("Hello\n")

            output_file = os.path.join(tmpdir, "results.json")

            result = run_cli_command(
                f"python -m autotrain.cli.autotrain llm "
                f"--inference "
                f"--model gpt2 "
                f"--inference-prompts {prompts_file} "
                f"--inference-output {output_file} "
                f"--inference-max-tokens 10 "
                f"--project-name test-inf",
                timeout=300,
            )

            # Check if ran (may fail on GPU/model issues in CI, but CLI works)
            if result.returncode == 0:
                assert os.path.exists(output_file)
                with open(output_file) as f:
                    results = json.load(f)
                    assert len(results) > 0


class TestWandbProjectPathHandling:
    """Test WANDB_PROJECT handles path-like project names correctly."""

    def test_wandb_project_uses_basename_not_full_path(self):
        """Test that WANDB_PROJECT is set to basename when project_name is a path.

        This tests the fix for the bug where using a path like
        '/workspace/trainings/hotel-sft' as project_name would cause W&B errors
        because paths contain invalid characters for W&B project names.
        """
        result = run_cli_command(
            """python -c "
import os
import sys
sys.path.insert(0, 'src')

from autotrain.trainers.clm.params import LLMTrainingParams

# Simulate CLI with path as project_name
config = LLMTrainingParams(
    model='gpt2',
    project_name='/workspace/trainings/hotel-sft-optuna-v2',
    log='wandb',
)

# Test the basename extraction logic
wandb_project = getattr(config, 'wandb_sweep_project', None) or os.path.basename(config.project_name)

# Verify
assert wandb_project == 'hotel-sft-optuna-v2', f'Expected hotel-sft-optuna-v2, got {wandb_project}'

# Verify no invalid W&B characters
invalid_chars = '/\\\\#?%:'
has_invalid = any(c in wandb_project for c in invalid_chars)
assert not has_invalid, f'wandb_project contains invalid characters: {wandb_project}'

print('WANDB_PROJECT basename test PASSED')
" """
        )
        assert result.returncode == 0, f"Test failed: {result.stderr}"
        assert "PASSED" in result.stdout

    def test_explicit_wandb_sweep_project_takes_precedence(self):
        """Test that explicit wandb_sweep_project is used over basename."""
        result = run_cli_command(
            """python -c "
import os
import sys
sys.path.insert(0, 'src')

from autotrain.trainers.clm.params import LLMTrainingParams

config = LLMTrainingParams(
    model='gpt2',
    project_name='/workspace/trainings/hotel-sft-optuna-v2',
    log='wandb',
    wandb_sweep_project='my-explicit-project',
)

wandb_project = getattr(config, 'wandb_sweep_project', None) or os.path.basename(config.project_name)

assert wandb_project == 'my-explicit-project', f'Expected my-explicit-project, got {wandb_project}'
print('Explicit wandb_sweep_project test PASSED')
" """
        )
        assert result.returncode == 0, f"Test failed: {result.stderr}"
        assert "PASSED" in result.stdout


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
