import json
import os

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langsmith import traceable
from openai import OpenAI

from script.extraction import extract_text_native
from script.ocr import extract_text_via_ocr
from script.errors import IngestionError

from script.config import CHUNK_SIZE, CHUNK_OVERLAP
from script.logger import logger

from dotenv import load_dotenv

load_dotenv()

TITLE_SUMMARY_MODEL = "gpt-5.6-luna"

# A line-based heuristic ("first line over 20 chars that doesn't start with
# a digit") is unreliable on real academic PDFs: license/attribution
# boilerplate, running headers, and author-affiliation blocks routinely
# satisfy that same shape before the actual title does. Extracting title
# and summary together in one structured call costs no more than the
# summary-only call this replaces, and is far more robust.
TITLE_SUMMARY_SCHEMA = {
    "type": "json_schema",
    "json_schema": {
        "name": "paper_title_and_summary",
        "schema": {
            "type": "object",
            "properties": {
                "title": {
                    "type": "string",
                    "description": (
                        "The paper's actual title as printed on the page — "
                        "not a running header, license/attribution notice, "
                        "author list, or venue name."
                    ),
                },
                "summary": {
                    "type": "string",
                    "description": "A 5-6 sentence summary covering the goal, method, and contribution.",
                },
            },
            "required": ["title", "summary"],
            "additionalProperties": False,
        },
        "strict": True,
    },
}

_openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))


def extract_title_and_summary(first_page_text: str) -> tuple[str, str]:
    response = _openai_client.chat.completions.create(
        model=TITLE_SUMMARY_MODEL,
        messages=[{
            "role": "user",
            "content": (
                "Extract the title and a 5-6 sentence summary (goal, method, "
                "contribution) of this research paper.\n\nText:\n"
                f"{first_page_text[:3000]}"
            ),
        }],
        response_format=TITLE_SUMMARY_SCHEMA,
    )
    result = json.loads(response.choices[0].message.content)
    return result["title"], result["summary"]

@traceable(name="pdf_ingestion")
def ingest_pdf(pdf_path: str) -> dict:
    """
    Ingest a PDF and return structured artifacts.

    Returns:
    {
        "vectorstore": FAISS,
        "title": str,
        "summary": str
    }
    """
    try:
        logger.info(f"Starting ingestion for PDF: {pdf_path}")

        # -------- 1. Load PDF --------
        try:
            documents = extract_text_native(pdf_path)
            extraction_method = "native"
        except IngestionError as e:
            logger.warning(
                f"Native extraction failed ({e.code}). Falling back to OCR."
            )
            documents = extract_text_via_ocr(pdf_path)
            extraction_method = "ocr"

        logger.info(
            f"Text extraction successful using method: {extraction_method}"
        )

        logger.info(f"Loaded {len(documents)} pages")

        # -------- 2. Extract title (deterministic) --------
        def get_first_meaningful_page(documents, min_chars=300):
            for doc in documents:
                text = doc.page_content.strip()
                if len(text) >= min_chars:
                    return text
            return ""
    
        first_page_text = get_first_meaningful_page(documents)

        # -------- 3. Extract title + generate summary (one call) --------
        if not first_page_text:
            logger.warning("No meaningful text found for title/summary extraction")
            title = "Unable to extract title"
            summary = "Summary unavailable due to insufficient extractable text."
        else:
            title, summary = extract_title_and_summary(first_page_text)
            logger.info("Extracted title and generated paper summary")

        logger.info(f"Extracted title: {title}")

        # -------- 4. Chunk full document --------
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=CHUNK_SIZE,
            chunk_overlap=CHUNK_OVERLAP,
        )

        chunks = splitter.split_documents(documents)

        if not chunks:
            raise IngestionError(
                code="ZERO_CHUNKS",
                message="Chunking failed even after extraction"
            )
        logger.info(
            f"Pages: {len(documents)} | "
            f"Total chars: {sum(len(d.page_content) for d in documents)}"
        )


        logger.info(f"Created {len(chunks)} chunks")

        # -------- 5. Build vectorstore --------
        embeddings = OpenAIEmbeddings()
        vectorstore = FAISS.from_documents(chunks, embeddings)

        logger.info("FAISS vectorstore created successfully")

        return {
            "vectorstore": vectorstore,
            "title": title,
            "summary": summary,
        }

    except Exception:
        logger.exception("PDF ingestion failed")
        raise