#!/usr/bin/env python3
"""
Post-process exported OGRE mesh files to fix normal/tangent seams
between body part meshes at matching vertex positions.

Also fixes pose normal offsets (shape key normals) so that animated
shape keys don't reintroduce seams at runtime.

Usage: python3 fix_ogre_mesh_seams.py <mesh_directory> [ogre_xml_converter_path]
"""

import os
import sys
import subprocess
import xml.etree.ElementTree as ET
import math

# Default path to OgreXMLConverter (can be overridden via argument or env var)
DEFAULT_OGRE_XML_CONVERTER = os.environ.get(
    'OGRETOOLS_XML_CONVERTER',
    '/media/slapin/library/ogre3/ogre-sdk/bin/OgreXMLConverter'
)

# Body part prefixes to process
BODY_PREFIXES = ['BodyTop', 'BodyBottom', 'BodyFeet']


def find_converter():
    """Find OgreXMLConverter executable."""
    if len(sys.argv) >= 3:
        return sys.argv[2]
    if os.path.exists(DEFAULT_OGRE_XML_CONVERTER):
        return DEFAULT_OGRE_XML_CONVERTER
    # Try PATH
    for path in os.environ.get('PATH', '').split(os.pathsep):
        exe = os.path.join(path, 'OgreXMLConverter')
        if os.path.exists(exe):
            return exe
    return None


def mesh_to_xml(mesh_path, converter):
    xml_path = mesh_path + '.xml'
    result = subprocess.run(
        [converter, mesh_path, xml_path],
        capture_output=True, text=True
    )
    if result.returncode != 0:
        print(f"WARNING: Failed to convert {mesh_path} to XML:")
        print(result.stdout)
        print(result.stderr)
        return None
    return xml_path


def xml_to_mesh(xml_path, converter):
    mesh_path = xml_path[:-4]  # remove .xml
    result = subprocess.run(
        [converter, xml_path, mesh_path],
        capture_output=True, text=True
    )
    if result.returncode != 0:
        print(f"WARNING: Failed to convert {xml_path} to binary:")
        print(result.stdout)
        print(result.stderr)
        return None
    return mesh_path


def parse_mesh_xml(xml_path):
    """Parse mesh XML and return (tree, vertices list, poses dict)."""
    try:
        tree = ET.parse(xml_path)
    except ET.ParseError as e:
        print(f"ERROR: Failed to parse {xml_path}: {e}")
        return None, None, None

    root = tree.getroot()
    sharedgeom = root.find('.//sharedgeometry')
    if sharedgeom is None:
        return None, None, None

    vbs = sharedgeom.findall('vertexbuffer')
    if len(vbs) < 1:
        return None, None, None

    # The first vertexbuffer has positions and normals
    pos_norm_vb = vbs[0]
    # The second vertexbuffer (if any) has texcoords and tangents
    tex_tan_vb = vbs[1] if len(vbs) > 1 else None

    pn_verts = pos_norm_vb.findall('vertex')
    tt_verts = tex_tan_vb.findall('vertex') if tex_tan_vb is not None else []

    vertices = []
    for i, pnv in enumerate(pn_verts):
        pos = pnv.find('position')
        normal = pnv.find('normal')
        if pos is None or normal is None:
            continue

        p = (float(pos.get('x')), float(pos.get('y')), float(pos.get('z')))
        n = (float(normal.get('x')), float(normal.get('y')), float(normal.get('z')))

        t = None
        if i < len(tt_verts):
            tangent = tt_verts[i].find('tangent')
            if tangent is not None:
                t = (float(tangent.get('x')), float(tangent.get('y')), float(tangent.get('z')))

        vertices.append({
            'index': i,
            'pos': p,
            'normal': n,
            'tangent': t,
        })

    # Parse poses: pose_name -> {vertex_index -> (nx, ny, nz)}
    poses = {}
    for pose in root.findall('.//poses/pose'):
        name = pose.get('name')
        pose_offsets = {}
        for po in pose.findall('poseoffset'):
            idx = int(po.get('index'))
            nx = float(po.get('nx', 0.0))
            ny = float(po.get('ny', 0.0))
            nz = float(po.get('nz', 0.0))
            pose_offsets[idx] = (nx, ny, nz)
        poses[name] = pose_offsets

    return tree, vertices, poses


def round_pos(p, tolerance=0.0001):
    return (
        round(p[0] / tolerance) * tolerance,
        round(p[1] / tolerance) * tolerance,
        round(p[2] / tolerance) * tolerance,
    )


def vector_normalize(v):
    length = math.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
    if length > 0.0001:
        return (v[0] / length, v[1] / length, v[2] / length)
    return v


