import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from pptx import Presentation
from pptx.util import Inches, Pt
from pptx.dml.color import RGBColor
import io, tempfile, os, re

st.set_page_config(page_title="PPTX Smart Enhancer", layout="wide")

# Load your hosted Hugging Face model
MODEL_ID = "Faisalkhany/llam1b_finetune"  # ← change this
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model     = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", torch_dtype="auto")
pipe      = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Templates & fonts
TEMPLATES = {
    "Office Default": {
        "font":"Calibri",
        "colors":{"primary":"#2B579A","secondary":"#5B9BD5","text":"#000000","background":"#FFFFFF"},
        "sizes":{"title":44,"heading":32,"body":18},"background":True
    },
    "Modern Office": {
        "font":"Segoe UI",
        "colors":{"primary":"#404040","secondary":"#A5A5A5","text":"#FFFFFF","background":"#121212"},
        "sizes":{"title":36,"heading":24,"body":16},"background":True
    },
    "Corporate": {
        "font":"Arial",
        "colors":{"primary":"#1976D2","secondary":"#BBDEFB","text":"#212121","background":"#FAFAFA"},
        "sizes":{"title":40,"heading":28,"body":18},"background":False
    }
}
FONT_CHOICES = {f:f for f in ["Arial","Calibri","Times New Roman","Segoe UI","Verdana","Georgia"]}

# helpers
def extract_slide_text(slide):
    return "\n".join(p.text for shp in slide.shapes if hasattr(shp,"text_frame")
                     for p in shp.text_frame.paragraphs)

def split_formatted_runs(text):
    segs,cur,bold,italic=[],[],False,False
    parts=re.split(r'(\*|_)',text)
    for t in parts:
        if t in ("*","_"):
            if cur:
                segs.append({"text":"".join(cur),"bold":bold,"italic":italic})
                cur=[]
            if t=="*": bold=not bold
            else:       italic=not italic
        else: cur.append(t)
    if cur: segs.append({"text":"".join(cur),"bold":bold,"italic":italic})
    return segs

def apply_formatting(tf,content,ds):
    tf.clear()
    for line in [l for l in content.split("\n") if l.strip()]:
        clean=re.sub(r'^#+\s*','',line).strip()
        if re.match(r'^(What is|Understanding|Key|Conclusion)',clean,re.IGNORECASE):
            p=tf.add_paragraph(); p.text=clean
            p.font.size=Pt(ds["heading_size"]); p.font.bold=True; p.space_after=Pt(12)
            continue
        if re.match(r'^[\-\*\•] ',clean):
            p=tf.add_paragraph(); p.level=0
            txt=clean[2:].strip()
            for seg in split_formatted_runs(txt):
                r=p.add_run(); r.text=seg["text"]; r.font.bold=seg["bold"]; r.font.italic=seg["italic"]
            p.font.size=Pt(ds["body_size"]); continue
        p=tf.add_paragraph()
        for seg in split_formatted_runs(clean):
            r=p.add_run(); r.text=seg["text"]; r.font.bold=seg["bold"]; r.font.italic=seg["italic"]
        p.font.size=Pt(ds["body_size"]); p.space_after=Pt(6)

def apply_design(slide,ds):
    cols=ds["colors"]
    for shp in slide.shapes:
        if hasattr(shp,"text_frame"):
            for p in shp.text_frame.paragraphs:
                for r in p.runs:
                    r.font.name=ds["font"]
                    r.font.color.rgb=RGBColor.from_string(cols["text"][1:])
                    if shp==getattr(slide.shapes,"title",None):
                        r.font.color.rgb=RGBColor.from_string(cols["primary"][1:])
                        r.font.size=Pt(ds["title_size"]); r.font.bold=True
    if ds["set_background"]:
        slide.background.fill.solid()
        slide.background.fill.fore_color.rgb=RGBColor.from_string(cols["background"][1:])

def enhance_slide(slide,prompt,ds):
    orig=extract_slide_text(slide)
    if not orig: return
    out=pipe(f"Improve this slide:\n{orig}\n\nInstructions: {prompt}",
             max_length=512, do_sample=False)[0]["generated_text"]
    title_sp=getattr(slide.shapes,"title",None)
    for shp in list(slide.shapes):
        if shp is not title_sp:
            shp._element.getparent().remove(shp._element)
    left,top=Inches(1),Inches(1.5 if title_sp else 0.5)
    tb=slide.shapes.add_textbox(left,top,Inches(8),Inches(5)); tf=tb.text_frame; tf.word_wrap=True
    apply_formatting(tf,out,ds); apply_design(slide,ds)

def process_presentation(uploaded,prompt,ds):
    with tempfile.NamedTemporaryFile(delete=False,suffix=".pptx") as tmp:
        tmp.write(uploaded.read()); path=tmp.name
    prs=Presentation(path)
    for s in prs.slides: enhance_slide(s,prompt,ds)
    os.unlink(path)
    return prs

def main():
    st.title("PPTX Smart Enhancer")
    uploaded=st.file_uploader("Upload PPTX",type=["pptx"])
    prompt  =st.text_area("Enhancement Instructions",height=100)
    with st.expander("🎨 Design Settings",True):
        tpl=st.selectbox("Template",list(TEMPLATES.keys()))
        if tpl=="Custom":
            font=st.selectbox("Font",list(FONT_CHOICES)); ts=st.slider("Title size",12,60,32)
            hs=st.slider("Heading size",12,48,24); bs=st.slider("Body size",8,36,18)
            sb=st.checkbox("Apply background",True)
            cols={k:st.color_picker(k.capitalize(),c) for k,c in TEMPLATES["Office Default"]["colors"].items()}
        else:
            t=TEMPLATES[tpl]; font,ts,hs,bs,sb=t["font"],*t["sizes"].values(),t["background"]
            cols=t["colors"]
    ds={"font":FONT_CHOICES[font],"colors":cols,
        "title_size":ts,"heading_size":hs,"body_size":bs,"set_background":sb}
    if st.button("Enhance")and uploaded and prompt:
        out_prs=process_presentation(uploaded,prompt,ds)
        buf=io.BytesIO(); out_prs.save(buf); buf.seek(0)
        st.download_button("Download Enhanced PPTX",data=buf,
            file_name="enhanced.pptx",
            mime="application/vnd.openxmlformats-officedocument.presentationml.presentation")

if __name__=="__main__":
    main()