# ============================================================
#  AutoAd Studio — projekt: tomasz-svd
#  Etap 1 (Scraper) + Etap 2 (LLM) + Etap 3 (Audio + FFmpeg)
#  Etap 4 (Video Composer — FFmpeg CPU)
#  Wersja stabilna dla Windows — gotowa do uruchomienia
# ============================================================

import os
import tempfile
import json
import subprocess
from io import BytesIO
import uuid

import requests
from urllib.parse import urljoin, urlparse
from bs4 import BeautifulSoup
import trafilatura
from PIL import Image
import base64

import gradio as gr

# ------------------------------------------------------------
#  MODUŁY OPCJONALNE (fallbacki jeśli brak)
# ------------------------------------------------------------

try:
    import colorgram
    _HAS_COLORGRAM = True
except:
    _HAS_COLORGRAM = False

try:
    from transformers import AutoTokenizer, AutoModelForCausalLM
    import torch
    _HAS_TRANSFORMERS = True
except:
    _HAS_TRANSFORMERS = False

try:
    from audiocraft.models import MusicGen
    from audiocraft.data.audio import audio_write
    _HAS_MUSICGEN = True
except:
    _HAS_MUSICGEN = False

# ------------------------------------------------------------
#  ŚCIEŻKI I KONFIGURACJA
# ------------------------------------------------------------

FFMPEG_PATH = r"F:\ffmpeg\bin\ffmpeg.exe"  # dostosuj jeśli masz inną lokalizację
TMPDIR = tempfile.gettempdir()

def tmp_path(name):
    return os.path.join(TMPDIR, name)

def unique(name):
    return tmp_path(f"{uuid.uuid4().hex}_{name}")

# ------------------------------------------------------------
#  ETAP 1 — SCRAPER
# ------------------------------------------------------------

def fetch_html(domain):
    if not domain:
        return None, "Brak domeny"
    if not domain.startswith("http"):
        domain = "https://" + domain
    try:
        r = requests.get(domain, timeout=8, headers={"User-Agent": "Mozilla/5.0"})
        r.raise_for_status()
        return r.text, domain
    except Exception as e:
        return None, str(e)

def extract_text(html):
    try:
        return trafilatura.extract(html) or ""
    except:
        return ""

def find_images(soup, base_url, limit=4):
    imgs = []
    for img in soup.find_all("img"):
        src = img.get("src") or img.get("data-src")
        if not src:
            continue
        imgs.append(urljoin(base_url, src))
        if len(imgs) >= limit:
            break
    return imgs

def download_image(url):
    try:
        r = requests.get(url, timeout=8, headers={"User-Agent": "Mozilla/5.0"})
        r.raise_for_status()
        return Image.open(BytesIO(r.content)).convert("RGB")
    except:
        return None

def extract_colors_from_image(pil_img, n=5):
    if not _HAS_COLORGRAM:
        return []
    try:
        path = tmp_path("temp_color.jpg")
        pil_img.save(path, format="JPEG")
        colors = colorgram.extract(path, n)
        return [f"#{c.rgb.r:02x}{c.rgb.g:02x}{c.rgb.b:02x}" for c in colors]
    except:
        return []