def fix_seams(mesh_dir, converter):
    """Fix normal/tangent/pose-normal seams across all body part meshes in directory."""
    # Find all body part mesh files
    mesh_files = []
    for f in os.listdir(mesh_dir):
        if not f.endswith('.mesh') or f.endswith('.mesh.xml'):
            continue
        for prefix in BODY_PREFIXES:
            if prefix in f:
                mesh_files.append(os.path.join(mesh_dir, f))
                break

    if not mesh_files:
        print(f"No body part meshes found in {mesh_dir}")
        return 0

    print(f"\n[OGRE Mesh Seam Fix] Found {len(mesh_files)} body part meshes")
    for mf in mesh_files:
        print(f"  - {os.path.basename(mf)}")

    # Convert all to XML
    xml_paths = []
    for mesh_path in mesh_files:
        print(f"\nConverting {os.path.basename(mesh_path)} to XML...")
        xml_path = mesh_to_xml(mesh_path, converter)
        if xml_path:
            xml_paths.append(xml_path)
        else:
            print(f"  SKIPPED (conversion failed)")

    if not xml_paths:
        print("No meshes could be converted to XML, aborting.")
        return 0

    # Parse all XMLs
    mesh_data = {}
    all_pose_names = set()
    for xml_path in xml_paths:
        name = os.path.basename(xml_path)[:-9]  # remove .mesh.xml
        tree, vertices, poses = parse_mesh_xml(xml_path)
        if vertices:
            mesh_data[name] = {
                'tree': tree,
                'vertices': vertices,
                'poses': poses,
                'xml_path': xml_path,
            }
            all_pose_names.update(poses.keys())
        else:
            print(f"WARNING: No vertices found in {name}")

    if not mesh_data:
        print("No mesh data could be parsed, aborting.")
        return 0

    # Build global position -> list of (normal, tangent) map for BASIS
    pos_map = {}
    for name, data in mesh_data.items():
        for v in data['vertices']:
            key = round_pos(v['pos'], 0.0001)
            if key not in pos_map:
                pos_map[key] = []
            pos_map[key].append((v['normal'], v['tangent']))

    # Compute averaged normals and tangents for positions appearing in 2+ meshes
    avg_data = {}
    multi_mesh_positions = 0
    for pos_key, entries in pos_map.items():
        if len(entries) < 2:
            continue
        multi_mesh_positions += 1

        # Average normals
        avg_n = [0.0, 0.0, 0.0]
        for n, t in entries:
            avg_n[0] += n[0]
            avg_n[1] += n[1]
            avg_n[2] += n[2]
        avg_n = vector_normalize(avg_n)

        # Average tangents (only if at least one entry has a tangent)
        avg_t = [0.0, 0.0, 0.0]
        tangent_count = 0
        for n, t in entries:
            if t is not None:
                avg_t[0] += t[0]
                avg_t[1] += t[1]
                avg_t[2] += t[2]
                tangent_count += 1

        if tangent_count > 0:
            avg_t = vector_normalize(avg_t)
        else:
            avg_t = None

        avg_data[pos_key] = {
            'normal': avg_n,
            'tangent': avg_t,
        }

    print(f"\nFound {multi_mesh_positions} positions shared across 2+ meshes")
    print(f"Averaging normals/tangents for {len(avg_data)} positions...")

    # Build global position -> list of pose normal offsets for EACH pose
    # pose_name -> {pos_key -> list of (nx, ny, nz)}
    pose_pos_maps = {}
    for pose_name in all_pose_names:
        pose_pos_map = {}
        for name, data in mesh_data.items():
            if pose_name not in data['poses']:
                continue
            pose_offsets = data['poses'][pose_name]
            for v in data['vertices']:
                if v['index'] not in pose_offsets:
                    continue
                key = round_pos(v['pos'], 0.0001)
                if key not in pose_pos_map:
                    pose_pos_map[key] = []
                pose_pos_map[key].append(pose_offsets[v['index']])
        pose_pos_maps[pose_name] = pose_pos_map

    # Compute averaged pose normal offsets for positions appearing in 2+ meshes
    avg_pose_data = {}
    for pose_name, pose_pos_map in pose_pos_maps.items():
        avg_offsets = {}
        for pos_key, entries in pose_pos_map.items():
            if len(entries) < 2:
                continue
            avg_n = [0.0, 0.0, 0.0]
            for nx, ny, nz in entries:
                avg_n[0] += nx
                avg_n[1] += ny
                avg_n[2] += nz
            avg_n = vector_normalize(avg_n)
            avg_offsets[pos_key] = avg_n
        if avg_offsets:
            avg_pose_data[pose_name] = avg_offsets

    if avg_pose_data:
        print(f"Averaging pose normal offsets for {len(avg_pose_data)} pose(s):")
        for pose_name in sorted(avg_pose_data.keys()):
            print(f"  - {pose_name}: {len(avg_pose_data[pose_name])} positions")
    else:
        print("No pose normal offsets to fix.")

    # Update all meshes with averaged values
    modified_count = 0
    for name, data in mesh_data.items():
        tree = data['tree']
        vertices = data['vertices']
        modified = False
        normal_fixes = 0
        tangent_fixes = 0
        pose_normal_fixes = 0

        root = tree.getroot()
        sharedgeom = root.find('.//sharedgeometry')
        if sharedgeom is None:
            continue
        vbs = sharedgeom.findall('vertexbuffer')
        if len(vbs) < 1:
            continue

        pos_norm_vb = vbs[0]
        tt_verts = vbs[1].findall('vertex') if len(vbs) > 1 else []
        pn_verts = pos_norm_vb.findall('vertex')

        # Fix basis normals and tangents
        for v in vertices:
            key = round_pos(v['pos'], 0.0001)
            if key not in avg_data:
                continue

            avg = avg_data[key]

            # Update normal
            pnv = pn_verts[v['index']]
            normal = pnv.find('normal')
            if normal is not None:
                old_n = (float(normal.get('x')), float(normal.get('y')), float(normal.get('z')))
                new_n = avg['normal']
                diff = math.sqrt(
                    (old_n[0]-new_n[0])**2 +
                    (old_n[1]-new_n[1])**2 +
                    (old_n[2]-new_n[2])**2
                )
                if diff > 0.000001:
                    normal.set('x', f"{new_n[0]:.6f}")
                    normal.set('y', f"{new_n[1]:.6f}")
                    normal.set('z', f"{new_n[2]:.6f}")
                    normal_fixes += 1

            # Update tangent
            if avg['tangent'] is not None and v['tangent'] is not None:
                if v['index'] < len(tt_verts):
                    ttv = tt_verts[v['index']]
                    tangent = ttv.find('tangent')
                    if tangent is not None:
                        old_t = (float(tangent.get('x')), float(tangent.get('y')), float(tangent.get('z')))
                        new_t = avg['tangent']
                        diff = math.sqrt(
                            (old_t[0]-new_t[0])**2 +
                            (old_t[1]-new_t[1])**2 +
                            (old_t[2]-new_t[2])**2
                        )
                        if diff > 0.000001:
                            tangent.set('x', f"{new_t[0]:.6f}")
                            tangent.set('y', f"{new_t[1]:.6f}")
                            tangent.set('z', f"{new_t[2]:.6f}")
                            tangent_fixes += 1

            modified = True

        # Fix pose normal offsets
        for pose_name, avg_offsets in avg_pose_data.items():
            if pose_name not in data['poses']:
                continue
            # Find the <pose> element
            pose_elem = None
            for p in root.findall('.//poses/pose'):
                if p.get('name') == pose_name:
                    pose_elem = p
                    break
            if pose_elem is None:
                continue

            for v in vertices:
                key = round_pos(v['pos'], 0.0001)
                if key not in avg_offsets:
                    continue
                if v['index'] not in data['poses'][pose_name]:
                    continue

                # Find the poseoffset element for this vertex
                for po in pose_elem.findall('poseoffset'):
                    if int(po.get('index')) == v['index']:
                        old_nx = float(po.get('nx', 0.0))
                        old_ny = float(po.get('ny', 0.0))
                        old_nz = float(po.get('nz', 0.0))
                        new_n = avg_offsets[key]
                        diff = math.sqrt(
                            (old_nx-new_n[0])**2 +
                            (old_ny-new_n[1])**2 +
                            (old_nz-new_n[2])**2
                        )
                        if diff > 0.000001:
                            po.set('nx', f"{new_n[0]:.6f}")
                            po.set('ny', f"{new_n[1]:.6f}")
                            po.set('nz', f"{new_n[2]:.6f}")
                            pose_normal_fixes += 1
                        break

            if pose_normal_fixes > 0:
                modified = True

        if modified:
            xml_path = data['xml_path']
            tree.write(xml_path, encoding='utf-8', xml_declaration=True)
            print(f"  {name}: fixed {normal_fixes} normals, {tangent_fixes} tangents, {pose_normal_fixes} pose normals")
            modified_count += 1
        else:
            print(f"  {name}: no changes needed")

    # Convert back to binary
    print(f"\nConverting {modified_count} modified meshes back to binary...")
    converted = 0
    for xml_path in xml_paths:
        mesh_path = xml_to_mesh(xml_path, converter)
        if mesh_path:
            converted += 1
        # Clean up XML file
        try:
            os.remove(xml_path)
        except OSError:
            pass

    print(f"[OGRE Mesh Seam Fix] Done. {converted}/{len(xml_paths)} meshes converted.")
    return converted


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: fix_ogre_mesh_seams.py <mesh_directory> [ogre_xml_converter_path]")
        sys.exit(1)

    mesh_dir = sys.argv[1]
    converter = find_converter()

    if not converter:
        print("ERROR: OgreXMLConverter not found. Set OGRETOOLS_XML_CONVERTER env var or pass path as 2nd argument.")
        sys.exit(1)

    if not os.path.isdir(mesh_dir):
        print(f"ERROR: {mesh_dir} is not a directory")
        sys.exit(1)

    fix_seams(mesh_dir, converter)
