"""
BLOOM LIMS - Subject Management Module

This module provides helper functions for managing Subjects in BLOOM LIMS.
Subjects are logical aggregates that decisions apply to, spanning multiple objects.

Key concepts:
- Object (fact): A concrete thing (sample, container, fileset, etc.)
- Subject (decision scope): A logical aggregate that decisions apply to
- Anchor: The primary object that defines the subject
- Members: Additional objects associated with the subject

Relationship types:
- subject_anchor: Subject → Primary anchor object
- subject_member: Subject → Member/evidence objects
"""

import logging
from typing import Optional, List, Dict, Any

logger = logging.getLogger(__name__)

# Relationship type constants
RELATIONSHIP_SUBJECT_ANCHOR = "subject_anchor"
RELATIONSHIP_SUBJECT_MEMBER = "subject_member"

# Template mapping for subject kinds
SUBJECT_TEMPLATE_MAP = {
    "analysis_bundle": "subject/generic/analysis-bundle-subject/1.0",
    "report": "subject/generic/report-subject/1.0",
    "generic": "subject/generic/generic-subject/1.0",
    # Workset subject templates
    "workset": "subject/generic/generic-subject/1.0",
    "workset_extraction": "subject/workset/extraction/1.0",
    "workset_sequencing": "subject/workset/sequencing/1.0",
}


def generate_subject_key(subject_kind: str, anchor_euid: str) -> str:
    """
    Generate a stable, deterministic subject key.
    
    Args:
        subject_kind: The kind of subject (accession, analysis_bundle, report, generic)
        anchor_euid: The EUID of the anchor object
        
    Returns:
        A stable subject key in format "{subject_kind}:{anchor_euid}"
    """
    return f"{subject_kind}:{anchor_euid}"


def find_subject_by_key(bob, subject_key: str):
    """
    Find a subject by its stable subject_key.
    
    Args:
        bob: BloomObj instance
        subject_key: The subject_key to search for
        
    Returns:
        The subject instance if found, None otherwise
    """
    try:
        results = bob.session.query(bob.Base.classes.generic_instance).filter(
            bob.Base.classes.generic_instance.category == "subject",
            bob.Base.classes.generic_instance.is_deleted == False,
        ).all()
        
        for result in results:
            props = result.json_addl.get("properties", {})
            if props.get("subject_key") == subject_key:
                return result
        
        return None
    except Exception as e:
        logger.error(f"Error finding subject by key {subject_key}: {e}")
        return None


def get_subject_template_euid(bob, subject_kind: str) -> Optional[str]:
    """
    Get the template EUID for a given subject kind.
    
    Args:
        bob: BloomObj instance
        subject_kind: The kind of subject
        
    Returns:
        The template EUID if found, None otherwise
    """
    template_path = SUBJECT_TEMPLATE_MAP.get(subject_kind, SUBJECT_TEMPLATE_MAP["generic"])
    parts = template_path.split("/")
    if len(parts) != 4:
        logger.error(f"Invalid template path: {template_path}")
        return None
    
    category, type_name, subtype, version = parts

    try:
        templates = bob.query_template_by_component_v2(
            category=category,
            type=type_name,
            subtype=subtype,
            version=version
        )
        if templates:
            return templates[0].euid
        return None
    except Exception as e:
        logger.error(f"Error getting subject template: {e}")
        return None


