import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import torch
import re
import io
import docx
import json
from pypdf import PdfReader
from transformers import AutoTokenizer, AutoModelForCausalLM

# ─── 1. AI ENGINE SETUP ───────────────────────────────────────────────────
@st.cache_resource
def load_ai_model():
    MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        device_map="cpu",
        torch_dtype=torch.float32,
        trust_remote_code=True
    )
    return tokenizer, model

def ask_ai(system_prompt, user_prompt):
    tokenizer, model = load_ai_model()
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        generated_ids = model.generate(**model_inputs, max_new_tokens=512, temperature=0.1)
    
    response = tokenizer.decode(generated_ids[0][len(model_inputs.input_ids[0]):], skip_special_tokens=True)
    return response

# ─── 2. FILE PROCESSING ───────────────────────────────────────────────────
def extract_text_from_file(uploaded_file):
    text = ""
    try:
        if uploaded_file.name.endswith('.pdf'):
            reader = PdfReader(uploaded_file)
            for page in reader.pages:
                text += page.extract_text() + "\n"
        elif uploaded_file.name.endswith('.docx'):
            doc = docx.Document(io.BytesIO(uploaded_file.read()))
            text = "\n".join([para.text for para in doc.paragraphs])
    except Exception as e:
        st.error(f"Error reading {uploaded_file.name}: {e}")
    return text

# ─── 3. UI & STYLING ──────────────────────────────────────────────────────
st.set_page_config(page_title="Smart ATS AI", page_icon="✨", layout="wide")

st.markdown("""
<style>
    header { visibility: hidden; }
    .stApp {
        background: linear-gradient(rgba(255,255,255,0.6), rgba(255,255,255,0.6)),
                    url("https://i.postimg.cc/k57CDHgB/tnzyl.jpg") !important;
        background-size: cover !important;
        background-attachment: fixed !important;
    }
    div.stButton > button {
        border-radius: 50px !important;
        padding: 10px 30px !important;
        font-weight: bold !important;
        background: linear-gradient(135deg, #a855f7 0%, #7c3aed 100%) !important;
        color: white !important;
        border: none !important;
    }
    .input-card {
        background: rgba(255,255,255,0.92) !important;
        padding: 30px;
        border-radius: 20px;
        border: 1px solid #e9d5ff;
        box-shadow: 0 10px 25px rgba(0,0,0,0.05);
    }
    .status-box {
        padding: 12px;
        border-radius: 10px;
        background: #f3e8ff;
        border-left: 5px solid #7c3aed;
        margin-bottom: 15px;
        color: #5b21b6;
    }
</style>
""", unsafe_allow_html=True)

if 'current_page' not in st.session_state: st.session_state.current_page = 'landing'
if 'results' not in st.session_state: st.session_state.results = None

# ─── PAGE 1: LANDING ──────────────────────────────────────────────────────
if st.session_state.current_page == 'landing':
    col1, col2 = st.columns([1.1, 0.9])
    with col1:
        st.markdown('<div style="padding: 10% 5% 5% 10%;">', unsafe_allow_html=True)
        st.markdown('<h1 style="font-size:48px; font-weight:800; color:#1e293b;">YOUR AI PARTNER FOR SMART HIRING</h1>', unsafe_allow_html=True)
        if st.button("START AI ANALYSIS"):
            st.session_state.current_page = 'input'
            st.rerun()
        st.markdown('</div>', unsafe_allow_html=True)
    with col2:
        st.image("https://i.postimg.cc/c4nJwQxz/Recruiter-looks-at-a-perfect-candidate-cv-illustration.jpg")

# ─── PAGE 2: INPUT ────────────────────────────────────────────────────────
elif st.session_state.current_page == 'input':
    col_main, _ = st.columns([1.2, 0.8])
    with col_main:
        st.markdown('<div style="padding: 40px;">', unsafe_allow_html=True)
        # تصحيح زر الرجوع
        if st.button("Back to Home"): 
            st.session_state.current_page = 'landing'
            st.rerun()
        
        st.markdown('<div class="input-card">', unsafe_allow_html=True)
        st.markdown('<h2>Recruitment Setup</h2>', unsafe_allow_html=True)
        
        j_title = st.text_input("Job Title")
        j_desc = st.text_area("Job Description", height=150)
        uploaded_files = st.file_uploader("Upload Resumes", type=["pdf", "docx"], accept_multiple_files=True)

        if st.button("Analyze Now"):
            if j_title and j_desc and uploaded_files:
                final_results = []
                with st.spinner("AI Engine is working..."):
                    status_placeholder = st.empty()
                    progress_bar = st.progress(0)
                    
                    for idx, f in enumerate(uploaded_files):
                        status_placeholder.markdown(f'<div class="status-box">🔍 Processing: {f.name}</div>', unsafe_allow_html=True)
                        text_raw = extract_text_from_file(f)
                        sys_p = "Analyze resume against job description. Return JSON only."
                        user_p = f"Job: {j_title}\nResume: {text_raw[:2500]}\nOutput: score, reason, skills, exp."
                        
                        try:
                            ai_raw = ask_ai(sys_p, user_p)
                            match = re.search(r'\{.*\}', ai_raw, re.DOTALL)
                            if match:
                                data = json.loads(match.group())
                                final_results.append({
                                    "name": f.name.split('.')[0],
                                    "score": data.get("score", 0),
                                    "skills": data.get("skills", []),
                                    "exp": data.get("exp", 0),
                                    "status": "Accepted" if data.get("score", 0) >= 70 else "Rejected",
                                    "reason": data.get("reason", "N/A")
                                })
                        except:
                            continue
                        progress_bar.progress((idx + 1) / len(uploaded_files))
                
                st.session_state.results = sorted(final_results, key=lambda x: x['score'], reverse=True)
                st.session_state.job_title = j_title
                st.session_state.current_page = 'results'
                st.rerun()
        st.markdown('</div></div>', unsafe_allow_html=True)

# ─── PAGE 3: RESULTS ───────────────────────────────────────────────────────
elif st.session_state.current_page == 'results':
    st.markdown('<div style="padding:40px;">', unsafe_allow_html=True)
    if st.button("New Analysis"):
        st.session_state.current_page = 'input'
        st.rerun()
    
    st.title(f"Results for {st.session_state.job_title}")
    df = pd.DataFrame(st.session_state.results)
    
    c1, c2, c3 = st.columns(3)
    c1.metric("Total", len(df))
    c2.metric("Accepted", len(df[df['status']=='Accepted']))
    c3.metric("Avg Score", f"{int(df['score'].mean())}%")

    fig = px.bar(df, x='name', y='score', color='status')
    st.plotly_chart(fig, use_container_width=True)
    st.dataframe(df, use_container_width=True)
    st.markdown('</div>', unsafe_allow_html=True)