def analyze_domain(domain):
    html, info = fetch_html(domain)
    if html is None:
        return {"error": f"Nie udało się pobrać strony: {info}"}

    soup = BeautifulSoup(html, "html.parser")

    title = soup.title.string.strip() if soup.title and soup.title.string else ""
    desc = ""
    meta = soup.find("meta", attrs={"name": "description"}) or soup.find("meta", attrs={"property": "og:description"})
    if meta and meta.get("content"):
        desc = meta["content"].strip()

    text = extract_text(html)
    short_text = text[:1000] + "..." if len(text) > 1000 else text

    base_url = info
    imgs = find_images(soup, base_url, limit=6)

    downloaded = []
    colors = []

    for url in imgs:
        img = download_image(url)
        if img:
            preview = img.copy()
            preview.thumbnail((320, 320))
            buf = BytesIO()
            preview.save(buf, format="JPEG")
            downloaded.append("data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode())

            if not colors:
                colors = extract_colors_from_image(img, n=5)

    domain_name = urlparse(base_url).netloc
    prompt = (
        f"Create a short energetic 15s ad for {domain_name}. "
        f"Tone: modern, friendly. Use brand colors {', '.join(colors) if colors else 'default colors'}. "
        f"Key message: {title or domain_name}. CTA: Visit {domain_name}."
    )

    return {
        "title": title,
        "description": desc,
        "text_snippet": short_text,
        "images": downloaded,
        "colors": colors,
        "prompt": prompt,
        "domain": domain_name,
    }

def ui_generate(domain):
    r = analyze_domain(domain)
    if "error" in r:
        return r["error"], "", "", "", []

    html = f"<h3>{r['title'] or r['domain']}</h3>"
    if r["description"]:
        html += f"<p><b>Meta description:</b> {r['description']}</p>"
    html += f"<p><b>Text snippet:</b> {r['text_snippet'][:600]}</p>"

    if r["colors"]:
        html += "<p><b>Detected colors:</b><br>"
        for c in r["colors"]:
            html += f"<span style='display:inline-block;width:28px;height:18px;background:{c};border:1px solid #ccc;margin-right:6px'></span> {c} "
        html += "</p>"

    if r["images"]:
        html += "<p><b>Images:</b><br>"
        for img in r["images"]:
            html += f"<img src='{img}' style='max-width:160px;margin-right:6px'/>"

    html += f"<h4>Auto prompt</h4><pre>{r['prompt']}</pre>"

    # dodatkowo zwracamy listę obrazów (base64) do Etapu 4
    return html, r["prompt"], r["domain"], r["text_snippet"][:800], r["images"]

# ------------------------------------------------------------
#  ETAP 2 — LLM (Phi-3)
# ------------------------------------------------------------

LLM_MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"

if _HAS_TRANSFORMERS:
    try:
        tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
        model = AutoModelForCausalLM.from_pretrained(LLM_MODEL_NAME, torch_dtype=torch.float32)
        model.eval()
    except:
        _HAS_TRANSFORMERS = False
        tokenizer = None
        model = None
else:
    tokenizer = None
    model = None

def generate_script(brand_prompt, domain, brand_text, length_sec, style):
    if not brand_prompt:
        return "Najpierw przeanalizuj domenę."

    if not _HAS_TRANSFORMERS:
        return json.dumps({
            "hook": f"{domain} — discover more!",
            "body": brand_text[:200],
            "cta": f"Visit {domain}",
            "overlay_text": ["Visit now", domain],
            "tone": style
        }, ensure_ascii=False, indent=2)

    try:
        length_sec = int(length_sec)
    except:
        length_sec = 15

    system_prompt = (
        "You are an ad script generator. "
        "Return JSON with: hook, body, cta, overlay_text, tone."
    )

    user_prompt = f"""
Brand: {domain}
Context: {brand_text}
Base prompt: {brand_prompt}
Length: {length_sec}s
Style: {style}
Return JSON only.
"""

    inp = tokenizer(f"<s>[INST] {system_prompt}\n{user_prompt} [/INST]", return_tensors="pt")

    with torch.no_grad():
        out = model.generate(
            **inp,
            max_new_tokens=400,
            do_sample=True,
            temperature=0.7,
            top_p=0.9
        )

    text = tokenizer.decode(out[0], skip_special_tokens=True)
    s = text.find("{")
    e = text.rfind("}")
    return text[s:e+1] if s != -1 and e != -1 else text

# ------------------------------------------------------------
#  ETAP 3 — AUDIO (MusicGen + FFmpeg + subprocess.run)
# ------------------------------------------------------------

def generate_silence(duration=15):
    path = tmp_path("silence.wav")
    import wave, struct
    sr = 22050
    n = int(sr * duration)
    with wave.open(path, "w") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        for _ in range(n):
            w.writeframes(struct.pack("<h", 0))
    return path

if _HAS_MUSICGEN:
    try:
        music_model = MusicGen.get_pretrained("facebook/musicgen-small")
    except:
        music_model = None
        _HAS_MUSICGEN = False
else:
    music_model = None

def generate_music(prompt, duration=15):
    out = tmp_path("music.wav")
    try:
        duration = int(duration)
    except:
        duration = 15

    if _HAS_MUSICGEN and music_model:
        try:
            music_model.set_generation_params(duration=duration)
            wav = music_model.generate([prompt])[0]
            audio_write(out, wav, music_model.sample_rate, format="wav")
            return out
        except:
            return generate_silence(duration)
    else:
        return generate_silence(duration)

def convert_to_mp3(wav_path):
    out = tmp_path("music.mp3")
    cmd = [
        FFMPEG_PATH, "-y",
        "-i", wav_path,
        "-c:a", "libmp3lame",
        "-q:a", "4",
        out
    ]
    try:
        subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=60)
        return out
    except:
        return wav_path

# ------------------------------------------------------------
#  ETAP 4 — VIDEO COMPOSER (FFmpeg CPU)
# ------------------------------------------------------------

def save_base64_image(b64, name):
    img_data = base64.b64decode(b64.split(",")[1])
    path = unique(name)
    with open(path, "wb") as f:
        f.write(img_data)
    return path

def create_slide(image_path, text, duration=3):
    out = unique("slide.mp4")

    draw = ""
    if text:
        safe_text = text.replace("'", "\\'")
        draw = (
            f"drawtext=text='{safe_text}':"
            f"fontcolor=white:fontsize=48:"
            f"x=(w-text_w)/2:y=h-200:"
            f"shadowcolor=black:shadowx=2:shadowy=2"
        )

    vf = draw if draw else "null"
    vf = vf + f",fade=t=in:st=0:d=0.5,fade=t=out:st={max(duration-0.5,0)}:d=0.5"

    cmd = [
        FFMPEG_PATH, "-y",
        "-loop", "1",
        "-i", image_path,
        "-t", str(duration),
        "-vf", vf,
        "-c:v", "libx264",
        "-pix_fmt", "yuv420p",
        out
    ]

    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return out

