#!/usr/bin/env python
import os
import sys
import subprocess
import click
import logging
import hashlib
import shutil


# Add the parent directory of 'neural' to sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))


from neural.parser.parser import create_parser, ModelTransformer
from neural.code_generation.code_generator import generate_code

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@click.group()
def cli():
    """Neural CLI: A compiler-like interface for .neural and .nr files."""
    pass

@cli.command()
@click.argument('file', type=click.Path(exists=True))
@click.option('--backend', default='tensorflow', help='Target backend: tensorflow or pytorch', type=click.Choice(['tensorflow', 'pytorch']))
@click.option('--verbose', is_flag=True, help='Show verbose output')
@click.option('--output', default=None, help='Output file path for generated code or visualizations')
def compile(file, backend, verbose, output):
    """
    Compile a .neural or .nr file into an executable Python script.
    
    Example:
        neural compile my_model.neural --backend pytorch
    """
    if verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    ext = os.path.splitext(file)[1].lower()
    # Choose the appropriate start rule based on extension
    if ext in ['.neural', '.nr']:
        parser_instance = create_parser('network')
    elif ext == '.rnr':
        parser_instance = create_parser('research')
    else:
        click.echo(f"Unsupported file type: {ext}")
        sys.exit(1)

    with open(file, 'r') as f:
        content = f.read()
    
    try:
        tree = parser_instance.parse(content)
    except Exception as e:
        click.echo(f"Error parsing {file}: {e}")
        sys.exit(1)

    transformer = ModelTransformer()
    try:
        model_data = transformer.transform(tree)
    except Exception as e:
        click.echo(f"Error transforming {file}: {e}")
        sys.exit(1)

    code = generate_code(model_data, backend)
    output_file = output or os.path.splitext(file)[0] + f"_{backend}.py"
    with open(output_file, 'w') as f:
        f.write(code)
    click.echo(f"Compiled {file} to {output_file} for backend {backend}")

@cli.command()
@click.argument('file', type=click.Path(exists=True))
def run(file):
    """
    Run an executable neural model.
    
    Example:
        neural run my_model_pytorch.py
    """
    # Here, we assume that the file is a Python file generated by the compile command.
    subprocess.run([sys.executable, file], check=True)


@cli.command()
@click.argument('file', type=click.Path(exists=True))
@click.option('--format', default='html', help='Output format: html or png')
def visualize(file, format):
    """
    Visualize network architecture and shape propagation.
    
    Example:
        neural visualize my_model.neural --format png
    """
    from neural.shape_propagation.shape_propagator import ShapePropagator
    from neural.parser.parser import create_parser
    from neural.dashboard.tensor_flow import create_animated_network

    def get_file_hash(file):
        with open(file, 'rb') as f:
            return hashlib.sha256(f.read()).hexdigest()

        cache_dir = ".neural_cache"
        os.makedirs(cache_dir, exist_ok=True)
        file_hash = get_file_hash(file)
        cache_file = os.path.join(cache_dir, f"viz_{file_hash}_{format}")

        if os.path.exists(cache_file):
            click.echo(f"Using cached visualization: {cache_file}")
        else:
            # Parse input file
            ext = os.path.splitext(file)[1].lower()
            parser_instance = create_parser('network' if ext in ['.neural', '.nr'] else 'research')
            
            with open(file, 'r') as f:
                content = f.read()
            
            try:
                tree = parser_instance.parse(content)
                model_data = ModelTransformer().transform(tree)
            except Exception as e:
                click.echo(f"Error processing {file}: {e}")
                sys.exit(1)

            # Propagate shapes
            propagator = ShapePropagator()
            input_shape = model_data['input']['shape']  # Get from parsed model
            if not input_shape:
                click.echo("Error: Input shape not defined in model!")
                sys.exit(1) 
            shape_history = []

            with click.progressbar(length=len(model_data['layers']), label="Propagating shapes") as bar:
                for layer in model_data['layers']:
                    input_shape = propagator.propagate(input_shape, layer, model_data['framework'])
                    shape_history.append({"layer": layer['type'], "output_shape": input_shape})
                    bar.update(1)

            # Generate visualizations
            report = propagator.generate_report()
            
            # Save Graphviz diagram
            dot = report['dot_graph']
            dot.format = 'png' if format == 'png' else 'svg'
            dot.render('architecture', cleanup=True)
            
            # Save Plotly visualization
            fig = report['plotly_chart']
            fig.write_html('shape_propagation.html')
            
            # Generate tensor flow animation
            tf_fig = create_animated_network(shape_history)
            tf_fig.write_html('tensor_flow.html')

            click.echo(f"""
            Visualizations generated:
            - architecture.{format} (Network architecture)
            - shape_propagation.html (Parameter count chart)
            - tensor_flow.html (Data flow animation)
            """)

