"""Audit the AME diagram against the current browser project; generate XML on demand.

Run from the repository root with .venv/bin/python. Does not run Amesim or use
the stale eight-branch compiled/result cache inside the four-branch archive.
"""
from __future__ import annotations

import argparse
from collections import Counter, defaultdict
from copy import deepcopy
import hashlib
import json
from pathlib import Path
import re
import sys
if '--check' in sys.argv:
    sys.dont_write_bytecode = True
import struct
import tarfile
import xml.etree.ElementTree as ET

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))
from app.simulation.registry import get_component_model_spec
from tests.test_test_mql_ame_contract import (
    _blocks, _required_element_text, _element_text, _resolve_globals,
    _public_model_type, _expected_parameters, _value_fields,
)

OUT = ROOT / 'test/model-audit/mql4'
AME = ROOT / 'tests/data/test_mql_4.ame'
XML = OUT / 'aligned-input.xml'
REFERENCE = ROOT / 'tests/baselines/simulation/test_mql_4/test-mql-4-amesim-reference.json'
PROJECT = ROOT / 'tests/data/test-mql-4-corrected.json'


def sha(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def public_port(model, index):
    if index == 0 and model in {'amesim_pnvo001', 'amesim_forc'}:
        return 'res'
    if index == 0 and model in {'amesim_step0', 'amesim_ud00'}:
        return 'out'
    return f'port_{index + 1}'


def parameter_key(values):
    return tuple(sorted((key, format(float(value), '.12g')) for key, value in values.items()))


def audit_source():
    with tarfile.open(AME) as archive:
        def read(suffix):
            names = [name for name in archive.getnames() if name.endswith(suffix)]
            if len(names) != 1:
                raise ValueError((suffix, names))
            return archive.extractfile(names[0]).read().decode('latin1')
        cir = read('.cir')
        sim = read('.sim')
        modelinfo = read('.modelinfo')
        variables = read('.var')
        perf = read('.ameperf')
        result_name = next(name for name in archive.getnames() if name.endswith('.results'))
        result_data = archive.extractfile(result_name).read()
        result_points, result_columns = struct.unpack('<2i', result_data[:8])
        member_hashes = {}
        for suffix in ['.cir', '.sim', '.modelinfo', '.var', '.results']:
            name = next(name for name in archive.getnames() if name.endswith(suffix))
            data = archive.extractfile(name).read()
            member_hashes[name] = {'bytes': len(data), 'sha256': hashlib.sha256(data).hexdigest()}
    globals_ = _resolve_globals(cir)
    components = _blocks(cir, 'COMP')
    lines = _blocks(cir, 'LINE')
    nodes = {}
    for kind, entities in [('component', components), ('line', lines)]:
        for index, body in enumerate(entities):
            submodel = _required_element_text(body, 'SUB_NAME')
            if submodel == 'DIRECT':
                continue
            model = _public_model_type(submodel)
            nodes[f'{kind}:{index}'] = {
                'modelType': model,
                'alias': _required_element_text(body, 'ALIAS'),
                'parameters': _expected_parameters(body, model, globals_),
                'sourceValues': {name: {'expression': value.expression, 'unit': value.unit}
                                 for name, value in _value_fields(body).items()},
                'position': _element_text(body, 'COMP_POS') or _element_text(body, 'LINE_POINTS'),
                'geometry': _element_text(body, 'COMP_GEOMETRY'),
                'outputType': _element_text(body, 'OUTPUT_TYPE'),
            }
    edges = set()
    contacts = set()
    for index, body in enumerate(components):
        for port_index, port_body in enumerate(_blocks(body, 'COMP_PORT')):
            if _required_element_text(port_body, 'PORT_CONNECT') != '1':
                continue
            for connection in _blocks(port_body, 'CONNECT'):
                other = int(_required_element_text(connection, 'CONNECT_ENTITY_NUM'))
                other_port = int(_required_element_text(connection, 'CONNECT_ENTITY_PORT'))
                contacts.add(((index, port_index), (other, other_port)))
    assert all((b, a) in contacts for a, b in contacts)
    def endpoint(index, port):
        key = f'component:{index}'
        return key, public_port(nodes[key]['modelType'], port)
    for a, b in contacts:
        edges.add(tuple(sorted((endpoint(*a), endpoint(*b)))))
    for index, body in enumerate(lines):
        assert _required_element_text(body, 'LINE_START_TYPE') == '0'
        assert _required_element_text(body, 'LINE_END_TYPE') == '0'
        start = endpoint(int(_required_element_text(body, 'LINE_START_ENTITY')),
                         int(_required_element_text(body, 'LINE_START_PORT')))
        end = endpoint(int(_required_element_text(body, 'LINE_END_ENTITY')),
                       int(_required_element_text(body, 'LINE_END_PORT')))
        if _required_element_text(body, 'SUB_NAME') == 'DIRECT':
            edges.add(tuple(sorted((start, end))))
        else:
            key = f'line:{index}'
            output_type = _required_element_text(body, 'OUTPUT_TYPE')
            assert output_type in {'1', '2'}
            first, second = ('port_1', 'port_2') if output_type == '1' else ('port_2', 'port_1')
            edges.add(tuple(sorted((start, (key, first)))))
            edges.add(tuple(sorted((end, (key, second)))))
    count = Counter(node['modelType'] for node in nodes.values())
    states = sum(count[model] * size for model, size in {
        'amesim_pnch012': 2, 'amesim_pnch023': 2, 'amesim_pnl0001': 2,
        'amesim_pnl0002': 2, 'amesim_pnl0003': 4, 'amesim_mecmas21': 2,
    }.items())
    return nodes, sorted(edges), {
        'archiveSha256': sha(AME), 'archiveBytes': AME.stat().st_size,
        'members': member_hashes, 'globals': globals_,
        'componentCount': len(components), 'modeledLineCount': len(nodes) - len(components),
        'directLineCount': len(lines) - (len(nodes) - len(components)),
        'contactCount': len(contacts) // 2, 'publicNodeCount': len(nodes),
        'publicParameterCount': sum(len(node['parameters']) for node in nodes.values()),
        'connectionCount': len(edges), 'typeCounts': dict(count),
        'diagramDynamicStateCount': states,
        'simulationFileRaw': sim, 'savedSimulationNumericRows': [[float(x) for x in row.split()] for row in sim.splitlines()],
        'cachedModelInfo': modelinfo,
        'cachedResultsSuitableForFourBranchComparison': False,
        'cacheRejectionReason': 'Diagram has 64 states / four pistons; cached modelinfo has 132 states, .ameperf contains eight LSTP00A instances, and .results declares 1116 saved variables while the current .var has only 640 rows.',
        'cachedResultHeader': {'pointCount': result_points, 'savedVariableCount': abs(result_columns)},
        'variableRowCount': len(variables.splitlines()),
        'cachedPerformanceContactInstances': sorted(set(re.findall(r'LSTP00A-([0-9]+)', perf))),
        'cachedPistonAliases': sorted(set(re.findall(r'Data_Path=\S+@(pn_brp2[^\s]*)', variables))),
    }


def project_xml():
    # The separately maintained XML fixture was deleted. Use the same input
    # conversion as the native CLI; the AME remains the independent authority.
    from app.simulation.native_codegen.input import load_input
    xml, _ = load_input(PROJECT)
    return xml


def xml_model(root=None):
    if root is None:
        root = ET.fromstring(project_xml())
    nodes = {item.attrib['id']: {
        'modelType': item.attrib['type'],
        'modelVersion': item.attrib['modelVersion'],
        'parameters': {p.attrib['name']: float(p.attrib['value']) for p in item.findall('Parameter')},
    } for item in root.findall('./Components/Component')}
    edges = [tuple(sorted((ep.attrib['component'], ep.attrib['port']) for ep in item.findall('Endpoint')))
             for item in root.findall('./Connections/Connection')]
    return nodes, edges, root


def find_mapping(first, first_edges, second, second_edges, *, ports=False, parameters=True):
    """Find an actual graph bijection, preserving node types and all parameters."""
    def edge_label(node, port):
        if not ports:
            return ''
        return port
    def graph(nodes, edges):
        adjacency = {key: defaultdict(Counter) for key in nodes}
        for (a, ap), (b, bp) in edges:
            adjacency[a][b][(edge_label(a, ap), edge_label(b, bp))] += 1
            adjacency[b][a][(edge_label(b, bp), edge_label(a, ap))] += 1
        return adjacency
    a, b = graph(first, first_edges), graph(second, second_edges)
    def signature(nodes, adj, key):
        return (nodes[key]['modelType'], parameter_key(nodes[key]['parameters']) if parameters else (),
                sum(sum(labels.values()) for labels in adj[key].values()))
    candidates = {key: [other for other in second if signature(first, a, key) == signature(second, b, other)]
                  for key in first}
    if any(not options for options in candidates.values()):
        return None
    mapping = {}
    used = set()
    def search():
        if len(mapping) == len(first):
            return dict(mapping)
        key = min((key for key in first if key not in mapping),
                  key=lambda key: (-sum(neighbor in mapping for neighbor in a[key]),
                                   len([other for other in candidates[key] if other not in used]), -len(a[key])))
        for other in candidates[key]:
            if other in used:
                continue
            if any(a[key].get(existing, Counter()) != b[other].get(target, Counter())
                   for existing, target in mapping.items()):
                continue
            mapping[key] = other
            used.add(other)
            result = search()
            if result is not None:
                return result
            used.remove(other)
            del mapping[key]
        return None
    return search()


def check_project_contract(project, xml_root):
    """Check generated port metadata and exact JSON/XML execution data, without writes.

    This covers the observed nullable optional-field failure and shared physical
    data, not a replacement for the complete frontend project parser.
    """
    xml_nodes, xml_edges, _ = xml_model(xml_root)
    json_nodes = {node['id']: node['data'] for node in project['nodes']}
    if len(json_nodes) != len(project['nodes']) or json_nodes.keys() != xml_nodes.keys():
        raise ValueError('JSON/XML component IDs differ or JSON contains duplicate IDs.')
    for node_id, data in json_nodes.items():
        expected = xml_nodes[node_id]
        if (data['componentType'] != expected['modelType'] or
                data['modelType'] != expected['modelType'] or
                data['modelVersion'] != expected['modelVersion']):
            raise ValueError(f'{node_id}: JSON/XML model type or version differs.')
        if data['parameters'] != expected['parameters']:
            raise ValueError(f'{node_id}: JSON/XML parameters differ.')
        for port in data['ports']:
            # frontend/src/App.tsx:isProjectPortDefinitionValue accepts omission,
            # or the literal intoComponent. An explicit null rejects the import.
            if ('positiveFlowDirection' in port and
                    port['positiveFlowDirection'] != 'intoComponent'):
                raise ValueError(f"{node_id}.{port['name']}: omit undefined positiveFlowDirection instead of null.")
    json_edges = []
    for edge in project['edges']:
        endpoints = []
        for side in ['source', 'target']:
            component, port = edge[side], edge[side + 'Handle']
            if component not in json_nodes:
                raise ValueError(f'Unknown connection component {component}.')
            data = json_nodes[component]
            active_ports = get_component_model_spec(data['modelType']).active_ports(data['parameters'])
            if port not in {item.name for item in active_ports}:
                raise ValueError(f'{component}.{port}: connection port is inactive or unknown.')
            endpoints.append((component, port))
        json_edges.append(tuple(sorted(endpoints)))
    if Counter(json_edges) != Counter(xml_edges):
        raise ValueError('JSON/XML connection endpoints differ.')
    simulation = xml_root.find('Simulation').attrib
    for key, xml_key in [('t_start', 'tStart'), ('t_stop', 'tStop'),
                         ('step', 'sampleStep'), ('max_step', 'maxStep')]:
        if project['simulation'][key] != float(simulation[xml_key]):
            raise ValueError(f'JSON/XML simulation {key} differs.')
    if project['simulation']['method'] != simulation['method']:
        raise ValueError('JSON/XML integration methods differ.')
    return {'valid': True, 'componentCount': len(json_nodes), 'connectionCount': len(json_edges),
            'parameterCount': sum(len(data['parameters']) for data in json_nodes.values()),
            'optionalPortFieldsValid': True, 'jsonXmlExecutionDataEqual': True}


def generate_project(source, source_edges, target, target_edges, xml_root, mapping):
    layout_source = PROJECT
    layout = json.loads(layout_source.read_text(encoding='utf-8'))
    layout_sha = sha(layout_source)
    layout_nodes = {node['id']: node for node in layout['nodes']}
    reverse_mapping = {value: key for key, value in mapping.items()}
    nodes = []
    for node_id, model in target.items():
        source_node = source[reverse_mapping[node_id]]
        template = layout_nodes.get(node_id)
        if template is None:
            template = next(node for node in layout['nodes'] if node['data']['modelType'] == model['modelType'])
        node = deepcopy(template)
        node['id'] = node_id
        node['selected'] = False
        node['data']['label'] = source_node['alias']
        node['data']['parameters'] = source_node['parameters']
        node['data']['modelVersion'] = model['modelVersion']
        node['data']['ports'] = get_component_model_spec(model['modelType']).as_catalog_dict()['ports']
        for port in node['data']['ports']:
            if port.get('positiveFlowDirection') is None:
                port.pop('positiveFlowDirection', None)
        if model['modelType'] == 'amesim_p4node2':
            node['data']['rotation'] = 180
            node['data']['mirrored'] = False
        if node_id not in layout_nodes:
            # The extra closing pipe sits below the four retained branches.
            node['position'] = {'x': 80.0, 'y': 900.0}
            node['data']['rotation'] = 0
        nodes.append(node)
    original_by_pair = {}
    for edge in layout['edges']:
        pair = tuple(sorted((edge['source'], edge['target'])))
        original_by_pair[pair] = edge
    xml_by_pair = {}
    for connection in xml_root.findall('./Connections/Connection'):
        pair = tuple(sorted(ep.attrib['component'] for ep in connection.findall('Endpoint')))
        xml_by_pair[pair] = connection
    edges = []
    changed_connections = []
    for ((a, ap), (b, bp)) in source_edges:
        a, b = mapping[a], mapping[b]
        pair = tuple(sorted((a, b)))
        existing = original_by_pair.get(pair)
        edge = deepcopy(existing) if existing else {'id': 'ame-' + a + '-' + ap + '-' + b + '-' + bp}
        if existing and existing['source'] == b:
            a, ap, b, bp = b, bp, a, ap
        old_connection = xml_by_pair[pair]
        old_endpoints = {ep.attrib['component']: ep.attrib['port'] for ep in old_connection.findall('Endpoint')}
        changed = old_endpoints != {a: ap, b: bp}
        if changed:
            changed_connections.append({'components': pair, 'before': old_endpoints, 'after': {a: ap, b: bp}})
        edge.update({'source': a, 'target': b, 'sourceHandle': ap, 'targetHandle': bp})
        edge.pop('selected', None)
        if any(target[node_id]['modelType'] == 'amesim_p4node2' for node_id in [a, b]):
            edge.setdefault('data', {}).pop('routePoints', None)
        if changed or existing is None:
            edge['data'] = {'isContactEdge': False}
        edges.append(edge)
        for endpoint in old_connection.findall('Endpoint'):
            endpoint.attrib['port'] = {a: ap, b: bp}[endpoint.attrib['component']]
    parameter_changes = []
    for component in xml_root.findall('./Components/Component'):
        node_id = component.attrib['id']
        parameters = source[reverse_mapping[node_id]]['parameters']
        for parameter in component.findall('Parameter'):
            key = parameter.attrib['name']
            before = float(parameter.attrib['value'])
            after = parameters[key]
            if abs(before - after) > max(1e-18, abs(after) * 1e-12):
                parameter_changes.append({'component': node_id, 'parameter': key, 'before': before, 'after': after})
            parameter.attrib['value'] = repr(after)
    simulation = {'t_start': 0.0, 't_stop': 10.0, 'step': 0.01, 'max_step': 1e30, 'method': 'BDF'}
    project = {'projectSchemaVersion': 1, 'name': 'test-mql-4-AME-aligned', 'nodes': nodes,
               'edges': edges, 'simulation': simulation}
    check_project_contract(project, xml_root)
    PROJECT.write_text(json.dumps(project, indent=2, ensure_ascii=False) + '\n', encoding='utf-8', newline='\n')
    generated_xml = OUT / 'aligned-input.xml'
    xml_root.attrib['name'] = project['name']
    ET.indent(xml_root)
    ET.ElementTree(xml_root).write(generated_xml, encoding='utf-8', xml_declaration=True)
    XML.write_bytes(generated_xml.read_bytes())
    # Check the delivered JSON itself, including every exact endpoint and number.
    json_nodes = {node['id']: node['data'] for node in json.loads(PROJECT.read_text(encoding='utf-8'))['nodes']}
    for source_id, project_id in mapping.items():
        assert json_nodes[project_id]['parameters'] == source[source_id]['parameters']
    expected_edges = {tuple(sorted(((mapping[a], ap), (mapping[b], bp)))) for ((a, ap), (b, bp)) in source_edges}
    actual_edges = {tuple(sorted(((edge['source'], edge['sourceHandle']), (edge['target'], edge['targetHandle'])))) for edge in edges}
    assert expected_edges == actual_edges
    for edge in edges:
        for side in ['source', 'target']:
            data = json_nodes[edge[side]]
            active_ports = get_component_model_spec(data['modelType']).active_ports(data['parameters'])
            assert edge[side + 'Handle'] in {port.name for port in active_ports}
    return {'path': str(PROJECT.relative_to(ROOT)), 'sha256': sha(PROJECT),
            'xmlPath': str(generated_xml.relative_to(ROOT)), 'xmlSha256': sha(generated_xml),
            'nodeCount': len(nodes), 'edgeCount': len(edges), 'parameterCount': sum(len(n['data']['parameters']) for n in nodes),
            'allMappedParametersExactlyMatchAme': True, 'allMappedConnectionPortsExactlyMatchAme': True,
            'parameterChangesFromLegacyXml': parameter_changes, 'connectionPortChangesFromLegacyXml': changed_connections,
            'layoutMetadataSource': {'path': str(layout_source.relative_to(ROOT)), 'sha256': layout_sha},
            'simulation': simulation,
            'simulationSettingsNote': 'AME archive stores historical fixed-step settings. This JSON uses native BDF; rtol is supplied by the backend and is not stored in the project. Initial validation used 1e-7; the backend default was subsequently changed to 1e-8. The earlier official Standard variable-step reference remains tolerance 1e-7. Inspect actual run diagnostics; solver algorithms are not identical.'}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--check', action='store_true', help='Check current JSON parameters and exact connection ports against AME; generate XML in memory, without file writes.')
    arguments = parser.parse_args()
    if arguments.check:
        target, target_edges, xml_root = xml_model()
        result = check_project_contract(json.loads(PROJECT.read_text(encoding='utf-8')), xml_root)
        source, source_edges, _ = audit_source()
        mapping = find_mapping(source, source_edges, target, target_edges, ports=True)
        if mapping is None or len(source) != len(target) or len(source_edges) != len(target_edges):
            raise ValueError('Current JSON differs from AME in components, parameters or connection ports.')
        for source_id, target_id in mapping.items():
            if source[source_id]['parameters'] != target[target_id]['parameters']:
                raise ValueError(f'{target_id}: parameters do not exactly match AME.')
        result['ameParametersAndConnectionPortsMatch'] = True
        result['executionXmlSource'] = str(PROJECT.relative_to(ROOT))
        print(json.dumps(result, indent=2))
        return
    OUT.mkdir(parents=True, exist_ok=True)
    source, source_edges, audit = audit_source()
    XML.write_bytes(project_xml())
    (OUT / 'input-before-audit.xml').write_bytes(XML.read_bytes())
    target, target_edges, xml_root = xml_model()
    mapping = find_mapping(source, source_edges, target, target_edges)
    audit['correctedXml'] = {'path': str(XML.relative_to(ROOT)), 'sha256': sha(XML),
                             'parameterAndComponentTopologyIsomorphic': mapping is not None}
    ref = json.loads(REFERENCE.read_text(encoding='utf-8'))
    audit['reference'] = {'path': str(REFERENCE.relative_to(ROOT)), 'sha256': sha(REFERENCE),
                          'source': ref['source'], 'timeCount': len(ref['times']), 'seriesCount': len(ref['series']),
                          'ameHashMatches': ref['originalAmeSha256'] == sha(AME),
                          'xmlHashMatches': ref['xmlSha256'] == sha(XML),
                          'newAmesimRunPerformed': False}
    (OUT / 'source-model.json').write_text(json.dumps({'nodes': source, 'edges': source_edges}, indent=2, ensure_ascii=False) + '\n', encoding='utf-8', newline='\n')
    if not mapping:
        mapping = find_mapping(source, source_edges, target, target_edges, parameters=False)
    if mapping:
        audit['componentMapping'] = {key: {'alias': source[key]['alias'], 'projectId': value} for key, value in mapping.items()}
        transformed = {tuple(sorted(((mapping[a], ap), (mapping[b], bp)))) for ((a, ap), (b, bp)) in source_edges}
        target_set = set(target_edges)
        audit['correctedXml']['exactPortDifferencesForTopologyMapping'] = {
            'ameOnly': sorted(transformed - target_set), 'xmlOnly': sorted(target_set - transformed)}
    if mapping:
        audit['generatedProject'] = generate_project(source, source_edges, target, target_edges, xml_root, mapping)
    audit['correctedXml']['scope'] = 'Input XML snapshot before this generation; generatedProject describes the delivered corrected files.'
    audit_text = json.dumps(audit, indent=2, ensure_ascii=False) + '\n'
    (OUT / 'model-audit.json').write_text(audit_text, encoding='utf-8', newline='\n')
    (ROOT / 'tests/model_audit/test-mql-4-audit.json').write_text(audit_text, encoding='utf-8', newline='\n')
    print(json.dumps({key: value for key, value in audit.items() if key not in {'componentMapping', 'members', 'globals'}}, indent=2))


if __name__ == '__main__':
    main()