def concat_videos(video_list):
    list_path = unique("list.txt")
    with open(list_path, "w") as f:
        for v in video_list:
            f.write(f"file '{v}'\n")

    out = unique("merged.mp4")
    cmd = [
        FFMPEG_PATH, "-y",
        "-f", "concat",
        "-safe", "0",
        "-i", list_path,
        "-c", "copy",
        out
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return out

def add_audio_to_video(video_path, audio_path):
    out = unique("final.mp4")
    cmd = [
        FFMPEG_PATH, "-y",
        "-i", video_path,
        "-i", audio_path,
        "-c:v", "copy",
        "-c:a", "aac",
        "-shortest",
        out
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return out

def generate_video_from_b64(images_b64, script_json, audio_path):
    try:
        data = json.loads(script_json)
    except:
        return "Błędny JSON scenariusza."

    hook = data.get("hook", "")
    body = data.get("body", "")
    cta = data.get("cta", "")

    if not images_b64:
        return "Brak obrazów."

    img_paths = [save_base64_image(b, "img.jpg") for b in images_b64]

    slides = []
    if img_paths:
        slides.append(create_slide(img_paths[0], hook, duration=3))
    if len(img_paths) > 1:
        slides.append(create_slide(img_paths[1], body, duration=4))
    if len(img_paths) > 2:
        slides.append(create_slide(img_paths[2], cta, duration=3))

    merged = concat_videos(slides)
    final = add_audio_to_video(merged, audio_path)
    return final

# ------------------------------------------------------------
#  UI GRADIO
# ------------------------------------------------------------

with gr.Blocks() as demo:
    gr.Markdown("# tomasz-svd — AutoAd Studio (Etap 1 + 2 + 3 + 4)")

    # ---------------- TAB 1 ----------------
    with gr.Tab("1. Analiza domeny"):
        domain_in = gr.Textbox(label="Domena", placeholder="example.com")
        btn = gr.Button("Analizuj")
        out_html = gr.HTML()
        auto_prompt = gr.Textbox(label="Auto prompt", lines=4)
        brand_domain = gr.Textbox(label="Domena (czysta)")
        brand_text = gr.Textbox(label="Tekst (skrót)", lines=6)
        images_state = gr.State([])  # przechowujemy base64 obrazów dla Etapu 4
        gallery = gr.Gallery(label="Obrazy (podgląd)")

        def _ui_generate(domain):
            html, prompt, dom, txt, imgs = ui_generate(domain)
            # do galerii konwertujemy base64 -> PIL
            pil_list = []
            for b in imgs:
                try:
                    img_data = base64.b64decode(b.split(",")[1])
                    pil_list.append(Image.open(BytesIO(img_data)))
                except:
                    pass
            return html, prompt, dom, txt, imgs, pil_list

        btn.click(
            _ui_generate,
            inputs=domain_in,
            outputs=[out_html, auto_prompt, brand_domain, brand_text, images_state, gallery]
        )

    # ---------------- TAB 2 ----------------
    with gr.Tab("2. Scenariusz (LLM)"):
        length = gr.Slider(5, 60, value=15, step=5, label="Długość")
        style = gr.Dropdown(
            ["energetyczny TikTok", "premium elegancki", "luźny młodzieżowy", "poważny biznesowy"],
            value="energetyczny TikTok",
            label="Styl"
        )
        btn2 = gr.Button("Generuj scenariusz")
        script_out = gr.Code(label="JSON")

        btn2.click(
            generate_script,
            inputs=[auto_prompt, brand_domain, brand_text, length, style],
            outputs=script_out
        )

    # ---------------- TAB 3 ----------------
    with gr.Tab("3. Audio (MusicGen + FFmpeg)"):
        music_prompt = gr.Textbox(label="Prompt muzyczny", value="energetic modern ad music")
        duration = gr.Slider(5, 30, value=15, step=1, label="Długość muzyki")
        btn3 = gr.Button("Generuj muzykę")
        audio_out = gr.Audio(label="Muzyka", type="filepath")

        def _gen_audio(p, d):
            wav = generate_music(p, d)
            mp3 = convert_to_mp3(wav)
            return mp3

        btn3.click(_gen_audio, inputs=[music_prompt, duration], outputs=audio_out)

    # ---------------- TAB 4 ----------------
    with gr.Tab("4. Wideo (FFmpeg CPU)"):
        gr.Markdown("Generowanie finalnego wideo MP4 z obrazów (Etap 1), scenariusza (Etap 2) i muzyki (Etap 3).")

        script_in = gr.Code(label="Scenariusz (JSON z Etapu 2)")
        audio_in = gr.Audio(label="Muzyka z Etapu 3", type="filepath")
        btn4 = gr.Button("Generuj wideo")
        video_out = gr.Video(label="Finalne wideo MP4")

        def _gen_video(imgs_b64, script, audio_path):
            if not imgs_b64:
                return None
            if not script:
                return None
            if not audio_path:
                return None
            return generate_video_from_b64(imgs_b64, script, audio_path)

        btn4.click(
            _gen_video,
            inputs=[images_state, script_out, audio_in],
            outputs=video_out
        )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)