import os
import shutil
from tqdm import tqdm
from dotenv import load_dotenv

from db import init_db
from ocr import ocr_file
from embed_and_store import process_pages

load_dotenv()

DOCUMENTS_DIR = "documents"
PROCESSED_DIR = "processed"

SUPPORTED_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".tif", ".gif", ".bmp", ".webp"}

# ── Tag your documents here ──────────────────────────────────────────────────
# Maps a filename substring → (nationality, corpus)
# First matching rule wins. Falls back to ("unknown", "unknown").
DOCUMENT_TAGS = {
    # German records
    "german":    ("german",   "german_records"),
    "deutsch":   ("german",   "german_records"),
    "Wehrmacht": ("german",   "german_records"),
    # Canadian war diaries
    "Calgary":   ("canadian", "war_diary_narrative"),
    "Blackwatch": ("canadian","war_diary_narrative"),
    "5CIB":      ("canadian", "war_diary_narrative"),
}

def tag_document(filename: str) -> tuple[str, str]:
    for key, (nationality, corpus) in DOCUMENT_TAGS.items():
        if key.lower() in filename.lower():
            return nationality, corpus
    return "unknown", "unknown"
# ─────────────────────────────────────────────────────────────────────────────


def main():
    os.makedirs(DOCUMENTS_DIR, exist_ok=True)
    os.makedirs(PROCESSED_DIR, exist_ok=True)
    init_db()

    files = [
        f for f in os.listdir(DOCUMENTS_DIR)
        if os.path.splitext(f)[1].lower() in SUPPORTED_EXTENSIONS
    ]

    if not files:
        print(f"No supported files found in '{DOCUMENTS_DIR}'. Drop PDFs or images there and re-run.")
        return

    print(f"Found {len(files)} file(s) to process.\n")

    for filename in tqdm(files, desc="Processing documents"):
        file_path = os.path.join(DOCUMENTS_DIR, filename)
        nationality, corpus = tag_document(filename)
        print(f"\n--- {filename} --- [{nationality} / {corpus}]")

        try:
            pages = ocr_file(file_path)
            print(f"  OCR complete: {len(pages)} page(s)")

            process_pages(filename, pages, nationality=nationality, corpus=corpus)

            shutil.move(file_path, os.path.join(PROCESSED_DIR, filename))
            print(f"  Moved to {PROCESSED_DIR}/")

        except Exception as e:
            print(f"  ERROR processing {filename}: {e}")

    print("\nDone. All files processed.")


def search(query: str, top_k: int = 5, nationality: str = None, corpus: str = None):
    from embed_and_store import embed_texts
    from db import semantic_search

    filters = []
    if nationality: filters.append(f"nationality={nationality}")
    if corpus:      filters.append(f"corpus={corpus}")
    filter_str = f" [{', '.join(filters)}]" if filters else ""
    print(f"\nSearching for: '{query}'{filter_str}\n")

    embedding = embed_texts([query])[0]
    results = semantic_search(embedding, top_k=top_k, nationality=nationality, corpus=corpus)

    for rank, (filename, page_num, chunk_index, text, nat, corp, similarity) in enumerate(results, 1):
        print(f"[{rank}] {filename} — page {page_num}, chunk {chunk_index} [{nat} / {corp}] (similarity: {similarity:.3f})")
        print(f"    {text[:200]}...")
        print()
def ask(question: str, top_k: int = 5, nationality: str = None, corpus: str = None):
    from embed_and_store import embed_texts
    from db import semantic_search
    import os
    from openai import OpenAI

    filters = []
    if nationality: filters.append(f"nationality={nationality}")
    if corpus:      filters.append(f"corpus={corpus}")
    filter_str = f" [{', '.join(filters)}]" if filters else ""
    print(f"\nAsking: '{question}'{filter_str}\n")

    # Step 1 - find relevant chunks
    embedding = embed_texts([question])[0]
    results = semantic_search(embedding, top_k=top_k, nationality=nationality, corpus=corpus)

    if not results:
        print("No relevant chunks found.")
        return

    # Step 2 - build context from chunks
    context = ""
    for filename, page_num, chunk_index, text, nat, corp, similarity in results:
        context += f"[Source: {filename}, page {page_num}]\n{text}\n\n"

    # Step 3 - send to LLM
    client = OpenAI(
        api_key=os.getenv("DEEPINFRA_API_KEY"),
        base_url="https://api.deepinfra.com/v1/openai",
    )

    prompt = f"""You are a WWII military historian. Answer the question below using only the provided source documents.
Cite the source file and page number for each claim you make.

QUESTION: {question}

SOURCE DOCUMENTS:
{context}

Answer:"""

    response = client.chat.completions.create(
        model="meta-llama/Meta-Llama-3.1-8B-Instruct",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=1000,
    )

    print("\nQUESTION:", question)
    print("\nANSWER:")
    print(response.choices[0].message.content)
    print("\nSOURCES USED:")
    for filename, page_num, chunk_index, text, nat, corp, similarity in results:
        print(f"  - {filename} page {page_num} [{nat} / {corp}] (similarity: {similarity:.3f})")

if __name__ == "__main__":
    # main()
    # search("is this document in german?", nationality="german")
    ask("What was the Calgary Highlanders doing in October 1944?", nationality="canadian", corpus="war_diary_narrative")