# main.py
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 Kegiatan",
    "proposal_budget": "Budget/Anggaran",
    "proposal_type": "Jenis Proposal"
}

user_inputs = {
    "proposal_created_date": datetime.now().strftime("%Y-%m-%d"),
    "proposal_title": None,
    "proposal_background": None,
    "proposal_objective": None,
    "proposal_description": None,
    "proposal_execution_date": None,
    "proposal_end_date": None,
    "proposal_project_location": None,
    "proposal_budget": None,
    "proposal_type": 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 save_proposal_to_file(proposal_text, filename=None):
    """Save the generated proposal to a file."""
    if filename is None:
        filename = f"proposal_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
    
    with open(filename, 'w', encoding='utf-8') as file:
        file.write(proposal_text)
    return filename

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 format_currency(amount_str):
    """Format currency string to standard format."""
    # Remove all non-digit characters
    amount = ''.join(filter(str.isdigit, amount_str))
    try:
        amount = int(amount)
        return f"Rp {amount:,}".replace(',', '.')
    except ValueError:
        return amount_str

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*(.+)",
        "proposal_type": r"(jenis\s+proposal|tipe\s+proposal)(nya| adalah| yaitu|:)?\s*(cash\s+advance|procurement|external)"
    }
    
    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)
            elif key == "proposal_budget":
                value = format_currency(value)
            elif key == "proposal_type":
                value = value.lower()
                if value not in ["cash advance", "procurement", "external"]:
                    print("Sistem: Jenis proposal tidak valid. Harap pilih antara Cash Advance, Procurement, atau External.")
                    continue
            
            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."""
    return [FIELD_NAMES_ID.get(key, key.replace('proposal_', '').replace('_', ' ').title()) 
            for key, value in user_inputs.items() 
            if value is None and key != "proposal_created_date"]

def generate_proposal():
    """Generate a detailed 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

    # Ask for proposal type if not already provided
    if user_inputs["proposal_type"] is None:
        print("Sistem: Sebelum membuat proposal, mohon tentukan jenis proposal:")
        print("1. Cash Advance")
        print("2. Procurement")
        print("3. External")
        proposal_type = input("Anda (ketik 1, 2, atau 3): ").strip()
        if proposal_type == "1":
            user_inputs["proposal_type"] = "cash advance"
        elif proposal_type == "2":
            user_inputs["proposal_type"] = "procurement"
        elif proposal_type == "3":
            user_inputs["proposal_type"] = "external"
        else:
            print("Sistem: Input tidak valid. Proposal tidak dapat dibuat.")
            return

    proposal_prompt = f"""
    Buatkan proposal lengkap dalam bahasa Indonesia berdasarkan informasi berikut. 
    Proposal harus detail dan komprehensif, minimal 2000 token, dengan struktur yang jelas dan pembahasan mendalam untuk setiap bagian.
    
    Format proposal sebagai berikut:
    
    PROPOSAL {user_inputs["proposal_type"].upper()}
    
    I. INFORMASI UMUM
    Tanggal Pembuatan: {user_inputs["proposal_created_date"]}
    Jenis Proposal: {user_inputs["proposal_type"]}
    
    II. JUDUL PROPOSAL
    {user_inputs["proposal_title"]}
    
    III. LATAR BELAKANG
    [Elaborasi latar belakang secara detail, minimal 3 paragraf, mencakup:]
    - Analisis situasi dan kondisi saat ini
    - Identifikasi masalah atau kebutuhan
    - Urgensi dan signifikansi proposal
    - Kesesuaian dengan visi dan misi organisasi
    - Dampak jangka panjang
    Gunakan dan kembangkan informasi ini: {user_inputs["proposal_background"]}
    
    IV. TUJUAN DAN SASARAN
    A. Tujuan Umum
    B. Tujuan Khusus (minimal 3 poin)
    C. Sasaran dan Target
    D. Indikator Keberhasilan
    E. Manfaat bagi Stakeholder
    [Kembangkan dari tujuan dasar ini: {user_inputs["proposal_objective"]}]
    
    V. DESKRIPSI PROGRAM
    A. Ringkasan Program
    B. Ruang Lingkup Program
    C. Metodologi Pelaksanaan
       1. Tahap Persiapan
       2. Tahap Implementasi
       3. Tahap Evaluasi
    D. Timeline dan Milestone (detail per minggu)
    E. Strategi Pelaksanaan
    F. Analisis Risiko dan Mitigasi
    [Elaborasi berdasarkan deskripsi ini: {user_inputs["proposal_description"]}]
    
    VI. WAKTU DAN TEMPAT PELAKSANAAN
    A. Waktu
       Tanggal Mulai: {user_inputs["proposal_execution_date"]}
       Tanggal Selesai: {user_inputs["proposal_end_date"]}
    B. Tempat: {user_inputs["proposal_project_location"]}
    C. Jadwal Detail Kegiatan (breakdown per minggu)
    D. Durasi dan Pembagian Waktu
    
    VII. RENCANA ANGGARAN
    A. Total Budget: {user_inputs["proposal_budget"]}
    B. Breakdown Anggaran:
       1. Biaya Operasional (40%)
       2. Biaya Material dan Peralatan (30%)
       3. Biaya Sumber Daya Manusia (20%)
       4. Biaya Tidak Terduga (10%)
    C. Justifikasi Anggaran per Komponen
    D. Analisis Cost-Benefit
    E. Rencana Pengelolaan Anggaran
    
    VIII. ORGANISASI PELAKSANA
    A. Struktur Tim
    B. Peran dan Tanggung Jawab
    C. Mekanisme Koordinasi
    D. Sistem Monitoring dan Evaluasi
    
    IX. PENUTUP
    A. Kesimpulan
    B. Harapan dan Dampak yang Diharapkan
    C. Ucapan Terima Kasih
    
    X. LAMPIRAN
    A. Timeline Detail
    B. Struktur Organisasi Tim Pelaksana
    C. Rincian Kebutuhan Teknis
    D. Format Laporan Evaluasi
    E. Checklist Persiapan
    
    Harap elaborasi setiap bagian secara detail dan profesional. Gunakan bahasa formal dan teknis yang sesuai dengan jenis proposal {user_inputs["proposal_type"]}. Setiap bagian harus memiliki penjelasan yang komprehensif dan saling terkait.
    """

    response = llm.invoke(proposal_prompt)
    
    # Save proposal to file
    filename = save_proposal_to_file(response)

    # Save to conversation 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)}")
    print(f"\nProposal telah disimpan dalam file: {filename}")

