import json
import logging
from soagila.db.sqlite_init import get_db
from soagila.core.idl.receipt import DecisionReceiptModel

logger = logging.getLogger("soagila.idl.trace")


async def forward_trace(receipt_id: str, max_depth: int = 10) -> list[DecisionReceiptModel]:
    """
    Forward trace: what decisions resulted from this one?
    Follows the Merkle chain forward from the given receipt.
    Also finds receipts that accessed memories created by the source receipt.
    """
    source = await _get_receipt_by_id(receipt_id)
    if not source:
        return []

    results = []
    db = await get_db()
    rows = await db.execute(
        """
            SELECT id, session_id, user_id, org_id, receipt_type,
                   memories_accessed, memories_created, arbitration_result,
                   context_hash, response_hash, merkle_hash, prev_merkle_hash,
                   metadata, created_at
            FROM decision_receipts
            WHERE created_at > ?
            ORDER BY created_at ASC
            LIMIT ?
        """,
        (source.created_at.isoformat(), max_depth),
    )
    rows = await rows.fetchall()

    for row in rows:
        receipt = _row_to_receipt(row)
        # Check if this receipt accessed any memories created by source
        if source.memories_created:
            accessed = set(receipt.memories_accessed)
            created = set(source.memories_created)
            if accessed & created:
                receipt.metadata["trace_connection"] = "direct_memory_link"
            else:
                receipt.metadata["trace_connection"] = "temporal_chain"
        results.append(receipt)

    return results


async def backward_trace(receipt_id: str, max_depth: int = 10) -> list[DecisionReceiptModel]:
    """
    Backward trace: what led to this decision?
    Follows the Merkle chain backward from the given receipt.
    Also finds receipts that created memories accessed by the source receipt.
    """
    source = await _get_receipt_by_id(receipt_id)
    if not source:
        return []

    results = []
    db = await get_db()
    rows = await db.execute(
        """
            SELECT id, session_id, user_id, org_id, receipt_type,
                   memories_accessed, memories_created, arbitration_result,
                   context_hash, response_hash, merkle_hash, prev_merkle_hash,
                   metadata, created_at
            FROM decision_receipts
            WHERE created_at < ?
            ORDER BY created_at DESC
            LIMIT ?
        """,
        (source.created_at.isoformat(), max_depth),
    )
    rows = await rows.fetchall()

    for row in rows:
        receipt = _row_to_receipt(row)
        # Check if this receipt created any memories accessed by source
        if source.memories_accessed:
            accessed = set(source.memories_accessed)
            created = set(receipt.memories_created)
            if accessed & created:
                receipt.metadata["trace_connection"] = "direct_memory_link"
            else:
                receipt.metadata["trace_connection"] = "temporal_chain"
        results.append(receipt)

    # Reverse to chronological order
    results.reverse()
    return results


async def _get_receipt_by_id(receipt_id: str) -> DecisionReceiptModel | None:
    """Helper to get a receipt by ID."""
    db = await get_db()
    result = await db.execute(
        """
            SELECT id, session_id, user_id, org_id, receipt_type,
                   memories_accessed, memories_created, arbitration_result,
                   context_hash, response_hash, merkle_hash, prev_merkle_hash,
                   metadata, created_at
            FROM decision_receipts WHERE id = ?
        """,
        (receipt_id,),
    )
    row = await result.fetchone()

    if not row:
        return None
    return _row_to_receipt(row)


def _row_to_receipt(row) -> DecisionReceiptModel:
    """Convert a database row to a DecisionReceiptModel."""
    return DecisionReceiptModel(
        id=str(row[0]),
        session_id=str(row[1]) if row[1] else None,
        user_id=str(row[2]) if row[2] else None,
        org_id=str(row[3]) if row[3] else None,
        receipt_type=row[4],
        memories_accessed=json.loads(row[5]) if row[5] else [],
        memories_created=json.loads(row[6]) if row[6] else [],
        arbitration_result=json.loads(row[7]) if row[7] else None,
        context_hash=row[8],
        response_hash=row[9],
        merkle_hash=row[10],
        prev_merkle_hash=row[11],
        metadata=json.loads(row[12]) if row[12] else {},
        created_at=row[13],
    )
