import argparse
import sys
import os
from openharness.harnesses.bugfix.harness import BugFixHarness
from openharness.harnesses.bugfix.tasks import get_sample_tasks, create_sample_repo


def play_bugfix(args):
    """Run the bugfix harness in interactive mode"""
    print("OpenHarness BugFix Harness")
    print("=" * 30)

    # Get tasks
    tasks = get_sample_tasks()

    # Create harness with workspace
    harness = BugFixHarness(tasks=tasks, workspace_dir=args.workspace)

    # Reset to get first task
    observation = harness.reset()

    if not observation["task"]:
        print("No tasks available")
        return

    print(f"Task: {observation['task']['description']}")
    print("\nAvailable tools:")
    for tool_name in harness.tools.keys():
        print(f"  - {tool_name}")

    print("\nEnter commands in JSON format:")
    print('  {"tool": "tool_name", "input": {"param": "value"}}')
    print("Type 'quit' to exit\n")

    # Import json at the function level to avoid scope issues
    import json

    step_count = 0
    max_steps = args.max_steps if hasattr(args, "max_steps") else 10

    while step_count < max_steps:
        try:
            user_input = input("Action: ").strip()

            if user_input.lower() in ["quit", "exit", "q"]:
                print("Exiting...")
                break

            if not user_input:
                continue

            # Parse JSON input
            action = json.loads(user_input)

            # Execute action
            observation, reward, done, info = harness.step(action)

            # Print results
            if "error" in observation.data:
                print(f"Error: {observation.data['error']}")
            else:
                print(f"Output: {observation.data.get('output', 'No output')}")

            print(f"Reward: {reward}")

            if done:
                print("Task completed!")
                break

            step_count += 1

        except json.JSONDecodeError:
            print("Invalid JSON format. Please try again.")
        except KeyboardInterrupt:
            print("\nExiting...")
            break
        except EOFError:
            print("\nExiting...")
            break
        except Exception as e:
            print(f"Error: {e}")

    # End episode and show final results
    episode_data = harness.end_episode()
    print(f"\nEpisode finished!")
    print(f"Total reward: {episode_data['reward']}")
    print(f"Actions taken: {len(episode_data['actions'])}")

    if args.save_trajectory:
        traj_id = harness.trajectory_store.save_trajectory(episode_data)
        print(f"Trajectory saved with ID: {traj_id}")


def train_agent(args):
    """Train an agent (placeholder for future implementation)"""
    print("Training functionality not yet implemented")
    print("This would involve RL algorithms like PPO or GRPO")


def run_benchmark(args):
    """Run benchmark (placeholder for future implementation)"""
    print("Benchmark functionality not yet implemented")


def main():
    parser = argparse.ArgumentParser(
        description="OpenHarness - AI Agent Environment Framework"
    )
    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # Play command
    play_parser = subparsers.add_parser("play", help="Run a harness environment")
    play_parser.add_argument(
        "harness_type", choices=["bugfix"], help="Type of harness to run"
    )
    play_parser.add_argument(
        "--max-steps", type=int, default=10, help="Maximum steps per episode"
    )
    play_parser.add_argument(
        "--save-trajectory", action="store_true", help="Save trajectory after episode"
    )
    play_parser.add_argument(
        "--workspace", type=str, default=".", help="Workspace directory"
    )

    # Train command
    train_parser = subparsers.add_parser("train", help="Train an agent")
    train_parser.add_argument(
        "harness_type", choices=["bugfix", "coding"], help="Type of harness to train on"
    )
    train_parser.add_argument(
        "--rollouts", type=int, default=1000, help="Number of training rollouts"
    )

    # Benchmark command
    benchmark_parser = subparsers.add_parser("benchmark", help="Run benchmarks")
    benchmark_parser.add_argument("benchmark_name", help="Name of benchmark to run")

    # Sample repo command
    sample_parser = subparsers.add_parser(
        "sample", help="Create sample repository for testing"
    )

    args = parser.parse_args()

    if args.command == "play":
        if args.harness_type == "bugfix":
            play_bugfix(args)
        else:
            print(f"Harness type {args.harness_type} not yet implemented")

    elif args.command == "train":
        train_agent(args)

    elif args.command == "benchmark":
        run_benchmark(args)

    elif args.command == "sample":
        create_sample_repo()

    else:
        parser.print_help()


if __name__ == "__main__":
    main()