def elang_proposal_response():
    """Main conversation loop."""
    print("Sistem: Halo! Saya Elang Proposal, siap membantu Anda membuat proposal.")
    print("\nSaya akan membantu Anda membuat proposal yang lengkap dan terstruktur.")
    print("Ketik 'bantuan' untuk melihat panduan, atau 'keluar' untuk mengakhiri.")
    
    conversation_history = load_history()
    
    while True:
        user_input = input("\nAnda: ").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
            
        if user_input in ["bantuan", "help", "tolong"]:
            print("\nSistem: Untuk membuat proposal, saya perlu informasi berikut:")
            print("1. Judul proposal")
            print("2. Latar belakang")
            print("3. Tujuan")
            print("4. Deskripsi")
            print("5. Tanggal pelaksanaan")
            print("6. Tanggal berakhir")
            print("7. Lokasi kegiatan")
            print("8. Budget/Anggaran")
            print("9. Jenis Proposal (Cash Advance, Procurement, atau External)")
            print("\nAnda bisa memberikan informasi dengan format:")
            print("- Judul: [isi judul]")
            print("- Latar belakang: [isi latar belakang]")
            print("- Tujuan: [isi tujuan]")
            print("Dan seterusnya.")
            continue

        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", "jenis proposal"]):
            extract_info(user_input)
            
            if all(value is not None for key, value in user_inputs.items() if key != "proposal_created_date" and key != "proposal_type"):
                generate_proposal()
                continue
            
            missing_info = check_missing_info()
            if missing_info:
                response = random.choice(missing_info_responses)
                print(f"Sistem: {response.format(missing_fields=', '.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:
            if "status" in lemmatized_input:
                missing = check_missing_info()
                if missing:
                    print("\nSistem: Informasi yang masih diperlukan:")
                    for i, field in enumerate(missing, 1):
                        print(f"{i}. {field}")
                else:
                    print("\nSistem: Semua informasi sudah lengkap! Saya bisa membuatkan proposal sekarang.")
            else:
                print("\nSistem: Mohon berikan informasi yang diperlukan untuk proposal.")
                print("Informasi yang masih diperlukan:")
                missing = check_missing_info()
                for i, field in enumerate(missing, 1):
                    print(f"{i}. {field}")

if __name__ == "__main__":
    elang_proposal_response()