import gradio as gr
import json
import time
import numpy as np
from pathlib import Path
from typing import List, Dict, Any, Optional
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import os
import re


# ====================================
# CONFIGURATION
# ====================================
MODEL_REPO = "xubayer/LlamaPIE-GGUF"
SMALL_MODEL_FILENAME = "llamapie-1b-q8_0.gguf"
LARGE_MODEL_FILENAME = "llamapie-8b-q8_0.gguf"


# ====================================
# MODEL LOADING
# ====================================
print("⏳ Downloading LlamaPIE models...")
small_model_path = hf_hub_download(repo_id=MODEL_REPO, filename=SMALL_MODEL_FILENAME)
large_model_path = hf_hub_download(repo_id=MODEL_REPO, filename=LARGE_MODEL_FILENAME)

print("🚀 Loading Classifier (1B)...")
classifier = Llama(
    small_model_path, 
    n_ctx=2048, 
    n_threads=2, 
    verbose=False,
    logits_all=True
)

print("🚀 Loading Generator (8B)...")
generator = Llama(large_model_path, n_ctx=4096, n_threads=4, verbose=False)


# ====================================
# SIMPLIFIED PROMPTING (Raw Text)
# ====================================
def classify_with_simple_continuation(model, dialogue_context: str) -> Dict:
    """
    Simple approach: Just continue the dialogue and see what the model generates.
    If it generates meaningful content, it's proactive. If not, it's reactive.
    """
    
    # Simple prompt: just the dialogue ending with |SILENCE >
    prompt = dialogue_context + " |SILENCE >"
    
    output = model(
        prompt,
        max_tokens=15,
        temperature=0.3,
        top_p=0.9,
        logprobs=5,
        stop=["User:", "Speaker", "\n\n", "|SILENCE"]
    )
    
    generated_text = output['choices'][0]['text'].strip()
    whisper_prob = 0.5
    
    # Analyze what was generated
    try:
        logprobs_data = output['choices'][0].get('logprobs', {})
        top_logprobs = logprobs_data.get('top_logprobs', [])
        
        if top_logprobs and len(top_logprobs) > 0:
            first_token_probs = top_logprobs[0]
            proactive_score = 0.0
            
            for token, logprob in first_token_probs.items():
                prob = np.exp(logprob)
                token_str = str(token).lower().strip()
                
                # Proactive: actual content words
                if len(token_str) > 2 and any(c.isalpha() for c in token_str):
                    proactive_score += prob * 2.0
                # Reactive: punctuation, whitespace
                elif token_str in ['.', ',', '!', '?', '', ' ', '\n', '<', '|', '▁']:
                    proactive_score -= prob * 1.5
                else:
                    proactive_score += prob * 0.3
            
            whisper_prob = 1.0 / (1.0 + np.exp(-proactive_score * 2.5))
            
    except Exception as e:
        print(f"   Logprob error: {e}")
    
    # Text-based heuristics
    if not generated_text or len(generated_text.strip()) < 2:
        whisper_prob = 0.1
    elif generated_text.strip() in ['.', ',', '!', '?', '|']:
        whisper_prob = 0.2
    elif len(generated_text.split()) >= 3:
        whisper_prob = max(whisper_prob, 0.9)
    elif len(generated_text.split()) >= 2:
        whisper_prob = max(whisper_prob, 0.8)
    elif any(c.isalpha() for c in generated_text) and len(generated_text.strip()) > 3:
        whisper_prob = max(whisper_prob, 0.7)
    
    silent_prob = 1.0 - whisper_prob
    decision = "WHISPER" if whisper_prob > 0.5 else "SILENT"
    confidence = max(whisper_prob, silent_prob)
    
    return {
        "decision": decision,
        "confidence": float(confidence),
        "whisper_probability": float(whisper_prob),
        "silent_probability": float(silent_prob),
        "generated_text": generated_text
    }