def create_subject(
    bob,
    anchor_euid: str,
    subject_kind: str,
    subject_key: Optional[str] = None,
    extra_props: Dict[str, Any] = None,
    template_euid: Optional[str] = None,
) -> Optional[str]:
    """
    Create or find existing subject. Idempotent - returns existing if found.
    
    Args:
        bob: BloomObj instance
        anchor_euid: EUID of the anchor object
        subject_kind: Kind of subject (accession, analysis_bundle, report, generic)
        subject_key: Optional custom subject key. Defaults to "{subject_kind}:{anchor_euid}"
        extra_props: Additional properties to set on the subject
        template_euid: Optional template EUID override
        
    Returns:
        Subject EUID if successful, None otherwise
    """
    if extra_props is None:
        extra_props = {}
    
    # Generate subject_key if not provided
    if subject_key is None:
        subject_key = generate_subject_key(subject_kind, anchor_euid)
    
    # Check for existing subject (idempotency)
    existing = find_subject_by_key(bob, subject_key)
    if existing:
        logger.info(f"Subject already exists with key {subject_key}: {existing.euid}")
        return existing.euid
    
    # Get or validate template
    if template_euid is None:
        template_euid = get_subject_template_euid(bob, subject_kind)
    
    if template_euid is None:
        logger.error(f"No template found for subject kind: {subject_kind}")
        return None
    
    try:
        # Create the subject instance
        result = bob.create_instances(template_euid)
        if not result or not result[0]:
            logger.error("Failed to create subject instance")
            return None

        subject = result[0][0]

        # Update properties
        props = subject.json_addl.get("properties", {})
        props["subject_kind"] = subject_kind
        props["subject_key"] = subject_key
        props["anchor_euid"] = anchor_euid
        props["status"] = extra_props.get("status", "active")
        props["name"] = extra_props.get("name", f"Subject: {subject_key}")

        # Merge extra properties
        for key, value in extra_props.items():
            if key not in ["subject_kind", "subject_key", "anchor_euid"]:
                props[key] = value

        subject.json_addl["properties"] = props
        from sqlalchemy.orm.attributes import flag_modified
        flag_modified(subject, "json_addl")

        bob.session.commit()

        # Create anchor relationship
        link_subject_anchor(bob, subject.euid, anchor_euid)

        logger.info(f"Created subject {subject.euid} with key {subject_key}")
        return subject.euid

    except Exception as e:
        logger.error(f"Error creating subject: {e}")
        bob.session.rollback()
        return None


def link_subject_anchor(bob, subject_euid: str, anchor_euid: str) -> bool:
    """
    Create subject → anchor relationship edge.

    Args:
        bob: BloomObj instance
        subject_euid: EUID of the subject
        anchor_euid: EUID of the anchor object

    Returns:
        True if successful, False otherwise
    """
    try:
        # Check if relationship already exists
        subject = bob.get_by_euid(subject_euid)
        for lineage in subject.parent_of_lineages:
            if lineage.is_deleted:
                continue
            if (lineage.child_instance.euid == anchor_euid and
                lineage.relationship_type == RELATIONSHIP_SUBJECT_ANCHOR):
                logger.info(f"Anchor relationship already exists: {subject_euid} -> {anchor_euid}")
                return True

        bob.create_generic_instance_lineage_by_euids(
            subject_euid,
            anchor_euid,
            relationship_type=RELATIONSHIP_SUBJECT_ANCHOR
        )
        bob.session.commit()
        logger.info(f"Created anchor relationship: {subject_euid} -> {anchor_euid}")
        return True

    except Exception as e:
        logger.error(f"Error creating anchor relationship: {e}")
        bob.session.rollback()
        return False


def add_subject_members(bob, subject_euid: str, member_euids: List[str]) -> Dict[str, bool]:
    """
    Add subject → member relationship edges.

    Args:
        bob: BloomObj instance
        subject_euid: EUID of the subject
        member_euids: List of member object EUIDs

    Returns:
        Dict mapping member EUIDs to success status
    """
    results = {}
    subject = bob.get_by_euid(subject_euid)

    # Get existing member relationships
    existing_members = set()
    for lineage in subject.parent_of_lineages:
        if lineage.is_deleted:
            continue
        if lineage.relationship_type == RELATIONSHIP_SUBJECT_MEMBER:
            existing_members.add(lineage.child_instance.euid)

    for member_euid in member_euids:
        try:
            if member_euid in existing_members:
                logger.info(f"Member relationship already exists: {subject_euid} -> {member_euid}")
                results[member_euid] = True
                continue

            bob.create_generic_instance_lineage_by_euids(
                subject_euid,
                member_euid,
                relationship_type=RELATIONSHIP_SUBJECT_MEMBER
            )
            results[member_euid] = True
            logger.info(f"Created member relationship: {subject_euid} -> {member_euid}")

        except Exception as e:
            logger.error(f"Error adding member {member_euid}: {e}")
            results[member_euid] = False

    bob.session.commit()
    return results