@cli.command()
def clean():
    """Remove generated files (e.g., .py, .png, .html, cache)."""
    for ext in ['.py', '.png', '.html']:
        for file in os.listdir('.'):
            if file.endswith(ext):
                os.remove(file)
                click.echo(f"Removed {file}")
    if os.path.exists(".neural_cache"):
        shutil.rmtree(".neural_cache")
        click.echo("Removed cache directory")


#######################
### Version Command ###
#######################

@cli.command()
def version():
    """Show the version of Neural CLI and dependencies."""
    click.echo(f"Neural CLI v0.1")
    click.echo(f"Python: {sys.version}")
    click.echo(f"Click: {click.__version__}")

########################################
### Dashboard Integration Or Command ###
########################################


@cli.command()
@click.argument('file', type=click.Path(exists=True))
def dashboard(file):
    """Launch the Neural dashboard for real-time monitoring of the model."""
    from neural.dashboard.dashboard import app, server, socketio
    from neural.shape_propagation.shape_propagator import ShapePropagator
    from neural.parser.parser import create_parser

    # Parse and propagate shapes to populate trace_data
    parser_instance = create_parser('network' if os.path.splitext(file)[1].lower() in ['.neural', '.nr'] else 'research')
    with open(file, 'r') as f:
        content = f.read()
    tree = parser_instance.parse(content)
    model_data = ModelTransformer().transform(tree)
    propagator = ShapePropagator()
    input_shape = model_data['input']['shape']
    for layer in model_data['layers']:
        input_shape = propagator.propagate(input_shape, layer, model_data['framework'])

    # Update trace_data for the dashboard
    import threading
    global trace_data
    trace_data = propagator.get_trace()
    threading.Thread(target=socketio.run, args=(server, "localhost", 5001), daemon=True).start()
    app.run_server(debug=True, port=8050)

#### --- Debugger --- ####

@cli.command()
@click.argument('file', type=click.Path(exists=True))
@click.option('--gradients', is_flag=True, help='Analyze gradient flow')
@click.option('--dead-neurons', is_flag=True, help='Detect dead neurons')
@click.option('--anomalies', is_flag=True, help='Detect training anomalies')
@click.option('--step', is_flag=True, help='Enable step debugging')
def debug(file, gradients, dead_neurons, anomalies, step):
    """Debug a neural network model with NeuralDbg."""
    from neural.shape_propagation.shape_propagator import ShapePropagator, register_gradient_hooks, detect_dead_neurons, detect_activation_anomalies, step_debug_hook
    from neural.parser.parser import create_parser

    parser_instance = create_parser('network' if os.path.splitext(file)[1].lower() in ['.neural', '.nr'] else 'research')
    with open(file, 'r') as f:
        content = f.read()
    tree = parser_instance.parse(content)
    model_data = ModelTransformer().transform(tree)
    
    propagator = ShapePropagator(debug=True)
    input_shape = model_data['input']['shape']
    for layer in model_data['layers']:
        input_shape = propagator.propagate(input_shape, layer, model_data['framework'])

    trace_data = propagator.get_trace()
    if gradients:
        model = ...  # Load or create the PyTorch/TensorFlow model from model_data
        gradient_trace = register_gradient_hooks(model)
        click.echo("Gradient flow analysis:")
        for entry in gradient_trace:
            click.echo(f"Layer {entry['layer']}: grad_norm = {entry['grad_norm']}")
    if dead_neurons:
        # Simulate or integrate with model to detect dead neurons
        click.echo("Dead neuron detection:")
        for entry in trace_data:
            if 'dead_ratio' in entry:
                click.echo(f"Layer {entry['layer']}: dead_ratio = {entry['dead_ratio']}")
    if anomalies:
        click.echo("Anomaly detection:")
        for entry in trace_data:
            if 'anomaly' in entry:
                click.echo(f"Layer {entry['layer']}: anomaly = {entry['anomaly']}, mean_activation = {entry['mean_activation']}")
    if step:
        # Implement step debugging (e.g., using hooks or breakpoints)
        click.echo("Step debugging mode: Pausing at each layer. Press Enter to continue...")
        for layer in model_data['layers']:
            input_shape = propagator.propagate(input_shape, layer, model_data['framework'])
            click.echo(f"Paused at {layer['type']}: input_shape = {input_shape}")
            suggestion = lm.suggest(f"network DebugNet {{ layers: {layer['type']}")
            click.echo(f"LM Suggestion: {suggestion}")
            input("Press Enter to continue...")