# ====================================
# CORE INFERENCE LOGIC
# ====================================
class LlamaPIEInference:
    def __init__(self, classifier, generator):
        self.classifier = classifier
        self.generator = generator
        
    def parse_dialogue(self, dialogue: str) -> List[Dict]:
        """Parse dialogue into segments ending with |SILENCE > markers."""
        silence_positions = [m.end() for m in re.finditer(r'\|\s*SILENCE\s*\>', dialogue)]
        segments = []
        
        start = 0
        for pos in silence_positions:
            segment = dialogue[start:pos].strip()
            if segment:
                segments.append({
                    "text": segment,
                    "end_pos": pos,
                    "context": dialogue[max(0, pos-500):pos]
                })
            start = pos
        return segments

    def classify_silence_position(self, context: str, confidence_threshold: float = 0.6) -> Dict:
        """
        Simple classification: continue the dialogue and check output.
        """
        print(f"\n🔍 Classifying...")
        print(f"   Last 80 chars: ...{context[-80:]}")
        
        try:
            classification = classify_with_simple_continuation(self.classifier, context)
            
            classification["threshold_met"] = (
                classification["decision"] == "WHISPER" and 
                classification["whisper_probability"] >= confidence_threshold
            )
            
            print(f"   Decision: {classification['decision']} (P={classification['whisper_probability']:.2f})")
            print(f"   Generated: '{classification['generated_text'][:60]}'")
            
            return classification
            
        except Exception as e:
            print(f"   ❌ Error: {e}")
            return {
                "decision": "SILENT",
                "confidence": 0.5,
                "whisper_probability": 0.5,
                "silent_probability": 0.5,
                "generated_text": "",
                "threshold_met": False
            }
    
    def generate_whisper(self, context: str, memory: str) -> str:
        """
        Generate whisper using memory context.
        """
        # Create a simple prompt with memory facts
        memory_text = ""
        if memory:
            try:
                memory_obj = json.loads(memory)
                profile = memory_obj.get("profile", "")
                events = memory_obj.get("events", {})
                
                memory_parts = []
                if profile:
                    memory_parts.append(f"Profile: {profile}")
                for event_key, event_val in events.items():
                    if isinstance(event_val, str):
                        memory_parts.append(f"Memory: {event_val}")
                
                if memory_parts:
                    memory_text = "\n".join(memory_parts) + "\n\n"
            except:
                memory_text = f"{memory}\n\n"
        
        # Simple prompt format
        prompt = f"""{memory_text}Conversation:
{context} |SILENCE >

Provide a helpful 1-3 word whisper:"""
        
        print(f"   🎤 Generating whisper...")
        
        output = self.generator(
            prompt,
            max_tokens=10,
            temperature=0.4,
            top_p=0.8,
            stop=["\n", "User:", "Speaker:", "Conversation:", "Profile:", "Memory:"],
        )
        
        raw_text = output['choices'][0]['text'].strip()
        
        # Extract and clean whisper
        whisper = raw_text.split('\n')[0].strip()
        whisper = re.sub(r'["\',;(){}[\]:]', '', whisper)
        
        # Remove common prefixes
        whisper = re.sub(r'^(whisper|hint|suggestion|reminder):\s*', '', whisper, flags=re.IGNORECASE)
        
        # Limit to 1-3 words
        words = [w for w in whisper.split() if len(w) > 0]
        if len(words) == 0:
            whisper = "Continue"
        elif len(words) > 3:
            whisper = " ".join(words[:3])
        else:
            whisper = " ".join(words)
        
        # Filter out useless whispers
        if whisper.lower() in ["the", "and", "or", "but", "a", "an", "is", "it", "to"]:
            whisper = "Continue"
        
        print(f"   Result: '{whisper}'")
        
        return whisper

    def process_conversation(self, dialogue: str, memory: str = "", 
                           confidence_threshold: float = 0.6) -> Dict[str, Any]:
        """Main pipeline: Process full conversation."""
        start_time = time.time()
        
        segments = self.parse_dialogue(dialogue)
        whispers = []
        
        print(f"\n📝 Processing {len(segments)} silence positions with threshold={confidence_threshold}")
        
        for i, segment in enumerate(segments):
            print(f"\n--- Position {i+1}/{len(segments)} ---")
            classification = self.classify_silence_position(
                segment["context"], 
                confidence_threshold
            )
            
            if classification["threshold_met"]:
                whisper_result = self.generate_whisper(segment["context"], memory)
                
                whispers.append({
                    "position": i + 1,
                    "text": whisper_result,
                    "confidence": classification["confidence"],
                    "whisper_probability": classification["whisper_probability"],
                    "silent_probability": classification["silent_probability"],
                    "context": segment["context"][-100:],
                    "classifier_output": classification.get("generated_text", "")[:50]
                })
        
        # Annotate dialogue
        annotated = dialogue
        for whisper in whispers:
            annotated = annotated.replace("|SILENCE >", f" [💬 {whisper['text']}] |SILENCE >", 1)
        
        print(f"\n✅ Complete! {len(whispers)}/{len(segments)} positions triggered")
        
        return {
            "input": {
                "dialogue": dialogue,
                "memory": memory,
                "confidence_threshold": confidence_threshold
            },
            "whispers": whispers,
            "annotated_dialogue": annotated,
            "stats": {
                "total_silence_positions": len(segments),
                "whisper_decisions": len(whispers),
                "processing_time_ms": int((time.time() - start_time) * 1000),
                "whisper_frequency": f"{len(whispers)/max(1, len(segments))*100:.1f}%"
            }
        }