def list_subjects_for_object(bob, object_euid: str) -> List[Dict[str, Any]]:
    """
    Find all subjects containing this object (as anchor or member).

    Args:
        bob: BloomObj instance
        object_euid: EUID of the object

    Returns:
        List of dicts with subject info: {euid, kind, role, status, subject_key}
    """
    subjects = []
    obj = bob.get_by_euid(object_euid)

    if not obj:
        return subjects

    # Templates don't have lineage relationships - only instances do
    if not hasattr(obj, 'child_of_lineages'):
        return subjects

    # Check child_of_lineages to find subjects where this object is a child
    for lineage in obj.child_of_lineages:
        if lineage.is_deleted:
            continue

        parent = lineage.parent_instance
        if parent.category != "subject":
            continue

        role = "unknown"
        if lineage.relationship_type == RELATIONSHIP_SUBJECT_ANCHOR:
            role = "anchor"
        elif lineage.relationship_type == RELATIONSHIP_SUBJECT_MEMBER:
            role = "member"

        props = parent.json_addl.get("properties", {})
        subjects.append({
            "euid": parent.euid,
            "kind": props.get("subject_kind", "unknown"),
            "role": role,
            "status": props.get("status", "unknown"),
            "subject_key": props.get("subject_key", ""),
            "name": props.get("name", parent.name),
        })

    return subjects


def list_members_for_subject(bob, subject_euid: str) -> Dict[str, List[Dict[str, Any]]]:
    """
    Return anchor and members for a subject.

    Args:
        bob: BloomObj instance
        subject_euid: EUID of the subject

    Returns:
        Dict with keys 'anchor' and 'members', each containing list of object info
    """
    result = {"anchor": [], "members": []}
    subject = bob.get_by_euid(subject_euid)

    if not subject:
        return result

    # Templates don't have lineage relationships - only instances do
    if not hasattr(subject, 'parent_of_lineages'):
        return result

    for lineage in subject.parent_of_lineages:
        if lineage.is_deleted:
            continue

        child = lineage.child_instance
        obj_info = {
            "euid": child.euid,
            "name": child.json_addl.get("properties", {}).get("name", child.name),
            "type": child.type,
            "subtype": child.subtype,
            "category": child.category,
        }

        if lineage.relationship_type == RELATIONSHIP_SUBJECT_ANCHOR:
            result["anchor"].append(obj_info)
        elif lineage.relationship_type == RELATIONSHIP_SUBJECT_MEMBER:
            result["members"].append(obj_info)

    return result


# =============================================================================
# Workset-specific functions
# =============================================================================


def create_workset(
    bob,
    anchor_euid: str,
    workset_type: str = "accession",
    workflow_euid: Optional[str] = None,
    executed_by: Optional[str] = None,
    name: Optional[str] = None,
) -> Optional[str]:
    """
    Create a workset subject for grouping workflow transaction objects.

    A workset groups workflow steps and objects created during a single
    transaction/batch operation (e.g., accessioning, extraction, sequencing).

    Args:
        bob: BloomObj instance
        anchor_euid: EUID of the anchor object (typically first workflow step)
        workset_type: Type of workset (accession, extraction, sequencing)
        workflow_euid: Optional EUID of the parent workflow
        executed_by: Optional username of who executed the workset
        name: Optional custom name for the workset

    Returns:
        Workset subject EUID if successful, None otherwise
    """
    from datetime import UTC, datetime

    # Map workset_type to template key
    template_key = f"workset_{workset_type}"
    if template_key not in SUBJECT_TEMPLATE_MAP:
        template_key = "workset"  # Fallback to default

    # Generate subject key
    subject_key = generate_subject_key("workset", anchor_euid)

    # Build extra properties
    now = datetime.now(UTC).isoformat().replace("+00:00", "Z")

    extra_props = {
        "status": "in_progress",
        "started_at": now,
        "workflow_euid": workflow_euid or "",
        "executed_by": executed_by or "",
    }

    if name:
        extra_props["name"] = name
    else:
        extra_props["name"] = f"Workset: {workset_type} - {anchor_euid}"

    # Get template EUID
    template_euid = get_subject_template_euid(bob, template_key)

    return create_subject(
        bob,
        anchor_euid=anchor_euid,
        subject_kind="workset",
        subject_key=subject_key,
        extra_props=extra_props,
        template_euid=template_euid,
    )


