import xml.etree.ElementTree as ET
from pyqasm.entrypoint import loads
from openqasm3 import ast
from .theme import ThemeManager

class SVGDrawer:
    def __init__(self, theme=None):
        self.theme_manager = ThemeManager(theme)

    def draw(self, program_str):
        if isinstance(program_str, str):
            module = loads(program_str)
        else:
            module = program_str

        module.unroll()
        module.remove_includes()

        quantum_registers = list(module._qubit_registers.keys()) if module._qubit_registers else []
        classical_registers = list(module._classical_registers.keys()) if module._classical_registers else []
        line_nums, sizes = self._compute_line_nums(module)
        statements = list(module._statements)
        # Collapse barriers into one:
        result = []
        for stmt in statements:
            if not (result and isinstance(stmt, ast.QuantumBarrier) and isinstance(result[-1], ast.QuantumBarrier)):
                result.append(stmt)
        statements = result
        moments, depths = self._compute_moments(statements, line_nums)
        
        n_lines = max(line_nums.values()) + 1 if line_nums else 0
        print("lines", n_lines)
        n_moments = len(moments)
        
        # No fallbacks allowed, these will raise KeyError if missing in theme
        gate_width = self.theme_manager.get_dimension('gate_width')
        gate_spacing = self.theme_manager.get_dimension('gate_spacing')
        line_spacing = self.theme_manager.get_dimension('line_spacing')
        padding = self.theme_manager.get_dimension('padding')
        label_offset = self.theme_manager.get_dimension('label_offset')
        
        width = (padding * 2) + (n_moments * (gate_width + gate_spacing)) + (label_offset * 1.5)
        height = (0.5 * padding) + (n_lines * line_spacing)
        
        svg = ET.Element('svg', {
            'xmlns': 'http://www.w3.org/2000/svg',
            'width': str(width),
            'height': str(height),
            'viewBox': f'0 0 {width} {height}'
        })
        
        bg_color = self.theme_manager.get_style('background')
        ET.SubElement(svg, 'rect', {
            'width': '100%',
            'height': '100%',
            'fill': bg_color
        })
        
        # Draw labels
        text_color = self.theme_manager.get_style('text')
        label_font = self.theme_manager.get_style('label_font')
        drawn_lines = set()
        # Sort items to ensure consistent label choice (e.g. prefer non-indexed for single bit)
        for (reg_name, reg_idx), line_idx in sorted(line_nums.items(), key=lambda x: x[0][1]):
            if line_idx in drawn_lines: continue
            drawn_lines.add(line_idx)
            
            y = padding + line_idx * line_spacing
            label = f"{reg_name}[{reg_idx}]" if reg_idx != -1 else reg_name
            text = ET.SubElement(svg, 'text', {
                'x': str(padding + label_offset - 10),
                'y': str(y),
                'fill': text_color,
                'font-family': label_font['family'],
                'font-size': str(label_font['size']),
                'text-anchor': 'end',
                'dominant-baseline': 'middle'
            })
            text.text = label

        # Draw wires (Stave)
        wire_config = self.theme_manager.get_style('qubit_wire')
        for q, line_idx in line_nums.items():
            y = padding + line_idx * line_spacing
            self._draw_line(svg, padding + label_offset, y, width - padding, y, wire_config)

        # Draw moments Phase 1: Vertical Lines
        gate_width = self.theme_manager.get_dimension('gate_width')
        gate_spacing = self.theme_manager.get_dimension('gate_spacing')
        x_start = padding + label_offset + gate_width/2
        
        x = x_start
        for moment in moments:
            for stmt in moment:
                self._draw_statement(svg, stmt, x, line_nums, phase='lines', qbits=quantum_registers, classical=classical_registers)
            x += gate_width + gate_spacing

        # Draw moments Phase 2: Shapes/Gates
        x = x_start
        for moment in moments:
            for stmt in moment:
                self._draw_statement(svg, stmt, x, line_nums, phase='shapes', qbits=quantum_registers, classical=classical_registers)
            x += gate_width + gate_spacing

        return ET.tostring(svg, encoding='unicode')

    def _draw_statement(self, svg, stmt, x, line_nums, phase='shapes', qbits=[], classical=[]):
        if isinstance(stmt, ast.QuantumGate):
            self._draw_gate(svg, stmt, x, line_nums, phase)
        elif isinstance(stmt, ast.QuantumMeasurementStatement):
            self._draw_measurement(svg, stmt, x, line_nums, phase)
        elif isinstance(stmt, ast.QuantumBarrier):
            if phase == 'lines': # Barrier is a line
                self._draw_barrier(svg, stmt, x, line_nums, qbits)

    def _draw_gate(self, svg, gate, x, line_nums, phase='shapes'):
        name = gate.name.name.lower()
        qubits = [self._identifier_to_key(q) for q in gate.qubits]
        lines = [line_nums[q] for q in qubits]
        
        sub = self.theme_manager.get_substitution(name)
        config = self.theme_manager.get_gate_config(name)
        
        if name == "swap" and len(lines) == 2:
            self._draw_swap(svg, lines[0], lines[1], x, phase)
            return
            
        # Determine if it's a controlled gate (starts with 'c', but not 'cz', 'ccx' etc handles specifically or generally)
        is_controlled = name.startswith('c') and name != 'curry' # simple heuristic
        if name in ['ccx', 'ccz']: is_controlled = True

        if is_controlled and len(lines) > 1:
            self._draw_controlled_gate(svg, name, lines, x, phase, gate)
            return

        if phase == 'lines':
            if len(lines) > 1:
                # Generic multi-qubit vertical connection
                padding = self.theme_manager.get_dimension('padding')
                line_spacing = self.theme_manager.get_dimension('line_spacing')
                y_min = padding + min(lines) * line_spacing
                y_max = padding + max(lines) * line_spacing
                conn_config = self.theme_manager.get_style('connection_line')
                self._draw_line(svg, x, y_min, x, y_max, conn_config)
            return

        padding = self.theme_manager.get_dimension('padding')
        line_spacing = self.theme_manager.get_dimension('line_spacing')

        # Evaluate parameters if any
        params = []
        if hasattr(gate, 'arguments') and gate.arguments:
            from pyqasm.expressions import Qasm3ExprEvaluator
            for arg in gate.arguments:
                try:
                    val = Qasm3ExprEvaluator.evaluate_expression(arg)[0]
                    params.append(val)
                except:
                    pass

        for line_idx in lines:
            y = padding + line_idx * line_spacing
            
            # Determine label: theme override > default name
            gate_label = config.get('label', name.upper())
            if 'text' in config:
                gate_label = config['text']
            
            if sub:
                self._draw_shape(svg, x, y, sub, label=gate_label, params=params)
            else:
                self._draw_shape(svg, x, y, config, label=gate_label, params=params)

    def _draw_controlled_gate(self, svg, name, lines, x, phase, gate):
        # By convention, the last qubit is the target, others are controls
        ctrl_lines = lines[:-1]
        target_line = lines[-1]
        
        if phase == 'lines':
            padding = self.theme_manager.get_dimension('padding')
            line_spacing = self.theme_manager.get_dimension('line_spacing')
            y_min = padding + min(lines) * line_spacing
            y_max = padding + max(lines) * line_spacing
            conn_config = self.theme_manager.get_style('connection_line')
            self._draw_line(svg, x, y_min, x, y_max, conn_config)
            return

        # Draw controls
        ctrl_config = self.theme_manager.get_shape_config('control_dot')
        padding = self.theme_manager.get_dimension('padding')
        line_spacing = self.theme_manager.get_dimension('line_spacing')
        for cl in ctrl_lines:
            y = padding + cl * line_spacing
            # Theme override for control dot per gate
            gate_config = self.theme_manager.theme.get('gates', {}).get(name, {})
            specific_ctrl = gate_config.get('control_shape', ctrl_config)
            if isinstance(specific_ctrl, str): # if it's just a shape name reference
                specific_ctrl = self.theme_manager.get_shape_config(specific_ctrl)
            self._draw_shape(svg, x, y, specific_ctrl)

        # Draw target
        y_target = padding + target_line * line_spacing
        
        # Determine target shape
        if name in ['cx', 'ccx']:
            target_config = self.theme_manager.get_shape_config('target_plus')
            self._draw_shape(svg, x, y_target, target_config)
        else:
            # Strip leading 'c' to find base gate config (e.g. 'cz' -> 'z')
            base_name = name[1:] if name.startswith('c') else name
            if name.startswith('cc'): base_name = name[2:]
            
            target_config = self.theme_manager.get_gate_config(base_name)
            
            # Evaluate parameters for target gate
            params = []
            if hasattr(gate, 'arguments') and gate.arguments:
                from pyqasm.expressions import Qasm3ExprEvaluator
                for arg in gate.arguments:
                    params.append(Qasm3ExprEvaluator.evaluate_expression(arg)[0])
            
            gate_label = target_config.get('label', base_name.upper())
            if 'text' in target_config: gate_label = target_config['text']
            
            self._draw_shape(svg, x, y_target, target_config, label=gate_label, params=params)

    def _draw_cx(self, svg, ctrl_line, target_line, x, phase='shapes'):
        # Just route to the new generalized controlled gate logic
        # This keeps the signature for backward compatibility if needed internally
        pass 

    def _draw_swap(self, svg, line1, line2, x, phase='shapes'):
        padding = self.theme_manager.get_dimension('padding')
        line_spacing = self.theme_manager.get_dimension('line_spacing')
        
        y1 = padding + line1 * line_spacing
        y2 = padding + line2 * line_spacing
        
        if phase == 'lines':
            conn_config = self.theme_manager.get_style('connection_line')
            self._draw_line(svg, x, y1, x, y2, conn_config)
            return
        
        cross_config = self.theme_manager.get_shape_config('swap_x')
        for y in [y1, y2]:
            self._draw_shape(svg, x, y, cross_config)

    def _draw_measurement(self, svg, stmt, x, line_nums, phase='shapes'):
        padding = self.theme_manager.get_dimension('padding')
        line_spacing = self.theme_manager.get_dimension('line_spacing')
        qubit = self._identifier_to_key(stmt.measure.qubit)
        line_idx = line_nums[qubit]
        y = padding + line_idx * line_spacing
        
        target_key = None
        if stmt.target:
            target_key = self._identifier_to_key(stmt.target)
        else:
            # Strictly assume classical registry named 'c', index 0
            if ('c', 0) in line_nums:
                target_key = ('c', 0)
            elif ('c', -1) in line_nums:
                target_key = ('c', -1)

        if phase == 'lines':
            # Draw connection to classical target
            if target_key and target_key in line_nums:
                target_idx = line_nums[target_key]
                y_target = padding + target_idx * line_spacing
                conn_config = self.theme_manager.get_style('measurement_line')
                self._draw_line(svg, x, y, x, y_target, conn_config)
            return

        config = self.theme_manager.get_gate_config('measurement')
        
        gate_label = config.get('label', 'M')
        if 'text' in config:
            gate_label = config['text']
            
        self._draw_shape(svg, x, y, config, label=gate_label)

    def _draw_barrier(self, svg, stmt, x, line_nums, qbits):
        padding = self.theme_manager.get_dimension('padding')
        line_spacing = self.theme_manager.get_dimension('line_spacing')
        barrier_config = self.theme_manager.get_style('barrier')
        barrier_padding = self.theme_manager.get_dimension('barrier_padding')
        
        lines = [v for k, v in line_nums.items() if k[0] in qbits]
        if not lines: return
        
        y_min = padding + min(lines) * line_spacing - barrier_padding
        y_max = padding + max(lines) * line_spacing + barrier_padding
        
        self._draw_line(svg, x, y_min, x, y_max, barrier_config)

    def _draw_line(self, svg, x1, y1, x2, y2, config):
        style = config['style']
        stroke = config['stroke']
        width = config['stroke_width']
        dash = config.get('dasharray', '')
        
        if style == 'wave':
            amp = config['amplitude']
            wl = config['wavelength']
            path_data = self._wave_path(x1, y1, x2, y2, amp, wl)
            ET.SubElement(svg, 'path', {
                'd': path_data,
                'stroke': stroke,
                'stroke-width': str(width),
                'fill': 'none',
                'stroke-dasharray': dash
            })
        else:
            ET.SubElement(svg, 'line', {
                'x1': str(x1), 'y1': str(y1), 'x2': str(x2), 'y2': str(y2),
                'stroke': stroke, 'stroke-width': str(width),
                'stroke-dasharray': dash
            })

    def _wave_path(self, x1, y1, x2, y2, amplitude, wavelength):
        import math
        dist = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
        angle = math.atan2(y2 - y1, x2 - x1)
        
        # For simplicity and precision, we generate points and join with L
        res = 10 # points per wavelength
        num_points = int(dist / wavelength * res)
        d_points = []
        for i in range(num_points + 1):
            t = i / num_points
            curr_dist = t * dist
            
            # Local wave offset
            wave_y = amplitude * math.sin(2 * math.pi * curr_dist / wavelength)
            
            # Rotate and translate
            gx = x1 + curr_dist * math.cos(angle) - wave_y * math.sin(angle)
            gy = y1 + curr_dist * math.sin(angle) + wave_y * math.cos(angle)
            d_points.append(f"{gx} {gy}")
            
        return "M " + " L ".join(d_points)

    def _draw_shape(self, svg, x, y, config, label=None, params=None):
        shape_type = config['type']
        
        if shape_type == 'circle':
            radius = config['radius']
            ET.SubElement(svg, 'circle', {
                'cx': str(x), 'cy': str(y), 'r': str(radius),
                'fill': config.get('fill', 'none'),
                'stroke': config.get('stroke', 'none'),
                'stroke-width': str(config.get('stroke_width', 1))
            })
            # Draw parametric arc if requested
            if params and config.get('parametric_mode') == 'arc':
                self._draw_parametric_arc(svg, x, y, radius, params[0], config)

        elif shape_type == 'rect':
            w = config['width']
            h = config['height']
            r = config.get('radius', 0)
            ET.SubElement(svg, 'rect', {
                'x': str(x - w/2), 'y': str(y - h/2),
                'width': str(w), 'height': str(h), 'rx': str(r),
                'fill': config.get('fill', 'none'),
                'stroke': config.get('stroke', 'none'),
                'stroke-width': str(config.get('stroke_width', 1))
            })
        elif shape_type == 'diamond':
            size = config.get('radius', config.get('size', 20))
            points = f"{x},{y-size} {x+size},{y} {x},{y+size} {x-size},{y}"
            ET.SubElement(svg, 'polygon', {
                'points': points,
                'fill': config.get('fill', 'none'),
                'stroke': config.get('stroke', 'none'),
                'stroke-width': str(config.get('stroke_width', 1))
            })
        elif shape_type == 'emoji':
            font_size = config.get('font_size', 24)
            text = ET.SubElement(svg, 'text', {
                'x': str(x), 'y': str(y),
                'font-size': str(font_size),
                'text-anchor': 'middle', 
                'dominant-baseline': 'middle'
            })
            text.text = config['value']
        elif shape_type == 'image':
            w = self.theme_manager.get_dimension('gate_width')
            h = self.theme_manager.get_dimension('gate_height')
            ET.SubElement(svg, 'image', {
                'href': config['value'],
                'x': str(x - w/2), 'y': str(y - h/2),
                'width': str(w), 'height': str(h)
            })
        elif shape_type == 'cross':
            size = config['size']
            line_config = {
                'style': config.get('style', 'straight'),
                'stroke': config['stroke'],
                'stroke_width': config['stroke_width'],
                'dasharray': config.get('dasharray', ''),
                'amplitude': config.get('amplitude', 2),
                'wavelength': config.get('wavelength', 4)
            }
            self._draw_line(svg, x - size, y - size, x + size, y + size, line_config)
            self._draw_line(svg, x - size, y + size, x + size, y - size, line_config)
        elif shape_type == 'plus_circle':
            radius = config['radius']
            stroke = config['stroke']
            sw = config['stroke_width']
            fill = config.get('fill', 'none')
            ET.SubElement(svg, 'circle', {
                'cx': str(x), 'cy': str(y), 'r': str(radius),
                'fill': fill, 'stroke': stroke, 'stroke-width': str(sw)
            })
            ET.SubElement(svg, 'line', {
                'x1': str(x - radius), 'y1': str(y), 'x2': str(x + radius), 'y2': str(y),
                'stroke': stroke, 'stroke-width': str(sw)
            })
            ET.SubElement(svg, 'line', {
                'x1': str(x), 'y1': str(y - radius), 'x2': str(x), 'y2': str(y + radius),
                'stroke': stroke, 'stroke-width': str(sw)
            })
        elif shape_type == 'svg':
             pass

        if label:
            text_color = self.theme_manager.get_style('text')
            txt = ET.SubElement(svg, 'text', {
                'x': str(x), 'y': str(y),
                'fill': text_color, 
                'font-family': 'sans-serif', 
                'font-size': '12',
                'text-anchor': 'middle', 
                'dominant-baseline': 'middle',
            })
            txt.text = label

    def _draw_parametric_arc(self, svg, x, y, base_radius, theta, config):
        import math
        arc_stroke_width = config.get('arc_stroke_width', 4)
        arc_stroke = config.get('arc_stroke', '#ff0000')
        
        # Determine radius for the arc
        if 'arc_radius' in config:
            radius = config['arc_radius']
        else:
            # Fallback to inner rim logic if not specified
            base_stroke_width = config.get('stroke_width', 1)
            radius = base_radius - base_stroke_width - arc_stroke_width/2 - 1
        
        # theta is in radians. 2*pi is full circle.
        # We start from top (3*pi/2)
        start_angle = -math.pi / 2
        end_angle = start_angle + theta
        
        x1 = x + radius * math.cos(start_angle)
        y1 = y + radius * math.sin(start_angle)
        x2 = x + radius * math.cos(end_angle)
        y2 = y + radius * math.sin(end_angle)
        
        large_arc_flag = 1 if abs(theta) > math.pi else 0
        sweep_flag = 1 if theta >= 0 else 0 # 1=clockwise (right), 0=anticlockwise (left)
        
        if abs(theta) >= 2 * math.pi:
            # Full circle doesn't work well with arc command, draw a circle instead
            ET.SubElement(svg, 'circle', {
                'cx': str(x), 'cy': str(y), 'r': str(radius),
                'fill': 'none', 'stroke': arc_stroke, 'stroke-width': str(arc_stroke_width)
            })
        else:
            d = f"M {x1} {y1} A {radius} {radius} 0 {large_arc_flag} {sweep_flag} {x2} {y2}"
            ET.SubElement(svg, 'path', {
                'd': d,
                'fill': 'none',
                'stroke': arc_stroke,
                'stroke-width': str(arc_stroke_width),
                'stroke-linecap': 'round'
            })

    def _compute_moments(self, statements, line_nums):
        # Full implementation of moment computation
        depths = {k: -1 for k in line_nums}
        moments = []
        
        for statement in statements:
            # Filter non-quantum statements
            if isinstance(statement, (ast.CalibrationGrammarDeclaration, ast.ClassicalDeclaration, 
                                     ast.ConstantDeclaration, ast.ExternDeclaration, 
                                     ast.IODeclaration, ast.QubitDeclaration)):
                continue

            if isinstance(statement, ast.QuantumGate):
                qubits = [self._identifier_to_key(q) for q in statement.qubits]
                lines = [line_nums[q] for q in qubits]
                
                # Clearance behavior: occupy all lines in the vertical span
                min_l, max_l = min(lines), max(lines)
                keys_in_span = [k for k, v in line_nums.items() if min_l <= v <= max_l]
                
                depth = 1 + max(depths[key] for key in keys_in_span)
                for key in keys_in_span:
                    depths[key] = depth
                    
            elif isinstance(statement, ast.QuantumMeasurementStatement):
                key = self._identifier_to_key(statement.measure.qubit)
                target_key = None
                if statement.target:
                    target_key = self._identifier_to_key(statement.target)
                else:
                    # Strictly target 'c[0]' or 'c'
                    if ('c', 0) in line_nums: target_key = ('c', 0)
                    elif ('c', -1) in line_nums: target_key = ('c', -1)

                keys_in_span = [k for k, v in line_nums.items()]
                depth = 1 + max(depths[k] for k in keys_in_span)
                for k in keys_in_span:
                    depths[k] = depth
            elif isinstance(statement, ast.QuantumBarrier):
                qubits = [self._identifier_to_key(q) for q in statement.qubits]
                lines = [line_nums[q] for q in qubits]
                
                # Clearance behavior: occupy all lines in the vertical span
                keys_in_span = [k for k, v in line_nums.items()]
                depth = 1 + max(depths[key] for key in keys_in_span)
                for key in keys_in_span:
                    depths[key] = depth
            else:
                # Skip other statements for now
                continue

            if depth >= len(moments):
                for _ in range(depth - len(moments) + 1):
                    moments.append([])
            
            moments[depth].append(statement)
            
        return moments, depths

    def _identifier_to_key(self, identifier):
        from pyqasm.expressions import Qasm3ExprEvaluator
        if isinstance(identifier, ast.Identifier):
            return identifier.name, -1

        indices = identifier.indices
        if len(indices) >= 1 and isinstance(indices[0], list) and len(indices[0]) >= 1:
            return (
                identifier.name.name,
                Qasm3ExprEvaluator.evaluate_expression(indices[0][0])[0],
            )
        raise ValueError(f"Unsupported identifier: {identifier}")

    def _compute_line_nums(self, module):
        line_nums = {}
        sizes = {}
        line_num = 0
        
        # Get ordering preferences from theme (dimensions or layout)
        try:
            rev_q = self.theme_manager.get_dimension('reverse_qubit_order')
        except:
            rev_q = False
            
        try:
            rev_c = self.theme_manager.get_dimension('reverse_classical_order')
        except:
            rev_c = False

        # Qubit registers first (at the top)
        for qubit_reg in module._qubit_registers:
            size = module._qubit_registers[qubit_reg]
            indices = list(range(size))
            if rev_q:
                indices.reverse()
            for i in indices:
                line_nums[(qubit_reg, i)] = line_num
                line_num += 1

        # Classical registers second (at the bottom)
        # Any measures on statements??:
        any_measures = any(isinstance(s, ast.QuantumMeasurementStatement) for s in module._statements)

        if any_measures and not module._classical_registers:
            module._classical_registers = {"c": 1}
        for k in module._classical_registers:
            size = module._classical_registers[k]
            indices = list(range(size))
            if rev_c:
                indices.reverse()
            
            if size == 1:
                line_nums[(k, -1)] = line_num
                line_nums[(k, 0)] = line_num # also map index 0 to same line
                sizes[(k, -1)] = 1
                line_num += 1
            else:
                for i in indices:
                    line_nums[(k, i)] = line_num
                    sizes[(k, i)] = 1
                    line_num += 1
            
        return line_nums, sizes

def draw(program, theme=None, filename=None):
    drawer = SVGDrawer(theme)
    svg_content = drawer.draw(program)
    if filename:
        with open(filename, 'w') as f:
            f.write(svg_content)
    return svg_content
