import os
import gradio as gr
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import pandas as pd
import plotly.express as px
from datetime import datetime

# =========================================================
# CONFIG
# =========================================================
EMBEDDING_MODEL = "all-MiniLM-L6-v2"
GENERATION_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
TOP_K = 3
MAX_NEW_TOKENS = 200

# =========================================================
# LOAD DOCUMENTS WITH METADATA
# =========================================================
def load_documents():
    documents = []
    folders = {
        "knowledge_base": "core_profile",
        "synthetic_data": "extended_profile"
    }

    for folder, category in folders.items():
        if not os.path.exists(folder):
            continue
        for file in os.listdir(folder):
            if file.endswith(".txt"):
                with open(os.path.join(folder, file), "r", encoding="utf-8") as f:
                    content = f.read().strip()
                    if content:
                        documents.append({
                            "content": content,
                            "source": file,
                            "category": category
                        })

    if not documents:
        documents.append({
            "content": "Name: Aniket Sirsikar. Education: PGDM SPJIMR. Skills: Product Management, Data Analytics, AI Systems.",
            "source": "fallback.txt",
            "category": "core_profile"
        })

    return documents

# =========================================================
# CHUNKING
# =========================================================
def chunk_documents(documents, max_words=100):
    chunks = []
    overlap = 20

    for doc in documents:
        words = doc["content"].split()
        for i in range(0, len(words), max_words - overlap):
            chunk_text = " ".join(words[i:i + max_words])
            if chunk_text.strip():
                chunks.append({
                    "text": chunk_text,
                    "source": doc["source"],
                    "category": doc["category"]
                })
    return chunks

# =========================================================
# VECTOR STORE
# =========================================================
def build_vector_store(chunks):
    embedder = SentenceTransformer(EMBEDDING_MODEL)
    texts = [c["text"] for c in chunks]
    embeddings = embedder.encode(texts, convert_to_numpy=True).astype("float32")

    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)

    return index, chunks, embedder

# =========================================================
# RETRIEVAL
# =========================================================
def retrieve(query, index, chunks, embedder):
    q_emb = embedder.encode([query], convert_to_numpy=True).astype("float32")
    _, idxs = index.search(q_emb, min(TOP_K, len(chunks)))
    return [chunks[i] for i in idxs[0]]

# =========================================================
# GENERATION ENGINE (LLAMA – CAUSAL)
# =========================================================
class AnswerGenerator:
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained(GENERATION_MODEL)
        self.model = AutoModelForCausalLM.from_pretrained(GENERATION_MODEL)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self.device)

    def generate(self, query, context_chunks):
        context = "\n".join([c["text"] for c in context_chunks[:3]])

        prompt = (
            "<|system|>\n"
            "You are an AI professional profile assistant.\n"
            "Answer using ONLY the information provided.\n"
            "Be concise and recruiter-facing.\n"
            "<|user|>\n"
            f"Information:\n{context}\n\n"
            f"Question: {query}\n"
            "<|assistant|>\n"
        )

        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=False,
                temperature=0.2,
                repetition_penalty=1.1,
                eos_token_id=self.tokenizer.eos_token_id
            )

        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        answer = decoded.split("<|assistant|>")[-1].strip()
        return answer

# =========================================================
# QUERY ENHANCEMENT
# =========================================================
def enhance_query(query):
    mapping = {
        "skills": "What are my professional and technical skills?",
        "roles": "What roles suit me based on my profile?",
        "projects": "What projects have I worked on?",
        "experience": "What is my professional experience?",
        "education": "What is my educational background?",
        "goals": "What are my career goals?"
    }
    for k, v in mapping.items():
        if k in query.lower():
            return v
    return query

# =========================================================
# INITIALIZE SYSTEM
# =========================================================
documents = load_documents()
chunks = chunk_documents(documents)
index, chunks, embedder = build_vector_store(chunks)
generator = AnswerGenerator()

