from app.pipeline.pdf_reader import PDFReader
from app.pipeline.layout_detector import LayoutDetector
from app.pipeline.article_segmenter import ArticleSegmenter
from app.pipeline.ocr_engine import OCREngine
from app.pipeline.headline import detect_headline
from app.pipeline.cropper import crop_article, crop_figures
from app.pipeline.categorizer import Categorizer
from app.pipeline.enricher import Enricher
from app.config import settings
import os

class Pipeline:
    def __init__(self):
        self.layout = LayoutDetector()
        self.segmenter = ArticleSegmenter()
        self.ocr = OCREngine(settings.OCR_LANGS)
        self.cat = Categorizer()
        self.enricher = Enricher()

    def process(self, pdf_path, newspaper, edition, date, job_dir):
        reader = PDFReader(pdf_path)
        results = []
        for page in reader.iter_pages():
            h, w = page["image"].shape[:2]
            blocks = self.layout.detect(page["image"])
            articles = self.segmenter.segment(blocks, w, h)
            for art in articles:
                text, conf = self.ocr.extract_text(page, art["bbox"], page["native_words"])
                if len(text.strip()) < 80:   # skip noise/ads
                    continue
                headline = detect_headline(art, page, self.ocr)
                img_path = crop_article(page["image"], art["bbox"], job_dir)
                fig_paths = crop_figures(page["image"], art, job_dir)
                
                # Semantic Embedding for RAG
                raw_text_for_embed = headline + " " + text[:1000]
                embedding = self.enricher.embed(raw_text_for_embed)
                
                # Fetch few-shot examples from DB
                few_shot_examples = []
                try:
                    from app.db.session import SessionLocal
                    from sqlalchemy import text as sqla_text
                    db = SessionLocal()
                    # Query for up to 3 verified examples using cosine distance
                    rows = db.execute(sqla_text("""
                        SELECT headline, text, category, subcategory, summary, sentiment, keywords, persons, locations, organizations 
                        FROM articles 
                        WHERE human_verified = TRUE 
                        ORDER BY embedding <=> :v LIMIT 3
                    """), {"v": str(embedding)}).fetchall()
                    for r in rows:
                        ex = {
                            "headline": r.headline,
                            "text": r.text,
                            "classification": {
                                "category": r.category,
                                "subcategory": r.subcategory,
                                "summary": r.summary,
                                "sentiment": r.sentiment,
                                "keywords": r.keywords,
                                "persons": r.persons,
                                "locations": r.locations,
                                "organizations": r.organizations
                            }
                        }
                        few_shot_examples.append(ex)
                except Exception as e:
                    print(f"RAG fetch failed: {e}")
                finally:
                    if 'db' in locals(): db.close()

                meta = self.cat.classify(headline, text, few_shot_examples=few_shot_examples)
                
                results.append({
                    "date": date, "newspaper": newspaper, "edition": edition,
                    "page": page["page_number"], "headline": headline,
                    "category": meta.get("category", "Others"), "subcategory": meta.get("subcategory"),
                    "text": text, "summary": meta.get("summary"),
                    "keywords": meta.get("keywords", []), "persons": meta.get("persons", []),
                    "locations": meta.get("locations", []), "organizations": meta.get("organizations", []),
                    "sentiment": meta.get("sentiment"), "language": "en",
                    "image_path": img_path, "figure_paths": fig_paths,
                    "bounding_box": art["bbox"],
                    "confidence": round(art["confidence"] * meta.get("confidence", 1), 3),
                    "embedding": embedding,
                })
        return results