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 DatabaseCreateOrErrorNode(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:
                    session.add(data)
                    await session.commit()
                    await session.refresh(data)
                    logger.info(f"✓ Created new record in {model_class.__tablename__}")
                    self.activate_output_edges(self.get_output_port_by_name('data'), data)
                else:
                    error_msg = f"Record already exists in {model_class.__tablename__}"
                    raise ValueError(error_msg)
                    
            except Exception as e:
                logger.error(f"✗ Failed to create/retrieve record: {e}")
                await session.rollback()
                raise