######################
### Export Command ###
######################
@cli.command()
@click.argument('file', type=click.Path(exists=True))
@click.option('--format', default='onnx', help='Export format: onnx')
@click.option('--output', default='model.onnx', help='Output file path')
def export(file, format, output):
    """Export a neural network model to ONNX or other formats."""
    from neural.parser.parser import create_parser
    from neural.code_generation.code_generator import generate_code
    parser_instance = create_parser('network' if os.path.splitext(file)[1].lower() in ['.neural', '.nr'] else 'research')
    with open(file, 'r') as f:
        content = f.read()
    tree = parser_instance.parse(content)
    model_data = ModelTransformer().transform(tree)
    generate_code(model_data, format)
    click.echo(f"Exported {file} to {output}")


###############################
### Hacky Mode for Debugger ###
###############################

@cli.command()
@click.argument('file', type=click.Path(exists=True))
@click.option('--hacky', is_flag=True, help='Run NeuralDbg in hacky mode for security analysis')
def debug(file, hacky):
  """Debug a neural network model with NeuralDbg."""
  # ... Existing debug code ...
  if hacky:
      from neural.hacky import hacky
      click.echo("Running NeuralDbg in hacky mode...")
      hacky_mode(propagator, model_data)
  else:
      # Normal debug mode
      click.echo("Running NeuralDbg in normal mode...")

######################
### Train Command ####
######################
@cli.command()
@click.argument('file', type=click.Path(exists=True))
@click.option('--backend', default='tensorflow', help='Target backend: tensorflow or pytorch')
@click.option('--log-dir', default='runs/neural', help='TensorBoard log directory')
def train(file, backend, log_dir):
    """Train a neural network model and log to TensorBoard."""
    from neural.parser.parser import create_parser
    from neural.code_generation.code_generator import generate_code
    parser_instance = create_parser('network' if os.path.splitext(file)[1].lower() in ['.neural', '.nr'] else 'research')
    with open(file, 'r') as f:
        content = f.read()
    tree = parser_instance.parse(content)
    model_data = ModelTransformer().transform(tree)
    code = generate_code(model_data, backend)
    output_file = f"model_{backend}.py"
    save_file(output_file, code)
    
    subprocess.run([sys.executable, output_file], check=True)
    click.echo(f"Training completed. Logs available in {log_dir}")


#############################
### Load Pretrained Models ##
#############################

@cli.command()
@click.argument('model_name', default='resnet50')
@click.option('--pretrained', is_flag=True, help='Load pretrained weights from Hugging Face')
@click.option('--output', default='model.pth', help='Output file for saved model')
def load(model_name, pretrained, output):
    """Load a pretrained model (e.g., ResNet50) from the hub."""
    from neural.pretrained import pretrained
    try:
        hub = PretrainedModelHub()
        model = hub.load(model_name, pretrained=pretrained)
        torch.save(model, output)
        click.echo(f"Loaded {model_name} and saved to {output}")
    except Exception as e:
        click.echo(f"Error loading model: {e}")
        sys.exit(1)


### ------ Chat ------- ####
@cli.command()
def chat():
    """Interact with NeuralChat to build models conversationally."""
    from neural.neural_chat import neural_chat
    chat = NeuralChat()
    click.echo("Welcome to NeuralChat! Type commands or 'exit' to quit.")
    while True:
        command = click.prompt("> ", type=str)
        if command.lower() == "exit":
            break
        response = chat.process_command(command)
        click.echo(response)



### ----- llm ------ ###
@cli.command()
@click.option('--prompt', default="network MyNet {", help='Input prompt for autocompletion')
def lm_suggest(prompt):
    """Get suggestions from Neurallm for .neural syntax."""
    from neural.neurallm import neurallm
    lm = Neurallm(model_path="./neurallm")
    suggestion = lm.suggest(prompt)
    click.echo(f"Suggestion: {suggestion}")

@cli.command()
@click.argument('dataset', type=click.Path(exists=True))
@click.option('--epochs', default=1, help='Number of training epochs')
def lm_train(dataset, epochs):
    """Fine-tune Neurallm on a dataset of .neural files."""
    from neural.neurallm import neurallm
    lm = Neurallm()
    lm.train(dataset, epochs)

@click.group()
def cli():  # <-- This must be named 'cli'
    pass

# Add commands
cli.add_command(compile)
cli.add_command(run)
# ... other commands

if __name__ == '__main__':
    cli()
