database_manager.py

from sqlmodel import SQLModel, create_engine, Session, select
import os
from datetime import datetime
from typing import Optional
from db_models import Publication, Hazard, RFEvent, Condition, Impact, ExternalDataset, SupportingDocument, ConfigParams, WorkflowLog, SearchQuery, Screening

DB_PATH = "mudcap_review.db"
_engine = None

def get_engine():
    global _engine
    if _engine is None:
        sqlite_url = f"sqlite:///{DB_PATH}"
        _engine = create_engine(sqlite_url)
    return _engine

def set_db_path(path):
    global DB_PATH, _engine
    DB_PATH = path
    # Reset engine if path changes
    _engine = None
    get_engine()

def initialize_database(force=False):
    global _engine
    if force and os.path.exists(DB_PATH):
        print(f"Forcing re-initialization of {DB_PATH}")
        # Dispose engine and set to None to release file handle
        if _engine:
            _engine.dispose()
            _engine = None
        # Try to delete file
        try:
            os.remove(DB_PATH)
        except PermissionError:
            print(f"Warning: Could not delete {DB_PATH}. It might be held by another process.")
    
    SQLModel.metadata.create_all(get_engine())
    print(f"Database initialized at {DB_PATH}")

def log_config(config_dict, version="1.0"):
    with Session(get_engine()) as session:
        config = ConfigParams(config_data=config_dict, version=version)
        session.add(config)
        session.commit()

def log_workflow(step_name, status, details=None, agent="ai"):
    with Session(get_engine()) as session:
        log = WorkflowLog(step_name=step_name, status=status, details=details, agent=agent)
        session.add(log)
        session.commit()

def log_search_query(database_name, query_text, parameters=None):
    with Session(get_engine()) as session:
        query = SearchQuery(database_name=database_name, query_text=query_text, parameters=parameters)
        session.add(query)
        session.commit()

def save_publication(pub_metadata):
    """
    Saves a PublicationMetadata object (from data_collector.base) to the database.
    Checks for duplicates based on DOI.
    If a duplicate is found, it is linked to the original.
    """
    with Session(get_engine()) as session:
        # Extract DOI
        doi = pub_metadata.doi
        
        duplicate_of = None
        # Check for duplicates if DOI exists
        if doi:
            # Find the original record if it exists
            # Note: We now look for any publication with this DOI.
            # If multiple exist, we just link to the first one found.
            existing = session.exec(
                select(Publication)
                .where(Publication.doi == doi)
                .limit(1)
            ).first()
            
            if existing:
                print(f"Duplicate DOI found: {doi}. Linking to existing ID {existing.id}")
                duplicate_of = existing.id
        
        db_pub = Publication(
            doi=doi,
            metadata_json=pub_metadata.model_dump(),
            duplicate_of_id=duplicate_of
        )
        session.add(db_pub)
        session.commit()
        session.refresh(db_pub)
        return db_pub

def add_screening(publication_id, operator_name, status, operator_type="human",
                  flood_event_relevance=None, geographic_relevance=None,
                  muddy_flood_relevance=None,
                  flood_confidence=None, flood_reasoning=None,
                  geo_confidence=None, geo_reasoning=None,
                  muddy_flood_confidence=None, muddy_flood_reasoning=None):
    """Adds or updates a screening decision for a specific operator and publication."""
    with Session(get_engine()) as session:
        existing = session.exec(
            select(Screening)
            .where(Screening.publication_id == publication_id)
            .where(Screening.operator_name == operator_name)
        ).first()

        if existing:
            existing.status = status
            existing.timestamp = datetime.now()
            existing.flood_event_relevance = flood_event_relevance
            existing.geographic_relevance = geographic_relevance
            existing.muddy_flood_relevance = muddy_flood_relevance
            existing.flood_confidence = flood_confidence
            existing.flood_reasoning = flood_reasoning
            existing.geo_confidence = geo_confidence
            existing.geo_reasoning = geo_reasoning
            existing.muddy_flood_confidence = muddy_flood_confidence
            existing.muddy_flood_reasoning = muddy_flood_reasoning
            session.add(existing)
        else:
            new_screening = Screening(
                publication_id=publication_id,
                operator_name=operator_name,
                operator_type=operator_type,
                status=status,
                flood_event_relevance=flood_event_relevance,
                geographic_relevance=geographic_relevance,
                muddy_flood_relevance=muddy_flood_relevance,
                flood_confidence=flood_confidence,
                flood_reasoning=flood_reasoning,
                geo_confidence=geo_confidence,
                geo_reasoning=geo_reasoning,
                muddy_flood_confidence=muddy_flood_confidence,
                muddy_flood_reasoning=muddy_flood_reasoning,
            )
            session.add(new_screening)

        session.commit()


def get_screenings_for_publications(pub_ids: list) -> dict:
    """Returns {pub_id: [Screening, ...]} for all given publication IDs."""
    if not pub_ids:
        return {}
    with Session(get_engine()) as session:
        screenings = session.exec(
            select(Screening).where(Screening.publication_id.in_(pub_ids))
        ).all()
        result: dict = {}
        for s in screenings:
            result.setdefault(s.publication_id, []).append(s)
        return result


