from fastapi import APIRouter, HTTPException, status
from pydantic import BaseModel

from app.api.deps import SessionDep
from app.knowledge_graph.graph_extractor import (
    get_chunk_by_uri,
    document_already_exists,
    SimpleGleaninger,
)
from app.knowledge_graph.misc import import_status_overview
from app.models import Chunk
from app.celery import extract_graph_from_chunk

router = APIRouter()


class BuildRequest(BaseModel):
    uri: str
    text: str


@router.post("/build")
def build(session: SessionDep, build_request: BuildRequest):
    if document_already_exists(session, build_request.uri):
        return {"status": "already exists, skip"}

    chunk = Chunk(
        doc_id=build_request.uri,
        text=build_request.text,
        meta={"source_uri": build_request.uri},
    )
    session.add(chunk)
    session.commit()
    session.refresh(chunk)

    extract_graph_from_chunk.delay(chunk.id)
    return {"status": "task submitted"}


@router.post("/glean")
def glean(session: SessionDep, build_request: BuildRequest):
    gleaner = SimpleGleaninger()
    entities_df, gleanings_triplets = gleaner.glean(
        session, build_request.text, build_request.uri
    )
    gleaner.save_to_db(session, build_request.uri, entities_df, gleanings_triplets)
    return {
        "status": "success",
        "entities": gleanings_triplets.entities,
        "relationships": gleanings_triplets.relationships,
    }


@router.get("/import/overview")
def import_overview(session: SessionDep):
    return import_status_overview(session)


class ImportCheckRequest(BaseModel):
    uri: str


@router.post("/import/check")
def import_check(session: SessionDep, import_check_request: ImportCheckRequest):
    chunk = get_chunk_by_uri(session, import_check_request.uri)
    if not chunk:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="document with this URI does not exist",
        )

    return {"import_status": chunk.import_status}