import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from simple_salesforce import Salesforce
import os
import base64
import datetime
from dotenv import load_dotenv
from fpdf import FPDF
import shutil
import html

# Load environment variables
load_dotenv()

# Required env vars check
required_env_vars = ['SF_USERNAME', 'SF_PASSWORD', 'SF_SECURITY_TOKEN']
missing_vars = [var for var in required_env_vars if not os.getenv(var)]
if missing_vars:
    raise EnvironmentError(f"Missing required environment variables: {missing_vars}")

# Load model and tokenizer
model_name = "distilgpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
model.config.pad_token_id = tokenizer.pad_token_id

PROMPT_TEMPLATE = """You are an AI assistant for construction supervisors. Given the role, project, milestones, and a reflection log, generate:

1. A Daily Checklist with clear and concise tasks based on the role and milestones. 
   Split the checklist into day-by-day tasks for a specified time period (e.g., one week).
2. Focus Suggestions based on concerns or keywords in the reflection log. Provide at least 2 suggestions.

Inputs:
Role: {role}
Project ID: {project_id}
Milestones: {milestones}
Reflection Log: {reflection}

Output Format:
Checklist (Day-by-Day):
- Day 1: 
   - Task 1
   - Task 2
- Day 2: 
   - Task 1
   - Task 2
...
Suggestions:
- 
"""

def clean_text_for_pdf(text):
    return html.unescape(text).encode('latin-1', 'replace').decode('latin-1')

def save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions):
    now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    filename = f"report_{supervisor_name}_{project_id}_{now}.pdf"
    file_path = f"./reports/{filename}"
    os.makedirs("reports", exist_ok=True)

    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", 'B', 14)
    pdf.cell(200, 10, txt="Supervisor Daily Report", ln=True, align="C")
    pdf.set_font("Arial", size=12)
    pdf.cell(200, 10, txt=clean_text_for_pdf(f"Role: {role}"), ln=True)
    pdf.cell(200, 10, txt=clean_text_for_pdf(f"Supervisor: {supervisor_name}"), ln=True)
    pdf.cell(200, 10, txt=clean_text_for_pdf(f"Project ID: {project_id}"), ln=True)
    pdf.ln(5)
    pdf.set_font("Arial", 'B', 12)
    pdf.cell(200, 10, txt="Daily Checklist", ln=True)
    pdf.set_font("Arial", size=12)
    for line in checklist.split("\n"):
        pdf.multi_cell(0, 10, clean_text_for_pdf(line))
    pdf.ln(5)
    pdf.set_font("Arial", 'B', 12)
    pdf.cell(200, 10, txt="Focus Suggestions", ln=True)
    pdf.set_font("Arial", size=12)
    for line in suggestions.split("\n"):
        pdf.multi_cell(0, 10, clean_text_for_pdf(line))
    pdf.output(file_path)

    temp_pdf_path = "/tmp/" + os.path.basename(file_path)
    shutil.copy(file_path, temp_pdf_path)
    return temp_pdf_path, filename

def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_path, pdf_name, checklist, suggestions):
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )
        with open(pdf_path, "rb") as f:
            encoded = base64.b64encode(f.read()).decode()
        content = sf.ContentVersion.create({
            'Title': pdf_name,
            'PathOnClient': pdf_name,
            'VersionData': encoded
        })
        content_id = content['id']
        download_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_id}"
        query = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
        if query['totalSize'] == 0:
            return ""
        supervisor_id = query['records'][0]['Id']
        project_query = sf.query(f"SELECT Id FROM Project__c WHERE Name = '{project_id}' LIMIT 1")
        if project_query['totalSize'] == 0:
            return ""
        project_id_sf = project_query['records'][0]['Id']
        sf.Supervisor_AI_Coaching__c.create({
            'Project_ID__c': project_id_sf,
            'Supervisor_ID__c': supervisor_id,
            'Daily_Checklist__c': checklist,
            'Suggested_Tips__c': suggestions,
            'Download_Link__c': download_url
        })
        return download_url
    except:
        return ""

def get_roles_from_salesforce():
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )
        result = sf.query("SELECT Role__c FROM Supervisor__c WHERE Role__c != NULL")
        return list(set(record['Role__c'] for record in result['records']))
    except:
        return []

def get_supervisor_name_by_role(role):
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )
        result = sf.query(f"SELECT Name FROM Supervisor__c WHERE Role__c = '{role}'")
        return [record['Name'] for record in result['records']]
    except:
        return []

def get_projects_for_supervisor(supervisor_name):
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )
        result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
        if result['totalSize'] == 0:
            return ""
        supervisor_id = result['records'][0]['Id']
        project_result = sf.query(f"SELECT Name FROM Project__c WHERE Supervisor_ID__c = '{supervisor_id}' LIMIT 1")
        return project_result['records'][0]['Name'] if project_result['totalSize'] > 0 else ""
    except:
        return ""

def get_dashboard_data_from_salesforce(supervisor_name, project_id):
    try:
        sf = Salesforce(
            username=os.getenv('SF_USERNAME'),
            password=os.getenv('SF_PASSWORD'),
            security_token=os.getenv('SF_SECURITY_TOKEN'),
            domain=os.getenv('SF_DOMAIN', 'login')
        )
        query = f"""
        SELECT Daily_Checklist__c, Suggested_Tips__c, Download_Link__c, CreatedDate 
        FROM Supervisor_AI_Coaching__c 
        WHERE Supervisor_ID__r.Name = '{supervisor_name}' AND Project_ID__r.Name = '{project_id}'
        ORDER BY CreatedDate DESC
        LIMIT 1
        """
        result = sf.query(query)
        if result['totalSize'] == 0:
            return "No dashboard data found.", "", "", ""
        record = result['records'][0]
        return (
            record.get('Daily_Checklist__c', 'N/A'),
            record.get('Suggested_Tips__c', 'N/A'),
            record.get('Download_Link__c', ''),
            record.get('CreatedDate', 'Unknown')
        )
    except Exception as e:
        return f"Error loading dashboard: {e}", "", "", ""