def complete_workset(bob, workset_euid: str, status: str = "complete") -> bool:
    """
    Mark a workset as complete and set completed_at timestamp.

    Args:
        bob: BloomObj instance
        workset_euid: EUID of the workset subject
        status: Final status (complete, failed, abandoned)

    Returns:
        True if successful, False otherwise
    """
    from datetime import UTC, datetime
    from sqlalchemy.orm.attributes import flag_modified

    if status not in ("complete", "failed", "abandoned"):
        logger.error(f"Invalid workset status: {status}")
        return False

    try:
        workset = bob.get_by_euid(workset_euid)
        if not workset:
            logger.error(f"Workset not found: {workset_euid}")
            return False

        if workset.category != "subject":
            logger.error(f"Object is not a subject: {workset_euid}")
            return False

        props = workset.json_addl.get("properties", {})
        if props.get("subject_kind") != "workset":
            logger.error(f"Subject is not a workset: {workset_euid}")
            return False

        # Update status and completed_at
        now = datetime.now(UTC).isoformat().replace("+00:00", "Z")

        props["status"] = status
        props["completed_at"] = now
        workset.json_addl["properties"] = props
        flag_modified(workset, "json_addl")

        bob.session.commit()
        logger.info(f"Workset {workset_euid} marked as {status}")
        return True

    except Exception as e:
        logger.error(f"Error completing workset {workset_euid}: {e}")
        bob.session.rollback()
        return False


def get_workset_by_anchor(bob, anchor_euid: str) -> Optional[Dict[str, Any]]:
    """
    Find a workset subject by its anchor EUID.

    Args:
        bob: BloomObj instance
        anchor_euid: EUID of the anchor object

    Returns:
        Dict with workset info if found, None otherwise
    """
    subject_key = generate_subject_key("workset", anchor_euid)
    workset = find_subject_by_key(bob, subject_key)

    if not workset:
        return None

    props = workset.json_addl.get("properties", {})
    return {
        "euid": workset.euid,
        "subject_key": props.get("subject_key", ""),
        "anchor_euid": props.get("anchor_euid", ""),
        "workflow_euid": props.get("workflow_euid", ""),
        "status": props.get("status", "unknown"),
        "started_at": props.get("started_at", ""),
        "completed_at": props.get("completed_at", ""),
        "executed_by": props.get("executed_by", ""),
        "name": props.get("name", workset.name),
    }


def list_worksets(
    bob,
    status: Optional[str] = None,
    workflow_euid: Optional[str] = None,
    limit: int = 100,
) -> List[Dict[str, Any]]:
    """
    List workset subjects with optional filters.

    Args:
        bob: BloomObj instance
        status: Filter by status (in_progress, complete, failed, abandoned)
        workflow_euid: Filter by workflow EUID
        limit: Maximum number of results

    Returns:
        List of workset info dicts
    """
    worksets = []

    try:
        query = bob.session.query(bob.Base.classes.generic_instance).filter(
            bob.Base.classes.generic_instance.category == "subject",
            bob.Base.classes.generic_instance.type == "workset",
            bob.Base.classes.generic_instance.is_deleted == False,
        )

        results = query.limit(limit).all()

        for workset in results:
            props = workset.json_addl.get("properties", {})

            # Apply filters
            if status and props.get("status") != status:
                continue
            if workflow_euid and props.get("workflow_euid") != workflow_euid:
                continue

            worksets.append({
                "euid": workset.euid,
                "subject_key": props.get("subject_key", ""),
                "anchor_euid": props.get("anchor_euid", ""),
                "workflow_euid": props.get("workflow_euid", ""),
                "status": props.get("status", "unknown"),
                "started_at": props.get("started_at", ""),
                "completed_at": props.get("completed_at", ""),
                "executed_by": props.get("executed_by", ""),
                "name": props.get("name", workset.name),
                "type": workset.type,
                "subtype": workset.subtype,
            })

        return worksets

    except Exception as e:
        logger.error(f"Error listing worksets: {e}")
        return []
