from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

# --- 1. Load Roadmap and Rules ---
def load_roadmap_and_rules(roadmap_path="roadmap.txt", rules_path="rules.txt"):
    roadmap_content = {}
    with open(roadmap_path, 'r') as f:
        current_section = None
        for line in f:
            line = line.strip()
            if line.startswith("#"):
                current_section = line[1:].strip()
                roadmap_content[current_section] = ""
            elif current_section:
                roadmap_content[current_section] += line + "\n"
    with open(rules_path, 'r') as f:
        rules_content = f.read()
    return roadmap_content, rules_content

roadmap, rules = load_roadmap_and_rules()

# --- 2. Load the AI Model (Mistral 7B) ---
model_name = "mistralai/Mistral-7B-Instruct-v0.2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# --- 3. Function to Get Chatbot Response ---
def get_chatbot_response(user_query, project_phase, roadmap, rules, model, tokenizer):
    phase_roadmap_section = roadmap.get(project_phase, "General Guidance")
    context = f"""
    Project Roadmap (Phase: {project_phase}):
    {phase_roadmap_section}

    Project Rules:
    {rules}

    User Query: {user_query}

    ---
    Provide helpful guidance for this project.
    """
    prompt = f"<s>[INST] {context} [/INST]"
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=300)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response_start_index = response.find("[/INST]") + len("[/INST]") if "[/INST]" in response else 0
    cleaned_response = response[response_start_index:].strip()
    return cleaned_response

# --- 4. Gradio Interface Function ---
def chatbot_interface(user_query, project_phase):
    return get_chatbot_response(user_query, project_phase, roadmap, rules, model, tokenizer)

# --- 5. Gradio Interface Setup ---
iface = gr.Interface(
    fn=chatbot_interface,
    inputs=[
        gr.Textbox(label="Your Question"),
        gr.Dropdown(list(roadmap.keys()), label="Project Phase", value=list(roadmap.keys())[0])
    ],
    outputs="text",
    title="Project Guidance Chatbot",
    description="Ask questions about your project phase and get guidance."
)

if __name__ == "__main__":
    iface.launch()