data_processor.py

"""
Data processing module for handling extracted Excel data
"""
from typing import List, Dict, Any
from database import db_manager
from models import TbArea, TbQuestaoBase, TbItemBase, TbAcaoBase
from sqlalchemy.orm import Session

class DataProcessor:
    """Processes extracted Excel data and saves to database"""
    
    def __init__(self):
        self.db_manager = db_manager
    
    def process_and_save_data(self, extracted_data: List[Dict[str, Any]]) -> bool:
        """Process extracted data and save to database"""
        try:
            session = self.db_manager.get_session()
            
            for record in extracted_data:
                # Process each record
                self._process_record(session, record)
            
            session.commit()
            session.close()
            print(f"Successfully processed and saved {len(extracted_data)} records")
            return True
            
        except Exception as e:
            print(f"Error processing data: {e}")
            if 'session' in locals():
                session.rollback()
                session.close()
            return False
    
    def _process_record(self, session: Session, record: Dict[str, Any]):
        """Process individual record and save to database following hierarchy: area -> question -> item -> action"""
        try:
            # Get or create Area
            area = self._get_or_create_area(session, record.get('area'))
            if not area:
                print(f"Skipping record {record.get('row_index', 'unknown')}: No area specified")
                return
            
            # Get or create Question
            question = self._get_or_create_question(
                session, 
                area.area_id, 
                record.get('questao'),
                record.get('questao_abreviada')
            )
            if not question:
                print(f"Skipping record {record.get('row_index', 'unknown')}: No question specified")
                return
            
            # Get or create Item
            item = self._get_or_create_item(
                session, 
                question.questao_base_id, 
                record.get('item_verificacao'),
                record.get('item_verificacao_abreviado')
            )
            if not item:
                print(f"Skipping record {record.get('row_index', 'unknown')}: No item specified")
                return
            
            # Get or create Action
            action = self._get_or_create_action(
                session, 
                item.item_base_id, 
                record.get('questao_auxiliar')
            )
            if not action:
                print(f"Skipping record {record.get('row_index', 'unknown')}: No action specified")
                return
            
        except Exception as e:
            print(f"Error processing record {record.get('row_index', 'unknown')}: {e}")
            raise
    
    def _get_or_create_area(self, session: Session, area_name: str) -> TbArea:
        """Get existing area or create new one based on uppercase area name"""
        if not area_name:
            return None
            
        # Use uppercase area name for uniqueness check
        area_upper = area_name.upper()
        area = session.query(TbArea).filter(TbArea.nome == area_upper).first()
        if not area:
            area = TbArea(nome=area_upper, status='ATIVO')
            session.add(area)
            session.flush()  # Get the ID
        return area
    
    def _get_or_create_question(self, session: Session, area_id: str, questao: str, questao_abreviada: str = None) -> TbQuestaoBase:
        """Get existing question or create new one based on questao_abreviada"""
        if not questao_abreviada or not area_id:
            return None
            
        # Try to find existing question by questao_abreviada and area
        question = session.query(TbQuestaoBase).filter(
            TbQuestaoBase.area_id == area_id,
            TbQuestaoBase.titulo == questao_abreviada
        ).first()
        
        if not question:
            question = TbQuestaoBase(
                area_id=area_id,
                titulo=questao_abreviada,
                descricao=questao,
                tipo_base='KATRU'
            )
            session.add(question)
            session.flush()
        return question
    
    def _get_or_create_item(self, session: Session, question_id: str, item_verificacao: str, item_verificacao_abreviado: str = None) -> TbItemBase:
        """Get existing item or create new one based on item_verificacao_abreviado"""
        if not item_verificacao_abreviado or not question_id:
            return None
            
        # Try to find existing item by item_verificacao_abreviado and question
        item = session.query(TbItemBase).filter(
            TbItemBase.questao_base_id == question_id,
            TbItemBase.titulo == item_verificacao_abreviado
        ).first()
        
        if not item:
            item = TbItemBase(
                questao_base_id=question_id,
                titulo=item_verificacao_abreviado,
                descricao=item_verificacao,
                tipo_base='KATRU'
            )
            session.add(item)
            session.flush()
        return item
    
    def _get_or_create_action(self, session: Session, item_id: str, questao_auxiliar: str) -> TbAcaoBase:
        """Get existing action or create new one based on questao_auxiliar"""
        if not questao_auxiliar or not item_id:
            return None
            
        # Try to find existing action by questao_auxiliar and item
        action = session.query(TbAcaoBase).filter(
            TbAcaoBase.item_base_id == item_id,
            TbAcaoBase.descricao == questao_auxiliar
        ).first()
        
        if not action:
            action = TbAcaoBase(
                item_base_id=item_id,
                titulo=questao_auxiliar[:200],  # Truncate to fit varchar(200)
                descricao=questao_auxiliar,
                tipo_base='KATRU'
            )
            session.add(action)
            session.flush()
        return action
    
    def get_data_summary(self) -> Dict[str, int]:
        """Get summary of data in database"""
        try:
            session = self.db_manager.get_session()
            
            summary = {
                'tb_area': session.query(TbArea).count(),
                'tb_questao_base': session.query(TbQuestaoBase).count(),
                'tb_item_base': session.query(TbItemBase).count(),
                'tb_acao_base': session.query(TbAcaoBase).count()
            }
            
            session.close()
            return summary
            
        except Exception as e:
            print(f"Error getting data summary: {e}")
            return {}