import gradio as gr
import re
import joblib
import csv
import os
from huggingface_hub import hf_hub_download
from datasets import load_dataset, Dataset, concatenate_datasets
from huggingface_hub import HfApi, HfFolder
import pandas as pd

DATASET_REPO = "SleepyTerr/entrepreneurial_readiness_v2"

# Load model
model_repo_id = "SleepyTerr/entrepreneurial_readiness"
model_file = hf_hub_download(repo_id=model_repo_id, filename="readiness_model.joblib", repo_type="model")
model = joblib.load(model_file)

# Conversation state
expected_fields = ["age", "risk_tolerance", "sales_skills", "dependents", 
                   "monthly_income", "monthly_expenses", "entertainment", 
                   "savings", "assets"]

field_questions = {
    "age": "How old are you?",
    "risk_tolerance": "On a scale of 1-10, how much risk are you willing to take?",
    "sales_skills": "On a scale of 1-10, how would you rate your sales skills?",
    "dependents": "How many dependents do you take care of?",
    "monthly_income": "What is your monthly income?",
    "monthly_expenses": "What are your monthly expenses?",
    "entertainment": "How much do you spend on entertainment each month?",
    "savings": "How much do you currently have saved?",
    "assets": "What is the total value of your assets?"
}

def parse_info(message):
    """Extract possible numbers from free text"""
    numbers = re.findall(r"\d+", message)
    return [int(n) for n in numbers]

def generate_tips(user_data, score):
    """Give personalized tips depending on weaknesses"""
    tips = []
    if user_data.get("monthly_expenses", 0) > user_data.get("monthly_income", 0) * 0.8:
        tips.append("Try to lower your expenses so they don’t eat up most of your income.")
    if user_data.get("savings", 0) < user_data.get("monthly_income", 0) * 2:
        tips.append("Build an emergency fund with at least 2–3 months of income saved.")
    if user_data.get("assets", 0) < user_data.get("monthly_income", 0) * 6:
        tips.append("Work on building long-term assets to increase stability.")
    if user_data.get("sales_skills", 0) < 5:
        tips.append("Consider improving your sales or communication skills—they’re key for entrepreneurship.")
    if user_data.get("risk_tolerance", 0) < 4:
        tips.append("Think about whether you’re comfortable taking risks—entrepreneurship often involves uncertainty.")
    
    if not tips:
        tips.append("Great job! You’re on a strong path to entrepreneurship.")
    
    return tips

def save_to_db(user_data, score):
    try:
        # Convert to pandas row
        row = {**user_data, "score": score}
        df = pd.DataFrame([row])

        # Load existing dataset
        ds = load_dataset(DATASET_REPO, split="train")

        # Convert new row into a Dataset
        new_ds = Dataset.from_pandas(df)

        # Concatenate
        updated = concatenate_datasets([ds, new_ds])

        # Push back to Hugging Face
        updated.push_to_hub(DATASET_REPO, split="train")

        print("✅ Response saved to Hugging Face dataset")
    except Exception as e:
        print("⚠️ Failed to save:", e)

def chatbot(message, history, user_data, current_field):
    history = history or []
    user_data = user_data or {}
    current_field = current_field or expected_fields[0]

    # If current field not filled, try to extract it
    if current_field in expected_fields:
        nums = parse_info(message)
        if nums:
            user_data[current_field] = nums[0]
            # Move to next field
            next_index = expected_fields.index(current_field) + 1
            if next_index < len(expected_fields):
                current_field = expected_fields[next_index]
                history.append((message, field_questions[current_field]))
            else:
                # All info collected → make prediction
                features = [user_data.get(f, 0) for f in expected_fields]
                score = model.predict([features])[0]
                tips = generate_tips(user_data, score)
                tips_text = "\n- " + "\n- ".join(tips)

                # Save to database (CSV here, but can be swapped)
                save_to_db(user_data, score)

                history.append((message, f"✅ Your entrepreneurial readiness score is **{score:.1f}/100**!\n\nHere are some tips to improve:\n{tips_text}"))
                return history, user_data, None
        else:
            # Ask again in case of missing number
            history.append((message, f"Please provide a number for {current_field}."))
    else:
        history.append((message, "Hmm... I need a bit more info. Can you clarify?"))

    return history, user_data, current_field

with gr.Blocks(css="""
    body {background-color: #1a002b;}
    .chatbot {background-color: #2d0a45; border-radius: 12px; padding: 10px;}
    .message.user {background-color: #5a2d82; color: white; border-radius: 8px; padding: 8px; margin: 5px;}
    .message.bot {background-color: #3b0f58; color: #eee; border-radius: 8px; padding: 8px; margin: 5px;}
""") as demo:
    gr.Markdown("<h1 style='color:white; text-align:center;'>💡 Entrepreneurial Readiness Chatbot</h1>")
    gr.Markdown("<p style='color:#ddd; text-align:center;'>Hello! I'm an entrepreneurial readiness chatbot! 🧑‍💻<br>Tell me a little about yourself and I'll give you a score between 0 and 100!</p>")
    
    chatbot_ui = gr.Chatbot(label="Chat", elem_classes="chatbot", height=400)
    msg = gr.Textbox(placeholder="Type your scenario here...", label="Your Message")
    state = gr.State({})
    current_field = gr.State(expected_fields[0])

    def respond(message, history, user_data, current_field):
        return chatbot(message, history, user_data, current_field)

    msg.submit(respond, [msg, chatbot_ui, state, current_field], [chatbot_ui, state, current_field])

demo.launch()