def answer_question(user_query):
    enhanced = enhance_query(user_query)
    retrieved = retrieve(enhanced, index, chunks, embedder)
    return generator.generate(enhanced, retrieved)



# =========================================================
# PREMIUM UI / UX (FULL VERSION – NO CORE CHANGES)
# =========================================================


# -----------------------------
# CV PATH
# -----------------------------
def get_cv_path():
    path = "cv/aniket_sirsikar.pdf"
    return path if os.path.exists(path) else None

# -----------------------------
# CONTACT FORM
# -----------------------------
def submit_contact(name, email, company, message):
    if not name or not email or not message:
        return "❌ Please fill in Name, Email and Message."

    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    entry = f"""
========================================
{timestamp}
Name: {name}
Email: {email}
Company: {company}
Message: {message}
========================================
"""

    try:
        with open("recruiter_contacts.txt", "a") as f:
            f.write(entry)
    except:
        pass

    return f"✅ Thank you {name}. I will reach out shortly."

# -----------------------------
# ANALYTICS CHARTS
# -----------------------------
def create_skills_chart():
    df = pd.DataFrame({
        "Skill": ["Python", "SQL", "Machine Learning", "Product Mgmt", "Analytics", "AI Systems"],
        "Level": [95, 90, 85, 85, 90, 80],
        "Category": ["Programming", "Programming", "AI/ML", "Product", "Analytics", "AI/ML"]
    })

    fig = px.bar(
        df,
        x="Skill",
        y="Level",
        color="Category",
        title="Core Skill Areas",
        color_discrete_map={
            "Programming": "#3498db",
            "AI/ML": "#e74c3c",
            "Analytics": "#f39c12",
            "Product": "#9b59b6"
        }
    )

    fig.update_layout(
        plot_bgcolor="white",
        paper_bgcolor="white",
        height=350
    )

    return fig


def create_experience_pie():
    df = pd.DataFrame({
        "Area": ["Product", "Analytics", "AI/ML", "Leadership"],
        "Share": [30, 30, 25, 15]
    })

    fig = px.pie(
        df,
        values="Share",
        names="Area",
        title="Experience Distribution"
    )

    fig.update_layout(
        plot_bgcolor="white",
        paper_bgcolor="white",
        height=400
    )

    return fig


def create_projects_timeline():
    df = pd.DataFrame({
        "Project": ["AI Twin", "Fraud Detection", "Analytics Dashboard", "AI Workflow Agents"],
        "Start": pd.to_datetime(["2024-01", "2023-01", "2023-06", "2025-01"]),
        "End": pd.to_datetime(["2024-04", "2023-09", "2023-12", "2025-04"]),
        "Impact": ["High", "High", "Medium", "Very High"]
    })

    fig = px.timeline(
        df,
        x_start="Start",
        x_end="End",
        y="Project",
        color="Impact"
    )

    fig.update_layout(
        plot_bgcolor="white",
        paper_bgcolor="white",
        height=350
    )

    return fig

# -----------------------------
# CUSTOM CSS
# -----------------------------
custom_css = """
.profile-summary {
    background: #2c3e50 !important;
    border: 2px solid #34495e;
    border-radius: 12px;
    padding: 1.5rem;
}
.profile-summary h3, 
.profile-summary p, 
.profile-summary strong {
    color: #ecf0f1 !important;
}
.section-header {
    background: #34495e;
    padding: 1rem 1.5rem;
    border-left: 5px solid #3498db;
    border-radius: 6px;
    margin-top: 1.5rem;
}
.section-header h3 {
    color: #ecf0f1 !important;
    margin: 0;
}
"""

# =========================================================
# UI LAYOUT
# =========================================================


# =========================================================
# CHAT RESPONSE HANDLER
# =========================================================
def respond(message, history):
    if not message or not message.strip():
        return history, history

    answer = answer_question(message)

    history = history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": answer}
    ]

    return history, history

    
