import os
import argparse
import json
import re
import torch
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from genre_policy import GenreDecisionPolicy

def clean_title(text):
    # Remove common junk from title string
    text = re.sub(r'\[.*?\]|\(.*?\)', '', text) # Remove brackets and parens completely, often contains series/format info, but maybe too aggressive? Let's just remove format tags.
    # Let's be conservative: remove known tags
    text = re.sub(r'\[(?i:epub|mobi|pdf|retail|calibre|libgen)\]', '', text)
    # Convert underscores to spaces
    text = text.replace('_', ' ')
    # Remove extra spaces
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def clean_author(text):
    text = text.replace('_', ' ')
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def infer_title_author_from_path(filepath):
    # Heuristics:
    # 1. /Author/Title.epub
    # 2. /Author/Series/Title.epub
    # 3. Author - Title.epub
    # 4. Title - Author.epub (Harder to guess, default to Author - Title if "-" present)
    
    filename = os.path.basename(filepath)
    name_no_ext, _ = os.path.splitext(filename)
    
    dir_path = os.path.dirname(filepath)
    parent_dir = os.path.basename(dir_path) if dir_path else ""
    grandparent_dir = os.path.basename(os.path.dirname(dir_path)) if os.path.dirname(dir_path) else ""
    
    author = "Unknown"
    title = name_no_ext

    # Check for " - " in filename
    if " - " in name_no_ext:
        parts = name_no_ext.split(" - ", 1)
        author = parts[0]
        title = parts[1]
    elif parent_dir and parent_dir.lower() not in ["books", "ebooks", "fiction", "nonfiction", "calibre library"]:
        # If parent dir looks like a name (e.g. not a generic folder)
        # We might use it as author, or grandparent if parent is a series.
        # Let's keep it simple: assume parent is author if no "-" in filename
        author = parent_dir
        
    return clean_title(title), clean_author(author)

def main():
    parser = argparse.ArgumentParser(description="Segment disk corpus by genre")
    parser.add_argument("--root", required=True, type=str, help="Root folder to scan")
    parser.add_argument("--output", required=True, type=str, help="Output JSONL manifest")
    parser.add_argument("--extensions", default=".epub,.mobi,.azw3,.pdf,.txt", type=str)
    parser.add_argument("--batch-size", default=256, type=int)
    parser.add_argument("--model-dir", default=".", type=str)
    parser.add_argument("--device", default="cuda", type=str)
    parser.add_argument("--dry-run", action="store_true", help="Scan and extract metadata but don't run model")
    args = parser.parse_args()

    exts = [e.strip().lower() for e in args.extensions.split(",")]
    
    print(f"Scanning {args.root} for {exts}...")
    files_to_process = []
    for root_dir, _, files in os.walk(args.root):
        for f in files:
            if any(f.lower().endswith(ext) for ext in exts):
                files_to_process.append(os.path.join(root_dir, f))
                
    print(f"Found {len(files_to_process)} files.")
    
    if args.dry_run:
        print("Dry run requested. Exiting after scan.")
        # Just show a few examples
        for f in files_to_process[:5]:
            t, a = infer_title_author_from_path(f)
            print(f"Path: {f}\nInferred Title: {t}\nInferred Author: {a}\n")
        return

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(args.model_dir)
    model.to(device)
    model.eval()

    policy_path = os.path.join(args.model_dir, "decision_policy.json")
    policy = GenreDecisionPolicy(policy_path)
    id2label = model.config.id2label

    # Prepare batches
    batches = [files_to_process[i:i + args.batch_size] for i in range(0, len(files_to_process), args.batch_size)]

    with open(args.output, "w") as f_out:
        for batch in tqdm(batches, desc="Processing batches"):
            texts = []
            metadata = []
            for filepath in batch:
                title, author = infer_title_author_from_path(filepath)
                txt = f"Title: {title}\nAuthor: {author}"
                texts.append(txt)
                metadata.append({
                    "path": filepath,
                    "filename": os.path.basename(filepath),
                    "inferred_title": title,
                    "inferred_author": author,
                    "input_text": txt
                })

            inputs = tokenizer(texts, return_tensors="pt", truncation=True, padding=True, max_length=128)
            inputs = {k: v.to(device) for k, v in inputs.items()}
            
            with torch.no_grad():
                outputs = model(**inputs)
                probs = torch.nn.functional.softmax(outputs.logits, dim=-1).cpu().numpy()

            for j, prob in enumerate(probs):
                top_indices = prob.argsort()[::-1]
                top_k = [[id2label[idx], float(prob[idx])] for idx in top_indices]
                
                decision_out = policy.decide(top_k)
                
                # Determine suggested folder
                suggested = "unknown"
                if decision_out["decision"] == "exact_auto_label":
                    suggested = decision_out["final_label"].lower().replace(" / ", "_").replace(" ", "_").replace("-", "_")
                elif decision_out["decision"] == "ambiguous_bucket":
                    suggested = "ambiguous"
                    
                res = metadata[j]
                res.update({
                    "top1_label": top_k[0][0],
                    "top1_confidence": round(top_k[0][1], 4),
                    "top2_label": top_k[1][0],
                    "top2_confidence": round(top_k[1][1], 4),
                    "margin": round(top_k[0][1] - top_k[1][1], 4),
                    "decision": decision_out["decision"],
                    "final_label": decision_out["final_label"],
                    "ambiguity_bucket": decision_out["ambiguity_bucket"],
                    "suggested_folder": suggested,
                    "model_version": "book_genre_v5_title_author"
                })
                
                f_out.write(json.dumps(res) + "\n")

    print(f"Segmentation manifest saved to {args.output}")

if __name__ == "__main__":
    main()