def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
    if not all([role, supervisor_name, project_id, milestones, reflection]):
        return "❗ Please fill all fields.", "", None, ""
    prompt = PROMPT_TEMPLATE.format(role=role, project_id=project_id, milestones=milestones, reflection=reflection)
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
    try:
        with torch.no_grad():
            outputs = model.generate(
                inputs['input_ids'],
                max_new_tokens=150,
                no_repeat_ngram_size=2,
                do_sample=True,
                top_p=0.9,
                temperature=0.7,
                pad_token_id=tokenizer.pad_token_id
            )
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    except:
        return "", "", None, ""

    def extract_between(text, start, end):
        s = text.find(start)
        e = text.find(end, s) if end else len(text)
        return text[s + len(start):e].strip() if s != -1 else ""

    checklist = extract_between(result, "Checklist:\n", "Suggestions:")
    suggestions = extract_between(result, "Suggestions:\n", None)

    if not checklist.strip():
        checklist = "- Perform daily safety inspection"
    if not suggestions.strip():
        suggestions = "- Monitor team coordination\n- Review safety protocols with the team"

    pdf_path, pdf_name = save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions)
    pdf_url = upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_path, pdf_name, checklist, suggestions)

    if pdf_url:
        suggestions += f"\n\n🔗 [Download PDF Report]({pdf_url})"

    return checklist, suggestions, pdf_path, pdf_name

def create_interface():
    roles = get_roles_from_salesforce()
    with gr.Blocks(theme="soft", css=".footer { display: none; }") as demo:
        gr.Markdown("## 🧠 AI-Powered Supervisor Assistant")

        with gr.Tab("✏️ Generate Checklist"):
            with gr.Row():
                role = gr.Dropdown(choices=roles, label="Role")
                supervisor_name = gr.Dropdown(choices=[], label="Supervisor Name")
                project_id = gr.Textbox(label="Project ID", interactive=False)

            milestones = gr.Textbox(label="Milestones (comma-separated KPIs)")
            reflection = gr.Textbox(label="Reflection Log", lines=4)

            with gr.Row():
                generate = gr.Button("Generate")
                clear = gr.Button("Clear")
                refresh = gr.Button("🔄 Refresh Roles")

            checklist_output = gr.Textbox(label="✅ Daily Checklist")
            suggestions_output = gr.Textbox(label="💡 Focus Suggestions")
            download_button = gr.File(label="⬇ Download Report")
            pdf_link = gr.HTML()

            role.change(fn=lambda r: gr.update(choices=get_supervisor_name_by_role(r)), inputs=role, outputs=supervisor_name)
            supervisor_name.change(fn=get_projects_for_supervisor, inputs=supervisor_name, outputs=project_id)

            def handle_generate(role, supervisor_name, project_id, milestones, reflection):
                checklist, suggestions, pdf_path, pdf_name = generate_outputs(role, supervisor_name, project_id, milestones, reflection)
                return checklist, suggestions, pdf_path, pdf_name

            generate.click(fn=handle_generate,
                           inputs=[role, supervisor_name, project_id, milestones, reflection],
                           outputs=[checklist_output, suggestions_output, download_button, pdf_link])

            clear.click(fn=lambda: ("", "", "", "", ""),
                        inputs=None,
                        outputs=[role, supervisor_name, project_id, milestones, reflection])

            refresh.click(fn=lambda: gr.update(choices=get_roles_from_salesforce()), outputs=role)

        with gr.Tab("📊 Supervisor Dashboard"):
            dash_supervisor = gr.Textbox(label="Supervisor Name", placeholder="e.g., SUP-056")
            dash_project = gr.Textbox(label="Project ID", placeholder="e.g., PROJ-078")
            load_dash = gr.Button("📥 Load Dashboard")
            dash_output = gr.HTML()

            def show_dashboard_html(sup_name, proj_id):
                checklist, suggestions, link, date = get_dashboard_data_from_salesforce(sup_name, proj_id)
                link_html = f'<a href="{link}" target="_blank">Download PDF</a>' if link else "No report available"
                html_card = f"""
                <div style='border:1px solid #ccc; border-radius:10px; padding:20px; background-color:#f9f9f9'>
                    <h3>📋 <u>Dashboard Summary</u></h3>
                    <p><b>Supervisor:</b> {sup_name}</p>
                    <p><b>Project ID:</b> {proj_id}</p>
                    <p><b>🗓️ Last Updated:</b> {date}</p>
                    <hr>
                    <h4>✅ <u>Daily Checklist</u></h4>
                    <pre style='background:#eef; padding:10px; border-radius:5px;'>{checklist}</pre>
                    <h4>💡 <u>Focus Suggestions</u></h4>
                    <pre style='background:#ffe; padding:10px; border-radius:5px;'>{suggestions}</pre>
                    <p>{link_html}</p>
                </div>
                """
                return html_card

            load_dash.click(fn=show_dashboard_html, inputs=[dash_supervisor, dash_project], outputs=dash_output)

    return demo

if __name__ == "__main__":
    app = create_interface()
    app.launch()