import re
import random
import json
import torch
from langchain_ollama import OllamaLLM
from langchain.memory import ConversationSummaryMemory
from Sastrawi.Stemmer.StemmerFactory import StemmerFactory

def check_gpu():
    """Check if GPU is available and return appropriate device."""
    if torch.cuda.is_available():
        print("GPU detected! Using:", torch.cuda.get_device_name(0))
        return "cuda"
    else:
        print("No GPU detected. Using CPU.")
        return "cpu"

# Get the appropriate device
device = check_gpu()

# Initialize the OllamaLLM with GPU support if available
llm_config = {
    "model": "llama3",
    "device": device,
    "gpu_layers": 50 if device == "cuda" else 0,  # Number of layers to offload to GPU
    "num_gpu": 1 if device == "cuda" else 0,  # Number of GPUs to use
    "num_thread": 8 if device == "cpu" else 4  # Adjust threads based on device
}

llm = OllamaLLM(**llm_config)

# Initialize conversation memory with session support
memory = ConversationSummaryMemory(llm=llm, memory_key="history")

# Files for storing conversation history and response templates
HISTORY_FILE = "conversation_history.json"
RESPONSE_FILES = {
    'general_questions': 'general_questions.json',
    'missing_info': 'missing_info_responses.json',
    'info_received': 'info_received_responses.json',
    'proposal_created': 'proposal_created_responses.json'
}

# Structure for storing proposal input status
user_inputs = {
    "judul": None,
    "deskripsi": None,
    "lokasi": None
}

# Initialize Sastrawi Stemmer
factory = StemmerFactory()
stemmer = factory.create_stemmer()

def load_responses(filename):
    """Load response templates from JSON files."""
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            data = json.load(file)
            # Return just the responses array for all except general_questions
            if filename != RESPONSE_FILES['general_questions']:
                return data['responses']
            return data
    except FileNotFoundError:
        print(f"Error: {filename} not found!")
        return {} if filename == RESPONSE_FILES['general_questions'] else []
    except json.JSONDecodeError:
        print(f"Error: Invalid JSON in {filename}")
        return {} if filename == RESPONSE_FILES['general_questions'] else []

# Load all response templates
general_questions = load_responses(RESPONSE_FILES['general_questions'])
missing_info_responses = load_responses(RESPONSE_FILES['missing_info'])
info_received_responses = load_responses(RESPONSE_FILES['info_received'])
proposal_created_responses = load_responses(RESPONSE_FILES['proposal_created'])

def load_history():
    """Load conversation history from file."""
    try:
        with open(HISTORY_FILE, "r", encoding="utf-8") as file:
            return json.load(file)
    except (FileNotFoundError, json.JSONDecodeError):
        return []

def save_history(history):
    """Save conversation history to file."""
    with open(HISTORY_FILE, "w", encoding="utf-8") as file:
        json.dump(history, file, ensure_ascii=False, indent=4)

def lemmatize_sentence(text):
    """Lemmatize each word in the input sentence using Sastrawi stemmer."""
    words = text.lower().split()
    lemmatized_words = [stemmer.stem(word) for word in words]
    return ' '.join(lemmatized_words)

def check_missing_info():
    """Check which information is still missing."""
    return [key for key, value in user_inputs.items() if value is None]

def clean_location(text):
    """Clean location text to extract main location name."""
    match = re.search(r"(di\s+)?([A-Za-z\s]+)$", text, re.IGNORECASE)
    return match.group(2).strip() if match else text.strip()

def extract_info(user_input):
    """Extract title, description, and location from user input if present."""
    patterns = {
        "judul": r"(judul|tema|topik|tentang)(nya| adalah| yaitu|:)?\s*(.+)",
        "deskripsi": r"(deskripsi|penjelasan|ringkasan)(nya| adalah| yaitu|:)?\s*(.+)",
        "lokasi": r"(lokasi|tempat)(nya| adalah| yaitu| di|:)?\s*(.+)"
    }
    
    for key, pattern in patterns.items():
        match = re.search(pattern, user_input, re.IGNORECASE)
        if match and user_inputs[key] is None:
            value = match.group(3).strip()
            if key == "lokasi":
                value = clean_location(value)
            user_inputs[key] = value
            print(f"Sistem: {random.choice(info_received_responses).format(key=key.capitalize(), value=value)}")

@torch.inference_mode()  # Enable inference mode for better GPU performance
def generate_proposal():
    """Generate proposal if all required information is available."""
    if not all(user_inputs.values()):
        missing = check_missing_info()
        print(f"Sistem: {random.choice(missing_info_responses).format(missing=', '.join(missing))}")
        return

    proposal_prompt = f"""
    Buatkan proposal dalam bahasa Indonesia berdasarkan informasi berikut:
    Judul: {user_inputs["judul"]}
    Deskripsi: {user_inputs["deskripsi"]}
    Lokasi: {user_inputs["lokasi"]}
    """

    response = llm.invoke(proposal_prompt)

    # Save to history
    conversation_history = load_history()
    conversation_history.append({
        "user": "User telah memberikan semua informasi.",
        "ai": response
    })
    save_history(conversation_history)

    print(f"Sistem: {random.choice(proposal_created_responses).format(response=response)}")

def elang_proposal_response():
    """Main conversation loop."""
    print(f"Sistem: Halo! Saya Elang Proposal, siap membantu Anda membuat proposal. (Running on {device.upper()})")
    conversation_history = load_history()
    
    while True:
        user_input = input("Anda: ").strip().lower()

        if not user_input:
            print("Sistem: Mohon berikan input yang valid. Ada yang bisa kubantu?")
            continue

        if user_input in ["quit", "keluar", "q"]:
            print("Sistem: Terima kasih telah menggunakan Elang Proposal. Sampai jumpa!")
            # Clean up GPU memory if using CUDA
            if device == "cuda":
                torch.cuda.empty_cache()
            break

        lemmatized_input = lemmatize_sentence(user_input)
        
        # Process proposal information
        if any(keyword in user_input for keyword in ["judul", "deskripsi", "lokasi", "proposal"]):
            extract_info(user_input)
            
            if all(user_inputs.values()):
                generate_proposal()
                continue
            
            missing_info = check_missing_info()
            if missing_info:
                print(f"Sistem: {random.choice(missing_info_responses).format(missing=', '.join(missing_info))}")
                continue
        
        # Check for general questions
        for question_key in general_questions.keys():
            if question_key in lemmatized_input:
                print(f"Sistem: {random.choice(general_questions[question_key])}")
                break
        else:
            print("Sistem: Untuk membuat proposal, saya perlu informasi berikut:")
            print("- Judul proposal")
            print("- Deskripsi atau penjelasan")
            print("- Lokasi kegiatan")
            print("Silakan berikan informasi tersebut satu per satu.")

if __name__ == "__main__":
    try:
        elang_proposal_response()
    finally:
        # Ensure GPU memory is cleared when the program exits
        if device == "cuda":
            torch.cuda.empty_cache()