#!/usr/bin/env python3
"""
FAR Chatbot - Simplified version for Hugging Face Spaces
"""

import os
import logging
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from openai import OpenAI
from typing import List, Tuple, Dict, Optional

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ConversationMemory:
    """Simple conversation memory"""
    def __init__(self, max_turns: int = 5):
        self.history = []
        self.max_turns = max_turns
        self.current_topics = []
    
    def add_turn(self, question: str, answer: str, topics: List[str] = None):
        self.history.append({"question": question, "answer": answer})
        if len(self.history) > self.max_turns:
            self.history.pop(0)
        if topics:
            self.current_topics = topics[:3]
    
    def get_context(self) -> str:
        if not self.history:
            return ""
        context = "Previous conversation:\n"
        for turn in self.history[-3:]:
            context += f"Q: {turn['question']}\nA: {turn['answer'][:200]}...\n\n"
        return context

class FARChatbot:
    """FAR Chatbot with RAG capabilities"""
    
    def __init__(self, faiss_index_path: str, texts_path: str, 
                 model_name: str = 'paraphrase-MiniLM-L6-v2',
                 openai_api_key: str = None, use_gpt5: bool = True):
        
        self.use_gpt5 = use_gpt5
        logger.info("Loading SentenceTransformer model...")
        self.model = SentenceTransformer(model_name)
        logger.info("Model loaded!")
        
        # Load FAISS index
        logger.info(f"Loading FAISS index from {faiss_index_path}")
        self.faiss_index = faiss.read_index(faiss_index_path)
        logger.info(f"FAISS index loaded with {self.faiss_index.ntotal} vectors")
        
        # Load texts
        logger.info(f"Loading texts from {texts_path}")
        with open(texts_path, 'r', encoding='utf-8') as f:
            self.texts = [line.strip() for line in f if line.strip()]
        logger.info(f"Loaded {len(self.texts)} text chunks")
        
        # OpenAI client
        api_key = openai_api_key or os.getenv('OPENAI_API_KEY')
        if not api_key:
            raise ValueError("OpenAI API key required")
        self.client = OpenAI(api_key=api_key)
        
        self.conversation = ConversationMemory()
    
    def search(self, query: str, top_k: int = 10) -> List[Tuple[str, str]]:
        """Search for relevant FAR sections"""
        query_embedding = self.model.encode([query])
        distances, indices = self.faiss_index.search(query_embedding.astype('float32'), top_k)
        
        results = []
        for idx in indices[0]:
            if 0 <= idx < len(self.texts):
                text = self.texts[idx]
                # Extract citation from text
                citation = "Unknown"
                if text.startswith("FAR "):
                    parts = text.split(":", 1)
                    if len(parts) > 1:
                        citation = parts[0].replace("FAR ", "").strip()
                results.append((citation, text))
        
        return results
    
    def chat(self, question: str, top_k: int = None) -> Dict:
        """Process a question and return response"""
        
        # Determine context size
        actual_top_k = 50 if self.use_gpt5 else (top_k or 10)
        
        # Search for relevant content
        search_results = self.search(question, top_k=actual_top_k)
        
        # Build context
        context = "\n\n".join([f"[{cit}]: {txt[:800]}" for cit, txt in search_results[:20]])
        
        # Get conversation history
        conv_context = self.conversation.get_context()
        
        # Build prompt
        system_prompt = """You are FAR Bot, an expert assistant for the Federal Acquisition Regulation (FAR).

INSTRUCTIONS:
1. Answer questions accurately based on the FAR content provided
2. ALWAYS cite specific FAR sections using [X.XXX] format
3. Be concise but thorough
4. If information isn't in the context, say so
5. Suggest follow-up questions when appropriate"""

        user_prompt = f"""Question: {question}

{conv_context}

Relevant FAR Sections:
{context}

Provide a clear, well-cited answer:"""

        # Call OpenAI
        try:
            response = self.client.chat.completions.create(
                model="gpt-4-turbo-preview",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.3,
                max_tokens=1500
            )
            answer = response.choices[0].message.content
        except Exception as e:
            logger.error(f"OpenAI error: {e}")
            answer = f"Error generating response: {e}"
        
        # Extract topics (simple extraction)
        topics = []
        topic_keywords = ["small business", "threshold", "competition", "contract", "bid", "proposal"]
        for kw in topic_keywords:
            if kw.lower() in question.lower():
                topics.append(kw.title())
        
        # Update conversation
        self.conversation.add_turn(question, answer, topics)
        
        # Generate suggestions
        suggestions = [
            f"What are the exceptions to this rule?",
            f"Can you provide more details about the thresholds?",
            f"What documentation is required?"
        ]
        
        return {
            'response': answer,
            'suggestions': suggestions[:3],
            'topics': topics,
            'sections': [cit for cit, _ in search_results[:5]],
            'search_results': search_results[:10],
            'context_size': len(search_results),
            'model_used': 'gpt-4-turbo'
        }