import argparse
import os
import sys
import webbrowser
from tn_lens.circuit import QASMParser
from tn_lens.analyzer import analyze_tensors
from tn_lens.renderer import render_graph

def main():
    parser = argparse.ArgumentParser(
        description="tn-lens: A Tensor Network Visualizer for Quantum Circuits."
    )
    
    parser.add_argument("input_file", type=str, help="Path to the input OpenQASM file (.qasm)")
    parser.add_argument("-o", "--output", type=str, default="tensor_network.html", help="Output HTML file path (default: tensor_network.html)")
    parser.add_argument("-m", "--memory", type=float, default=500.0, help="Initial RAM threshold in MB for the slider (default: 500.0)")
    parser.add_argument("-O", "--optimizer", type=str, default="greedy", choices=['greedy', 'optimal', 'branch-all', 'branch-2', 'auto'], help="Contraction path optimizer to use (default: greedy). Note: cotengra requires library installation.")
    parser.add_argument("--no-browser", action="store_true", help="Prevent the browser from opening automatically")

    args = parser.parse_args()

    if not os.path.exists(args.input_file):
        print(f"Error: File '{args.input_file}' not found.")
        sys.exit(1)

    with open(args.input_file, "r", encoding="utf-8") as f:
        qasm_content = f.read()

    print(f"Parsing {args.input_file}...")
    
    qasm_parser = QASMParser()
    try:
        inputs, output, shapes, gate_info = qasm_parser.parse(qasm_content)
        circuit_svg = qasm_parser.generate_svg(gate_info)
    except Exception as e:
        print(f"Error parsing QASM: {e}")
        sys.exit(1)

    print(f"Analyzing Tensor Network (Optimizer: {args.optimizer})...")

    try:
        nodes, edges, path, info = analyze_tensors(inputs, output, shapes, optimizer=args.optimizer)
    except Exception as e:
        print(f"Error during tensor contraction analysis: {e}")
        sys.exit(1)

    print(f"Rendering interactive UI to {args.output}...")

    render_graph(
        nodes, 
        edges, 
        circuit_svg=circuit_svg, 
        memory_limit_mb=args.memory, 
        output_html=args.output
    )

    max_mem = max([n.get('memory_mb', 0) for n in nodes if n['group'] != 'initial'])
    print(f"Success! Peak memory required: {max_mem:.2f} MB")

    if not args.no_browser:
        filepath = 'file://' + os.path.realpath(args.output)
        webbrowser.open(filepath)

if __name__ == "__main__":
    main()