"""
AYON_register.py
================
Python Script TOP — runs once per work item.

Reads AYON entity attribs stamped by TOPs_2_AYONassets, discovers USD
files in the publish directory, then creates or updates the AYON
representation for the version. Stamps ayon_repre_id and status back
onto the work item.

Safe to re-run: existing representations are updated in place.
"""

import os
import sys
import uuid
import hashlib
import platform
import traceback

import ayon_api
from ayon_api.operations import OperationsSession
from ayon_core.pipeline import get_current_project_name
from ayon_core.pipeline.anatomy import Anatomy


# ================================================================
# CONFIGURATION
# ================================================================

PRODUCT_TYPE     = "usd"
REPRESENTATION   = "usd"
USD_EXTENSIONS   = {".usd", ".usda", ".usdc", ".usdz"}
HASH_ALGORITHM   = "md5"
ANATOMY_ROOT     = "work"


# ================================================================
# HELPERS
# ================================================================

def compute_file_hash(filepath):
    h = hashlib.new(HASH_ALGORITHM)
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def make_rootless_path(abs_path, root_value):
    abs_path   = abs_path.replace("\\", "/")
    root_value = root_value.replace("\\", "/").rstrip("/")
    if abs_path.lower().startswith(root_value.lower()):
        return "{root[%s]}%s" % (ANATOMY_ROOT, abs_path[len(root_value):])
    return abs_path


def build_file_entry(abs_path, root_value):
    return {
        "id":       uuid.uuid4().hex,
        "name":     os.path.basename(abs_path),
        "path":     make_rootless_path(abs_path, root_value),
        "size":     os.path.getsize(abs_path),
        "hash":     compute_file_hash(abs_path),
        "hashType": HASH_ALGORITHM,
    }


def discover_usd_files(publish_dir):
    """Return all USD files in publish_dir, with the primary .usd listed first."""
    if not os.path.isdir(publish_dir):
        raise RuntimeError("Publish directory not found on disk: %s" % publish_dir)

    all_files = sorted(
        os.path.join(publish_dir, f).replace("\\", "/")
        for f in os.listdir(publish_dir)
        if os.path.splitext(f)[1].lower() in USD_EXTENSIONS
        and os.path.isfile(os.path.join(publish_dir, f))
    )

    if not all_files:
        raise RuntimeError("No USD files found in: %s" % publish_dir)

    # Prefer a .usd as the primary file
    main_files = [f for f in all_files if f.endswith(".usd")]
    primary    = main_files[0] if main_files else all_files[0]
    return primary, all_files


def resolve_anatomy(project_name, folder_path, product_name, version_num):
    """Return (publish_dir, root_value) via AYON Anatomy template."""
    anatomy      = Anatomy(project_name)
    folder_parts = folder_path.strip("/").split("/")
    project_data = ayon_api.get_project(project_name)

    template_data = {
        "root":           anatomy.roots,
        "project":        {"name": project_name, "code": project_data.get("code", project_name)},
        "folder":         {"name": folder_parts[-1], "path": folder_path},
        "hierarchy":      "/".join(folder_parts[:-1]),
        "product":        {"name": product_name, "type": PRODUCT_TYPE},
        "version":        version_num,
        "representation": REPRESENTATION,
        "ext":            REPRESENTATION,
        "user":           "",
        "comment":        "",
    }

    result      = anatomy.get_template_item("publish", "default", "directory").format_strict(template_data)
    publish_dir = result.normalized().replace("\\", "/")
    root_value  = str(anatomy.roots[ANATOMY_ROOT]).replace("\\", "/").rstrip("/")
    return publish_dir, root_value


def build_context_data(project_name, project_code, folder_path, product_name, version_num, root_value):
    folder_parts = folder_path.strip("/").split("/")
    return {
        "project":        {"name": project_name, "code": project_code},
        "folder":         {"name": folder_parts[-1], "path": folder_path},
        "hierarchy":      "/".join(folder_parts[:-1]),
        "product":        {"name": product_name, "type": PRODUCT_TYPE},
        "version":        version_num,
        "representation": REPRESENTATION,
        "root":           {ANATOMY_ROOT: root_value},
        "ext":            REPRESENTATION,
    }


