# core.py
# Wspólna logika: Scraper, LLM, Audio, Video (Etap 1–4)

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

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

# opcjonalne moduły
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

FFMPEG_PATH = r"F:\ffmpeg-2026-03-30-git-e54e117998-full_build\ffmpeg-2026-03-30-git-e54e117998-full_build\bin\ffmpeg.exe"
TMPDIR = tempfile.gettempdir()

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

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

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

def fetch_html(domain: str):
    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: str) -> str:
    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: str):
    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: str):
    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: str):
    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>"

    return html, r["prompt"], r["domain"], r["text_snippet"][:800], r["images"]

# ---------------- ETAP 2 — LLM ----------------

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 ----------------

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 ----------------

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 None

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

    if not images_b64:
        return None

    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