# Instantiate pipeline
pipeline = LlamaPIEInference(classifier, generator)


# ====================================
# GRADIO UI
# ====================================
css = """
.gradio-container { max-width: 1100px !important; }
.output-box { background: #f8f9fa; border-radius: 8px; padding: 12px; }
"""


def run_inference(dialogue: str, memory: str, confidence_threshold: float):
    """Gradio inference wrapper."""
    try:
        result = pipeline.process_conversation(dialogue, memory, confidence_threshold)
        
        whisper_lines = [
            f"• Pos {w['position']}: '{w['text']}' (P={w['whisper_probability']:.2f})"
            for w in result["whispers"]
        ]
        whispers_text = "\n".join(whisper_lines) if whisper_lines else "No whispers triggered"
        
        annotated = result["annotated_dialogue"]
        if len(annotated) > 2000:
            annotated = annotated[:2000] + "\n... (truncated)"
        
        output = (
            f"### 📊 **Results** ({result['stats']['processing_time_ms']}ms)\n"
            f"**Whispers:** {result['stats']['whisper_decisions']}/{result['stats']['total_silence_positions']} positions "
            f"({result['stats']['whisper_frequency']})\n\n"
            f"**Generated Whispers:**\n"
            f"```\n{whispers_text}\n```\n\n"
            f"**Annotated Dialogue:**\n"
            f"```\n{annotated}\n```\n"
        )
        
        return output, json.dumps(result, indent=2)
        
    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        return f"❌ Error: {str(e)}\n\n```\n{error_details}\n```", ""


# ====================================
# GRADIO INTERFACE
# ====================================
with gr.Blocks(css=css, title="LlamaPIE Whisper Inference") as demo:
    gr.Markdown("# 🤖 LlamaPIE Whisper Inference")
    gr.Markdown("*Simple proactive assistant with memory-aware whispers*")
    
    with gr.Row():
        with gr.Column(scale=1):
            memory_input = gr.Textbox(
                label="📚 Memory Context (JSON)",
                placeholder='{"profile": "...", "events": {...}}',
                lines=8,
                value='{"profile": "User is interested in technology", "events": {}}'
            )
            confidence_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.65,
                step=0.05,
                label="🎯 Confidence Threshold (higher = fewer whispers)"
            )
            
        with gr.Column(scale=2):
            dialogue_input = gr.Textbox(
                label="💬 Dialogue (with |SILENCE > markers)",
                placeholder="User: Hello |SILENCE > Agent: Hi there |SILENCE >",
                lines=12
            )
            
    submit_btn = gr.Button("🚀 Generate Whispers", variant="primary", size="lg")
    
    with gr.Row():
        with gr.Column():
            output_display = gr.Markdown(label="Results")
            
    with gr.Row():
        with gr.Column():
            json_output = gr.JSON(label="Full JSON Output")
    
    submit_btn.click(
        fn=run_inference,
        inputs=[dialogue_input, memory_input, confidence_slider],
        outputs=[output_display, json_output]
    )
    
    gr.Examples(
        examples=[
            [
                "User: I've been working on a project. |SILENCE > Agent: That's great! |SILENCE > User: I added the new database |SILENCE > feature. |SILENCE >",
                '{"profile": "Software developer working on web app", "events": {"event0": "Struggled with database optimization last week"}}',
                0.65
            ],
            [
                "User: I need to call mom. |SILENCE > Agent: Sure! |SILENCE > User: And that important |SILENCE > meeting tomorrow. |SILENCE >",
                '{"profile": "Busy schedule, forgets appointments", "events": {"event0": "Mom\'s birthday coming up"}}',
                0.7
            ],
        ],
        inputs=[dialogue_input, memory_input, confidence_slider],
        label="Example Conversations"
    )


if __name__ == "__main__":
    demo.launch()