#!/usr/bin/env python3
"""
Phase 1.3: High-Speed Parallel LLM Semantic Relabeling (Ollama + LMS Dual Engine)
"""

import json
import os
import sys
import time
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed

INPUT_FILE = "/opt/vox/sandbox/datasets/semantic_raw.jsonl"
OUTPUT_FILE = "/opt/vox/sandbox/datasets/semantic_relabeled.jsonl"

OLLAMA_URL = "http://localhost:11434/api/generate"
LMS_URL = "http://localhost:1234/v1/chat/completions"

SYSTEM_PROMPT = """Classify the user query into EXACTLY ONE category:
- "User": Personal identity, persona, preferences, user constraints ("My name is Emily", "I am a software engineer", "I prefer async Rust").
- "Domain": Codebases, technical Q&A, active tasks, programming, architecture, bugs ("Fix Tokio deadlock", "How does stage 3 pipeline work?").
- "Temporal": Session recency, context recaps, session continuity ("What did we work on yesterday?", "Summarize last turn").

Output JSON ONLY: {"scope": "User" | "Domain" | "Temporal"}"""

def classify_query_llm(item, worker_id):
    text = item["text"]
    
    prompt = f"{SYSTEM_PROMPT}\n\nQuery: \"{text}\"\nJSON Output:"
    
    # Alternate between Ollama and LMS depending on worker_id to balance load
    use_lms = (worker_id % 2 == 1)
    
    scope = None
    if use_lms:
        try:
            payload = {
                "model": "llama-3.1-8b-instruct",
                "messages": [{"role": "user", "content": prompt}],
                "temperature": 0.0,
                "response_format": {"type": "json_object"}
            }
            res = requests.post(LMS_URL, json=payload, timeout=10)
            if res.status_code == 200:
                content = res.json()["choices"][0]["message"]["content"]
                parsed = json.loads(content)
                scope = parsed.get("scope")
        except Exception:
            pass

    if not scope: # Try Ollama fallback
        try:
            payload = {
                "model": "llama3.1:8b",
                "prompt": prompt,
                "stream": False,
                "options": {"temperature": 0.0}
            }
            res = requests.post(OLLAMA_URL, json=payload, timeout=10)
            if res.status_code == 200:
                resp_text = res.json().get("response", "").strip()
                s = resp_text.find("{")
                e = resp_text.rfind("}")
                if s != -1 and e != -1:
                    parsed = json.loads(resp_text[s:e+1])
                    scope = parsed.get("scope")
        except Exception:
            pass

    # Heuristic fast-path fallback if LLM times out
    if not scope or scope not in ["User", "Domain", "Temporal"]:
        text_lower = text.lower()
        if any(w in text_lower for w in ["yesterday", "last session", "previous", "earlier", "recap", "summary", "kal", "pichle"]):
            scope = "Temporal"
        elif any(w in text_lower for w in ["my name", "i am", "i live", "i prefer", "my role", "mera name", "main", "meri"]):
            scope = "User"
        else:
            scope = "Domain" # Primary default

    item_copy = dict(item)
    item_copy["scope"] = scope
    return item_copy

def main():
    print("=== Phase 1.3: Parallel LLM Semantic Relabeling ===", flush=True)
    
    if not os.path.exists(INPUT_FILE):
        print(f"Error: {INPUT_FILE} missing!", flush=True)
        sys.exit(1)
        
    items = []
    with open(INPUT_FILE, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                items.append(json.loads(line.strip()))
                
    total = len(items)
    print(f"Loaded {total} raw semantic items to relabel.", flush=True)
    
    results = []
    start_time = time.time()
    
    num_workers = 16
    print(f"Starting {num_workers} parallel worker threads across Ollama & LMS...", flush=True)
    
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        futures = {executor.submit(classify_query_llm, item, idx): idx for idx, item in enumerate(items)}
        
        completed_count = 0
        for future in as_completed(futures):
            res = future.result()
            results.append(res)
            completed_count += 1
            if completed_count % 500 == 0 or completed_count == total:
                elapsed = time.time() - start_time
                rate = completed_count / elapsed
                print(f"  Progress: {completed_count}/{total} ({completed_count/total*100:.1f}%) | {rate:.1f} items/sec", flush=True)
                
    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        for item in results:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")
            
    # Distribution tally
    counts = {"User": 0, "Domain": 0, "Temporal": 0}
    for r in results:
        sc = r.get("scope", "Domain")
        counts[sc] = counts.get(sc, 0) + 1
        
    print(f"\n✅ Relabeling complete! Saved {len(results)} items to {OUTPUT_FILE}", flush=True)
    print("Relabeled Scope Distribution:")
    for sc, cnt in counts.items():
        print(f"  - {sc}: {cnt} ({cnt/len(results)*100:.1f}%)", flush=True)

if __name__ == "__main__":
    main()