import traceback
import pickle
import dspy
from io import BytesIO
from celery import Celery
from celery.utils.log import get_task_logger
from sqlalchemy.orm import Session
from sqlalchemy.orm.attributes import flag_modified

from app.core import settings
from app.core.db import engine
from app.models import ImportStatus, Chunk
from app.knowledge_graph.graph_extractor import (
    SimpleGraphExtractor,
)

app = Celery(
    "knowledge_graph",
    broker=settings.CELERY_BROKER_URL,
    backend=settings.CELERY_RESULT_BACKEND,
)

logger = get_task_logger(__name__)


@app.task
def extract_graph_from_chunk(chunk_id: int):
    # reconfigure dspy to use a new instance of the OpenAI API client,
    # so that we can track usage per task
    if settings.GRAPH_EXTRACT_LLM_PROVIDER == "openai":
        oai_client = dspy.OpenAI(
            model="gpt-4o", api_key=settings.OPENAI_API_KEY, max_tokens=4096
        )
    else:
        oai_client = dspy.Google(
            model="gemini-1.5-flash",
            api_key=settings.GOOGLE_API_KEY,
            max_output_tokens=8192,
        )
    dspy.settings.configure(lm=oai_client)
    with Session(engine) as session:
        chunk = session.query(Chunk).get(chunk_id)

        if chunk.import_status != ImportStatus.PENDING:
            logger.info(
                f"Chunk {chunk_id}[status: {chunk.import_status}] already processed, skipping"
            )
            return

        chunk.import_status = ImportStatus.EXTRACTING_GRAPH
        session.commit()

        extractor = SimpleGraphExtractor()
        try:
            entities_df, relationships_df = extractor.extract(
                chunk.text,
                chunk.doc_id,
            )
        except Exception:
            chunk.import_status = ImportStatus.FAILED
            error_msg = f"Failed to extract graph: {traceback.format_exc()}"
            logger.error(error_msg)
            chunk.import_result = error_msg
            session.commit()
            return

        entities_df_bytes = BytesIO()
        entities_df.to_pickle(entities_df_bytes)
        entities_df_bytes.seek(0)
        relationships_df_bytes = BytesIO()
        relationships_df.to_pickle(relationships_df_bytes)
        relationships_df_bytes.seek(0)
        chunk.entities_df = entities_df_bytes.read()
        chunk.relationships_df = relationships_df_bytes.read()
        chunk.import_status = ImportStatus.EXTRACTED_GRAPH
        if settings.GRAPH_EXTRACT_LLM_PROVIDER == "openai":
            chunk.llm_usage = {"extract": dspy.settings.lm.history}
        session.add(chunk)
        session.commit()

        save_graph_to_db.delay(chunk_id)


@app.task
def save_graph_to_db(chunk_id: int):
    # reconfigure dspy to use a new instance of the OpenAI API client,
    # so that we can track usage per task
    if settings.GRAPH_EXTRACT_LLM_PROVIDER == "openai":
        oai_client = dspy.OpenAI(
            model="gpt-4o", api_key=settings.OPENAI_API_KEY, max_tokens=4096
        )
    else:
        oai_client = dspy.Google(
            model="gemini-1.5-flash",
            api_key=settings.GOOGLE_API_KEY,
            max_output_tokens=8192,
        )
    dspy.settings.configure(lm=oai_client)
    with Session(engine) as session:
        chunk = session.query(Chunk).get(chunk_id)

        if chunk.import_status != ImportStatus.EXTRACTED_GRAPH:
            logger.info(
                f"Chunk {chunk_id}[status: {chunk.import_status}] not ready for saving, skipping"
            )
            return

        chunk.import_status = ImportStatus.SAVING_TO_DB
        session.commit()

        entities_df = pickle.loads(chunk.entities_df)
        relationships_df = pickle.loads(chunk.relationships_df)

        extractor = SimpleGraphExtractor()
        try:
            extractor.save_to_db(
                session,
                chunk.text,
                chunk.doc_id,
                entities_df,
                relationships_df,
            )
        except Exception:
            chunk.import_status = ImportStatus.FAILED
            error_msg = f"Failed to save graph to db: {traceback.format_exc()}"
            logger.error(error_msg)
            chunk.import_result = error_msg
            session.commit()
            return

        chunk.import_status = ImportStatus.DONE
        if settings.GRAPH_EXTRACT_LLM_PROVIDER == "openai":
            chunk.llm_usage["save_db"] = dspy.settings.lm.history
            flag_modified(chunk, "llm_usage")
        session.commit()