def create_or_update_representation(project_name, version_id, files_list, attrib, context_data):
    """Create a new representation or update the existing one in place."""
    existing = ayon_api.get_representation_by_name(
        project_name, REPRESENTATION, version_id,
        fields={"id", "attrib", "context", "files", "active"},
    )

    if existing:
        session = OperationsSession()
        try:
            session.update_representation(
                project_name, existing["id"],
                attrib=attrib,
                context=context_data,
                files=files_list,
                active=True,
            )
        except TypeError:
            # Older ayon_api versions don't support all kwargs — fall back to PATCH
            ayon_api.get_server_api_connection().patch(
                "projects/%s/representations/%s" % (project_name, existing["id"]),
                attrib=attrib,
                context=context_data,
                files=files_list,
                active=True,
            )
        session.commit()
        print("  updated existing representation (id: %s)" % existing["id"])
        return existing["id"], "updated"

    try:
        repre_id = ayon_api.create_representation(
            project_name,
            name=REPRESENTATION,
            version_id=version_id,
            files=files_list,
            attrib=attrib,
            context=context_data,
            data={"context": context_data},
            tags=[],
            status=None,
            active=True,
        )
    except TypeError:
        # Older ayon_api — fall back to POST
        repre_id = uuid.uuid4().hex
        ayon_api.get_server_api_connection().post(
            "projects/%s/representations" % project_name,
            id=repre_id,
            name=REPRESENTATION,
            versionId=version_id,
            files=files_list,
            attrib=attrib,
            context=context_data,
            data={"context": context_data},
            active=True,
        )
    print("  created representation (id: %s)" % repre_id)
    return repre_id, "created"


# ================================================================
# MAIN
# ================================================================

try:
    project_name = get_current_project_name()
    folder_path  = work_item.stringAttribValue("ayon_folder_path")
    product_name = work_item.stringAttribValue("ayon_product_name")
    version_id   = work_item.stringAttribValue("ayon_version_id")
    version_num  = work_item.intAttribValue("ayon_version")
    asset_name   = work_item.stringAttribValue("asset_name")

    print("\n%s" % ("=" * 60))
    print("  project : %s" % project_name)
    print("  asset   : %s" % asset_name)
    print("  product : %s" % product_name)
    print("  version : v%03d" % version_num)
    print("=" * 60)

    if not version_id:
        raise RuntimeError("Missing ayon_version_id — did TOPs_2_AYONassets run successfully?")

    print("\n[1] Resolving anatomy")
    publish_dir, root_value = resolve_anatomy(project_name, folder_path, product_name, version_num)
    print("  %s" % publish_dir)

    print("\n[2] Discovering USD files")
    primary_file, all_files = discover_usd_files(publish_dir)
    print("  primary : %s" % os.path.basename(primary_file))
    print("  total   : %d file(s)" % len(all_files))

    primary_rootless = make_rootless_path(primary_file, root_value)
    files_list       = [build_file_entry(f, root_value) for f in all_files]

    attrib = {
        "template": primary_rootless,
        "path":     primary_rootless,
    }

    print("\n[3] Building context")
    project_data  = ayon_api.get_project(project_name)
    context_data  = build_context_data(
        project_name, project_data.get("code", project_name),
        folder_path, product_name, version_num, root_value,
    )

    print("\n[4] Registering representation")
    repre_id, status = create_or_update_representation(
        project_name, version_id, files_list, attrib, context_data
    )

    work_item.setStringAttrib("ayon_repre_id",   repre_id)
    work_item.setStringAttrib("ayon_reg_status", status)
    work_item.setStringAttrib("ayon_reg_error",  "")

    print("\n%s" % ("=" * 60))
    print("  DONE  %s  |  %s  |  %s" % (asset_name, status, repre_id))
    print("=" * 60 + "\n")

except Exception as e:
    print("\n%s\nSCRIPT FAILED\n%s" % ("=" * 60, "=" * 60))
    print(traceback.format_exc())

    try:
        work_item.setStringAttrib("ayon_reg_status", "failed")
        work_item.setStringAttrib("ayon_reg_error",  str(e)[:500])
        work_item.setStringAttrib("ayon_repre_id",   "")
    except Exception:
        pass

    raise