def get_included_publications(operator_name: str) -> list:
    """Returns publications where the given operator marked both flood and geographic relevance as True."""
    with Session(get_engine()) as session:
        included_ids = session.exec(
            select(Screening.publication_id)
            .where(Screening.operator_name == operator_name)
            .where(Screening.flood_event_relevance == True)
            .where(Screening.geographic_relevance == True)
        ).all()
        if not included_ids:
            return []
        pubs = session.exec(
            select(Publication).where(Publication.id.in_(included_ids))
        ).all()
        return pubs


def get_next_unscreened_publication(operator_name: str, offset: int = 0) -> Optional[Publication]:
    """Returns the next unscreened publication for an operator, skipping duplicates."""
    with Session(get_engine()) as session:
        statement = (
            select(Publication)
            .where(Publication.duplicate_of_id == None)
            .where(~Publication.id.in_(
                select(Screening.publication_id).where(Screening.operator_name == operator_name)
            ))
            .offset(offset)
            .limit(1)
        )
        result = session.exec(statement).first()
        if result:
            # Access metadata_json while session is open to avoid DetachedInstanceError
            _ = result.metadata_json
        return result


def get_publications_page(status_filter: str, op_name: str, limit: int = 20, offset: int = 0) -> list:
    """Returns a page of publications for review modes."""
    with Session(get_engine()) as session:
        if status_filter == "Relevant":
            statement = (
                select(Publication)
                .join(Screening)
                .where(Screening.operator_name == op_name)
                .where(Screening.status == "relevant")
                .offset(offset)
                .limit(limit)
            )
        elif status_filter == "Irrelevant":
            statement = (
                select(Publication)
                .join(Screening)
                .where(Screening.operator_name == op_name)
                .where(Screening.status == "irrelevant")
                .offset(offset)
                .limit(limit)
            )
        elif status_filter == "Duplicates":
            statement = (
                select(Publication)
                .where(Publication.duplicate_of_id != None)
                .offset(offset)
                .limit(limit)
            )
        else:  # All
            statement = select(Publication).offset(offset).limit(limit)
        results = session.exec(statement).all()
        # Access metadata_json while session is open
        for r in results:
            _ = r.metadata_json
        return results


def get_all_screenings_with_publications(operator_name: Optional[str] = None):
    """
    Returns all screening records joined with their corresponding publication data.
    If operator_name is provided, filters by that operator.
    Returns a list of tuples (Screening, Publication).
    """
    with Session(get_engine()) as session:
        statement = select(Screening, Publication).join(
            Publication, Screening.publication_id == Publication.id
        )
        if operator_name:
            statement = statement.where(Screening.operator_name == operator_name)

        results = session.exec(statement).all()
        # Ensure metadata is loaded for all publications
        for _, pub in results:
            _ = pub.metadata_json
        return results


def get_extraction(pub_id: int) -> Optional[dict]:
    """Returns existing extraction data for a publication, or None if not extracted yet."""
    with Session(get_engine()) as session:
        hazard = session.exec(
            select(Hazard).where(Hazard.publication_id == pub_id)
        ).first()
        if not hazard:
            return None
        rf_event = session.exec(
            select(RFEvent).where(RFEvent.hazard_id == hazard.id)
        ).first()
        conditions = session.exec(
            select(Condition).where(Condition.hazard_id == hazard.id)
        ).all()
        impacts = session.exec(
            select(Impact).where(Impact.hazard_id == hazard.id)
        ).all()
        return {
            "hazard": hazard,
            "rf_event": rf_event,
            "conditions": list(conditions),
            "impacts": list(impacts),
        }


def save_extraction(pub_id: int, hazard_data: dict, rfevent_data: Optional[dict],
                    conditions: list, impacts: list) -> int:
    """Saves extraction data for a publication. Replaces any existing extraction. Returns hazard ID."""
    with Session(get_engine()) as session:
        # Remove existing extraction for this publication
        old_hazard = session.exec(
            select(Hazard).where(Hazard.publication_id == pub_id)
        ).first()
        if old_hazard:
            session.exec(
                select(RFEvent).where(RFEvent.hazard_id == old_hazard.id)
            )
            for rf in session.exec(select(RFEvent).where(RFEvent.hazard_id == old_hazard.id)).all():
                session.delete(rf)
            for cond in session.exec(select(Condition).where(Condition.hazard_id == old_hazard.id)).all():
                session.delete(cond)
            for imp in session.exec(select(Impact).where(Impact.hazard_id == old_hazard.id)).all():
                session.delete(imp)
            session.delete(old_hazard)
            session.flush()

        hazard = Hazard(publication_id=pub_id, **hazard_data)
        session.add(hazard)
        session.flush()

        if rfevent_data:
            rfevent = RFEvent(hazard_id=hazard.id, **rfevent_data)
            session.add(rfevent)

        for cond in conditions:
            session.add(Condition(hazard_id=hazard.id, **cond))

        for imp in impacts:
            session.add(Impact(hazard_id=hazard.id, **imp))

        pub = session.get(Publication, pub_id)
        if pub:
            pub.manual_extraction = True
            pub.has_events = True
            session.add(pub)

        session.commit()
        return hazard.id