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

# Initialize the OllamaLLM with the desired model
llm = OllamaLLM(model="llama3")

# 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'
}
FIELD_NAMES_ID = {
    "proposal_title": "Judul",
    "proposal_background": "Latar Belakang",
    "proposal_objective": "Tujuan",
    "proposal_description": "Deskripsi",
    "proposal_execution_date": "Tanggal Pelaksanaan",
    "proposal_end_date": "Tanggal Berakhir",
    "proposal_project_location": "Lokasi",
    "proposal_budget": "Budget"
}

user_inputs = {
    "proposal_created_date": datetime.now().strftime("%Y-%m-%d"),  # Automatically set current date
    "proposal_title": None,
    "proposal_background": None,
    "proposal_objective": None,
    "proposal_description": None,
    "proposal_execution_date": None,
    "proposal_end_date": None,  # New field for end date
    "proposal_project_location": None,
    "proposal_budget": None  # New field for budget
}

# 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 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 all proposal information from user input if present."""
    patterns = {
        "proposal_title": r"(judul|tema|topik|tentang)(nya| adalah| yaitu|:)?\s*(.+)",
        "proposal_background": r"(latar\s+belakang|background)(nya| adalah| yaitu|:)?\s*(.+)",
        "proposal_objective": r"(tujuan|objektif|sasaran)(nya| adalah| yaitu|:)?\s*(.+)",
        "proposal_description": r"(deskripsi|penjelasan|ringkasan)(nya| adalah| yaitu|:)?\s*(.+)", 
        "proposal_execution_date": r"(tanggal\s+pelaksanaan|waktu\s+pelaksanaan|jadwal|mulai|tanggal\s+mulai|dimulai)(nya| adalah| yaitu|:)?\s*(\d{1,2}\s+[a-zA-Z]+\s+\d{4})",
        "proposal_end_date": r"(tanggal\s+selesai|tanggal\s+berakhir|akhir\s+kegiatan|selesai)(nya| adalah| yaitu|:)?\s*(\d{1,2}\s+[a-zA-Z]+\s+\d{4})",
        "proposal_project_location": r"(lokasi|tempat)(nya| adalah| yaitu| di|:)?\s*(.+)",
        "proposal_budget": r"(budget|anggaran|biaya)(nya| adalah| yaitu|:)?\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 == "proposal_project_location":
                value = clean_location(value)
            user_inputs[key] = value
            print(f"Sistem: {FIELD_NAMES_ID.get(key, key.replace('proposal_', '').replace('_', ' ').title())} sudah tercatat: {value}.")

def check_missing_info():
    """Check which information is still missing."""
    missing_info_mapping = {
        "proposal_title": "Judul",
        "proposal_background": "Latar Belakang",
        "proposal_objective": "Tujuan",
        "proposal_description": "Deskripsi",
        "proposal_execution_date": "Tanggal Pelaksanaan",
        "proposal_end_date": "Tanggal Berakhir",
        "proposal_project_location": "Lokasi",
        "proposal_budget": "Budget"
    }
    
    return [missing_info_mapping[key] 
            for key, value in user_inputs.items() 
            if value is None and key != "proposal_created_date"]

def generate_proposal():
    """Generate proposal if all required information is available."""
    if not all(value is not None for key, value in user_inputs.items() if key != "proposal_created_date"):
        missing = check_missing_info()
        print(f"Sistem: Mohon lengkapi informasi berikut ya: {', '.join(missing)}. Supaya proposalnya bisa dibuat dengan baik.")
        return

    proposal_prompt = f"""
    Buatkan proposal dalam bahasa Indonesia berdasarkan informasi berikut:
    
    Tanggal Pembuatan: {user_inputs["proposal_created_date"]}
    Judul: {user_inputs["proposal_title"]}
    Latar Belakang: {user_inputs["proposal_background"]}
    Tujuan: {user_inputs["proposal_objective"]}
    Deskripsi: {user_inputs["proposal_description"]}
    Tanggal Pelaksanaan: {user_inputs["proposal_execution_date"]}
    Tanggal Berakhir: {user_inputs["proposal_end_date"]}
    Lokasi: {user_inputs["proposal_project_location"]}
    Budget: {user_inputs["proposal_budget"]}
    """

    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("Sistem: Halo! Saya Elang Proposal, siap membantu Anda membuat proposal.")
    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!")
            break

        lemmatized_input = lemmatize_sentence(user_input)
        
        # Process proposal information
        if any(keyword in user_input for keyword in ["judul", "latar belakang", "tujuan", "deskripsi", "tanggal", "lokasi", "proposal", "budget", "anggaran", "biaya", "selesai", "berakhir"]):
            extract_info(user_input)
            
            if all(value is not None for key, value in user_inputs.items() if key != "proposal_created_date"):
                generate_proposal()
                continue
            
            missing_info = check_missing_info()
            if missing_info:
                print(f"Sistem: Mohon lengkapi informasi berikut ya: {', '.join(missing_info)}. Supaya proposalnya bisa dibuat dengan baik.")
                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("- Latar belakang")
            print("- Tujuan") 
            print("- Deskripsi")
            print("- Tanggal pelaksanaan")
            print("- Tanggal berakhir")
            print("- Lokasi kegiatan")
            print("- Budget/Anggaran")
            print("Silakan berikan informasi tersebut satu per satu.")

if __name__ == "__main__":
    elang_proposal_response()