from typing import Any
from loguru import logger
from sqlalchemy import select, inspect
from service_forge.workflow.node import Node
from service_forge.workflow.port import Port

class DatabaseUpdateNode(Node):
    DEFAULT_INPUT_PORTS = [
        Port("data", Any),
        Port("primary_keys", list[str] | None),
    ]

    DEFAULT_OUTPUT_PORTS = [
        Port("data", Any),
    ]

    def __init__(self, name: str):
        super().__init__(name)

    async def _run(self, data: Any, primary_keys: list[str] | None = None) -> None:
        if not hasattr(data, '__table__') or not hasattr(data, '__mapper__'):
            raise ValueError(f"data must be a SQLAlchemy model instance (inherited from declarative_base()), got {type(data)}")
        
        session_factory = await self.default_postgres_database.get_session_factory()
        async with session_factory() as session:
            try:
                model_class = type(data)
                mapper = inspect(model_class)
                
                if primary_keys:
                    pk_columns = []
                    for pk_name in primary_keys:
                        if not hasattr(model_class, pk_name):
                            raise ValueError(f"Field '{pk_name}' does not exist in model {model_class.__name__}")
                        pk_columns.append((pk_name, getattr(model_class, pk_name)))
                else:
                    pk_columns = [(pk_col.name, getattr(model_class, pk_col.name)) for pk_col in mapper.primary_key]
                
                conditions = []
                for pk_name, pk_attr in pk_columns:
                    pk_value = getattr(data, pk_name)
                    if pk_value is None:
                        raise ValueError(f"Primary key '{pk_name}' is None, cannot check existence")
                    conditions.append(pk_attr == pk_value)
                
                query = select(model_class).where(*conditions)
                result = await session.execute(query)
                existing_record = result.scalar_one_or_none()
                
                if existing_record is None:
                    error_msg = f"Record does not exist in {model_class.__tablename__} with the provided primary keys"
                    logger.error(f"✗ {error_msg}")
                    raise ValueError(error_msg)
                
                # Update all fields from data to existing_record, excluding primary keys and timestamp fields
                excluded_fields = {'created_at', 'updated_at'}
                pk_field_names = {pk_col[0] for pk_col in pk_columns}
                
                for column in mapper.columns:
                    if column.name not in pk_field_names and column.name not in excluded_fields:
                        setattr(existing_record, column.name, getattr(data, column.name))
                
                await session.commit()
                await session.refresh(existing_record)
                logger.info(f"✓ Updated record in {model_class.__tablename__}")
                self.activate_output_edges(self.get_output_port_by_name('data'), existing_record)
                    
            except Exception as e:
                logger.error(f"✗ Failed to update record: {e}")
                await session.rollback()
                raise