with gr.Blocks(css=custom_css, title="Aniket Sirsikar – Digital Twin") as demo:

    # Header
    gr.HTML("""
        <div style='background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                    padding: 2.5rem; border-radius: 12px; color: white; 
                    text-align: center; margin-bottom: 1.5rem;'>
            <h1 style='margin:0; font-size: 2.5rem;'>Aniket Sirsikar</h1>
            <p style='margin:0.8rem 0 0 0; font-size: 1.2rem;'>
                Product Manager | Data Analytics | AI Systems
            </p>
        </div>
    """)

    with gr.Row():
        # Chat column
        with gr.Column(scale=6):

            gr.HTML("<div class='section-header'><h3>💬 Chat with Aniket's digital twin </h3></div>")

            chatbot = gr.Chatbot(label="AI Twin", type="messages")
            state = gr.State([])

            user_input = gr.Textbox(
                placeholder="Ask about skills, projects, experience...",
                label="Your Question"
            )

            with gr.Row():
                send = gr.Button("Ask", variant="primary")
                clear = gr.Button("Clear")

            gr.Markdown("**Quick Questions:**")

            with gr.Row():
                q1 = gr.Button("My skills ?")
                q2 = gr.Button("My projects ?")

            with gr.Row():
                q3 = gr.Button("My interests ?")
                q4 = gr.Button("Suitable roles ?")

        # Profile column
        with gr.Column(scale=4):

            gr.HTML("""
                <div class="profile-summary">
                <h3>📋 Profile Summary</h3>
                    <p>
                        I am an MBA candidate at SPJIMR specializing in Product Management and AI Systems.
                        My work combines analytical depth, structured thinking, and execution rigor to design
                        products that solve high-impact, large-scale problems.
                    </p>
                    <p>
                        With experience spanning public sector operations and AI-driven technology environments,
                        I have led initiatives involving large-scale data analysis, workflow optimization,
                        and the development of insight-driven dashboards and decision tools.
                    </p>
                    <p>
                        I operate comfortably at the intersection of data, product, and systems thinking —
                        translating ambiguity into clarity, insights into strategy, and strategy into execution.
                    </p>
                    <p>
                        My long-term focus is on building scalable, AI-enabled platforms that integrate analytics,
                        automation, and user experience into cohesive, high-leverage product ecosystems.
                    </p>
                </div>
            """)

            

            cv_path = get_cv_path()
            if cv_path:
                gr.DownloadButton("📥 Download CV", value=cv_path)

    # Analytics
    gr.HTML("<div class='section-header'><h3>📊 Recruiter Insights</h3></div>")

    with gr.Accordion("View Analytics", open=False):
        gr.Plot(create_skills_chart())
        gr.Plot(create_experience_pie())
        gr.Plot(create_projects_timeline())

    # Contact Form
    gr.HTML("<div class='section-header'><h3>📬 Contact</h3></div>")

    with gr.Row():
        with gr.Column():
            name = gr.Textbox(label="Name *")
            email = gr.Textbox(label="Email *")
        with gr.Column():
            company = gr.Textbox(label="Company")
            message = gr.Textbox(label="Message *", lines=3)

    submit = gr.Button("Send Message")
    status = gr.Textbox(show_label=False)

    # Events
    send.click(respond, [user_input, state], [chatbot, state])
    user_input.submit(respond, [user_input, state], [chatbot, state])
    clear.click(lambda: ([], []), None, [chatbot, state])

    q1.click(lambda: respond("What are my skills?", state.value), None, [chatbot, state])
    q2.click(lambda: respond("What projects have I done?", state.value), None, [chatbot, state])
    q3.click(lambda: respond("What are my interests?", state.value), None, [chatbot, state])
    q4.click(lambda: respond("What role suits me?", state.value), None, [chatbot, state])

    submit.click(submit_contact, [name, email, company, message], status)

demo.queue().launch()