import os
import re
import time
import json
import logging
from groq import Groq
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

api_key    = os.getenv("GROQ_API_KEY")
MODEL_NAME = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
client     = Groq(api_key=api_key)

MIN_DURATION    = 60    # seconds
MAX_DURATION    = 180   # seconds
TARGET_DURATION = 90    # ideal segment length

# ── Backtick fence markers ──
_FENCE_JSON = "```json"
_FENCE      = "```"


# ─────────────────────────────────────────────────────────────────────────────
# VALIDATION — no mechanical padding, only natural segments pass
# ─────────────────────────────────────────────────────────────────────────────

def validate_segments(segments, video_duration=None):
    """
    Accept only segments the AI correctly identified as complete scenes.
    Short segments are REJECTED — not padded — to avoid random cuts.
    The AI is instructed to return full scenes, so if it returns something
    short it means it only found the punchline, not the full story.
    """
    valid = []
    for seg in segments:
        start = seg.get("start_time", 0)
        end   = seg.get("end_time",   0)
        dur   = end - start

        # Clamp end_time to video duration if known
        if video_duration and end > video_duration:
            end = video_duration
            seg["end_time"] = round(end, 2)
            dur = end - start

        # Too long → hard reject
        if dur > MAX_DURATION:
            logger.warning(
                f"⚠️  Too long ({dur:.1f}s) — skipped: "
                f"'{seg.get('title', '')}' [{start}s–{end}s]"
            )
            continue

        # Too short → reject and log so we can tune the prompt
        if dur < MIN_DURATION:
            logger.warning(
                f"⚠️  Too short ({dur:.1f}s) — AI returned only punchline, not full scene. "
                f"Skipped: '{seg.get('title', '')}' [{start}s–{end}s]"
            )
            continue

        seg["duration"] = round(dur, 2)
        valid.append(seg)

    logger.info(f"✅ Valid segments after filter: {len(valid)}/{len(segments)}")
    return valid


# ─────────────────────────────────────────────────────────────────────────────
# FALLBACK — heuristic when AI returns nothing
# ─────────────────────────────────────────────────────────────────────────────

def _fallback_segments_from_transcript(transcript: str, video_duration: float) -> list:
    """
    If AI returns nothing useful, generate segments from the transcript using
    the densest word windows around timestamp markers.
    Only used as a last resort.
    """
    logger.warning("⚠️  Using fallback segment generator from transcript timestamps")

    lines = []
    for match in re.finditer(r'\[(\d+\.?\d*)\s*-\s*(\d+\.?\d*)\]\s*(.*)', transcript):
        lines.append({
            "start": float(match.group(1)),
            "end":   float(match.group(2)),
            "text":  match.group(3).strip(),
        })

    if not lines:
        return []

    candidates = []
    step   = 60
    window = TARGET_DURATION
    t      = lines[0]["start"]
    max_t  = lines[-1]["end"]

    while t + MIN_DURATION <= max_t:
        w_end      = min(t + window, max_t)
        in_window  = [l for l in lines if l["start"] >= t and l["end"] <= w_end]
        word_count = sum(len(l["text"].split()) for l in in_window)
        candidates.append({
            "start_time":  round(t, 2),
            "end_time":    round(w_end, 2),
            "word_count":  word_count,
            "title":       f"Highlight at {int(t // 60)}m{int(t % 60):02d}s",
            "description": "Auto-detected highlight segment",
            "reason":      "Fallback: highest word-density window",
            "viral_score": word_count,
        })
        t += step

    if not candidates:
        return []

    candidates.sort(key=lambda x: x["word_count"], reverse=True)
    top = candidates[:3]
    top.sort(key=lambda x: x["start_time"])

    deduped = []
    for c in top:
        if deduped and c["start_time"] < deduped[-1]["end_time"] - 20:
            continue
        deduped.append(c)

    logger.info(f"🔧 Fallback generated {len(deduped)} segment(s)")
    return deduped


# ─────────────────────────────────────────────────────────────────────────────
# CORE ANALYSIS — single chunk
# ─────────────────────────────────────────────────────────────────────────────

def analyze_transcript(transcript, video_duration=None):
    """
    Send one transcript chunk to Groq and return validated segments.
    """

    prompt = f"""
You are a viral short-form video editor specializing in TikTok, Reels, and YouTube Shorts.
Your job is to find COMPLETE, PUBLISH-READY segments — not just funny lines or punchlines.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
⛔ CRITICAL DURATION RULE — VIOLATIONS WILL BE DELETED
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
- end_time − start_time MUST be between {MIN_DURATION} and {MAX_DURATION} seconds
- Segments shorter than {MIN_DURATION}s will be DELETED, not fixed
- "I found a funny line at 150s" is NOT a valid segment
- "I found a complete story from 95s to 220s that includes setup + peak + conclusion" IS valid

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
✅ MANDATORY THINKING PROCESS
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
For every candidate moment, follow these steps:

Step 1 — Find the INTERESTING MOMENT (the peak/punchline/reveal)
Step 2 — Scroll BACK in the transcript to find where the SETUP begins
          (usually 30–90 seconds before the peak)
          Ask: "When does the viewer first need to start watching to understand this?"
Step 3 — Scroll FORWARD to find where the NATURAL CONCLUSION ends
          (usually 15–40 seconds after the peak)
          Ask: "Is the story/joke/argument fully resolved? Is there a reaction?"
Step 4 — Check: end_time − start_time is between {MIN_DURATION} and {MAX_DURATION} seconds

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
📌 CORRECT vs WRONG EXAMPLE
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Funny moment at 150s:
  ✅ CORRECT: start=95s (setup), end=220s (reaction ends) → 125s duration
  ❌ WRONG:   start=145s (just before punchline), end=165s → 20s duration

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
📋 A COMPLETE SEGMENT MUST HAVE ALL OF THESE
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
1. HOOK        — something in the first 5s that makes viewers keep watching
2. SETUP       — enough context so a stranger understands without watching anything before
3. BUILD-UP    — tension, humor, conflict, or emotion growing
4. PAYOFF      — the peak moment (punchline / reveal / climax)
5. CONCLUSION  — resolution, reaction, or natural pause — NOT an abrupt cut
6. STANDALONE  — makes complete sense on its own

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
🔢 OUTPUT FORMAT — raw JSON only, no markdown, no explanation
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{{
  "segments": [
    {{
      "start_time":  <float — where the SETUP begins, NOT the funny moment>,
      "end_time":    <float — where the CONCLUSION ends, NOT just the punchline>,
      "title":       "<punchy YouTube Shorts title, max 60 chars>",
      "description": "<1-2 sentences describing the full arc: setup → peak → conclusion>",
      "reason":      "<why this is complete and viral-worthy, max 25 words>"
    }}
  ]
}}

If no complete story qualifies: {{"segments": []}}

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
TRANSCRIPT:
{transcript}
"""

    max_retries = 3
    base_delay  = 5

    for attempt in range(max_retries):
        try:
            response = client.chat.completions.create(
                model    = MODEL_NAME,
                messages = [
                    {
                        "role":    "system",
                        "content": (
                            "You are a JSON-only assistant. "
                            "Output raw JSON only — no markdown, no code blocks, no explanation. "
                            f"Every segment MUST be between {MIN_DURATION} and {MAX_DURATION} seconds. "
                            "Short segments (punchlines without setup) will be DELETED — not fixed. "
                            "Always include: full setup before the peak + reaction/conclusion after. "
                            "Use the timestamps in the transcript to pick natural scene boundaries."
                        ),
                    },
                    {"role": "user", "content": prompt},
                ],
                temperature = 0.3,
            )

            content = response.choices[0].message.content.strip()
            logger.info(f"🤖 AI Raw Response (first 300 chars): {content[:300]}...")

            # Strip markdown fences if model ignored system prompt
            if _FENCE_JSON in content:
                content = content.split(_FENCE_JSON)[1].split(_FENCE)[0].strip()
            elif _FENCE in content:
                content = content.split(_FENCE)[1].split(_FENCE)[0].strip()

            data         = json.loads(content)
            raw_segments = data.get("segments", [])

            valid_segments = validate_segments(raw_segments, video_duration=video_duration)

            # Fallback if AI returned nothing valid
            if not valid_segments:
                logger.warning("⚠️  AI returned 0 valid segments — trying fallback generator")
                fallback = _fallback_segments_from_transcript(
                    transcript, video_duration or 0
                )
                valid_segments = validate_segments(fallback, video_duration=video_duration)

            data["segments"] = valid_segments
            content          = json.dumps(data)
            logger.info(
                f"🤖 Parsed: {len(raw_segments)} raw → {len(valid_segments)} valid segments"
            )
            return {"content": content}

        except Exception as e:
            logger.error(f"❌ Error in Groq analysis (attempt {attempt + 1}): {e}")
            if attempt < max_retries - 1:
                wait = base_delay * (2 ** attempt)
                logger.warning(f"⚠️  Retrying in {wait}s...")
                time.sleep(wait)

    logger.error("❌ All retry attempts failed.")
    return {"content": '{"segments": []}'}


# ─────────────────────────────────────────────────────────────────────────────
# SCENE-AWARE CHUNKING
# ─────────────────────────────────────────────────────────────────────────────

def _sentence_chunk(transcript: str, max_tokens: int) -> list[str]:
    """
    Fallback chunker when transcript has no timestamps.
    Splits at sentence boundaries with overlap.
    """
    sentences      = transcript.replace('\n', ' ').split('. ')
    chunks         = []
    current_chunk  = []
    current_length = 0

    for sentence in sentences:
        sentence_length = len(sentence.split())
        if current_length + sentence_length > max_tokens and current_chunk:
            chunks.append('. '.join(current_chunk) + '.')
            current_chunk  = current_chunk[-5:] + [sentence]
            current_length = sum(len(s.split()) for s in current_chunk)
        else:
            current_chunk.append(sentence)
            current_length += sentence_length

    if current_chunk:
        chunks.append('. '.join(current_chunk) + '.')
    return chunks


def smart_chunk_transcript(transcript: str, max_tokens: int = 4000) -> list[str]:
    """
    Split transcript into coherent chunks at SCENE BOUNDARIES.

    Strategy:
    - Parse all [start - end] timestamp lines
    - Cut only when the chunk is full AND there is a natural gap (>3s) between
      consecutive lines — a sign of a scene or topic change
    - Keep 10-line overlap between chunks so cross-boundary stories aren't lost
    - Falls back to sentence-level splitting if no timestamps found
    """
    lines = []
    for match in re.finditer(r'\[(\d+\.?\d*)\s*-\s*(\d+\.?\d*)\]\s*(.*)', transcript):
        lines.append({
            "start": float(match.group(1)),
            "end":   float(match.group(2)),
            "text":  match.group(3).strip(),
        })

    if not lines:
        logger.info("📦 No timestamps found — using sentence chunker")
        return _sentence_chunk(transcript, max_tokens)

    chunks        = []
    current_lines = []
    current_words = 0

    for i, line in enumerate(lines):
        current_lines.append(line)
        current_words += len(line["text"].split())

        is_last = (i + 1 == len(lines))
        is_full = (current_words >= max_tokens)

        # Detect scene break: gap > 3 seconds between this line and the next
        next_gap = (
            lines[i + 1]["start"] - line["end"]
            if not is_last else 999
        )
        at_scene_break = next_gap > 3.0

        if is_last or (is_full and at_scene_break):
            chunk_text = "\n".join(
                f"[{l['start']} - {l['end']}] {l['text']}"
                for l in current_lines
            )
            chunks.append(chunk_text)
            logger.info(
                f"📦 Chunk {len(chunks)}: "
                f"{current_lines[0]['start']:.0f}s – {current_lines[-1]['end']:.0f}s "
                f"({current_words} words)"
            )
            # Overlap: keep last 10 lines for context continuity
            current_lines = current_lines[-10:]
            current_words = sum(len(l["text"].split()) for l in current_lines)

    logger.info(f"📦 Split into {len(chunks)} scene-aware chunk(s)")
    return chunks


# ─────────────────────────────────────────────────────────────────────────────
# MAIN ENTRY POINT
# ─────────────────────────────────────────────────────────────────────────────

def analyze_transcript_with_chunking(transcript, video_duration=None):
    """
    Analyze transcript using scene-aware chunking for long content.
    Processes each chunk separately and merges + deduplicates results.
    """
    word_count = len(transcript.split())

    if word_count > 3000:
        logger.info(f"📦 Transcript is {word_count} words — using scene-aware chunking...")
        chunks       = smart_chunk_transcript(transcript, max_tokens=3000)
        all_segments = []

        for i, chunk in enumerate(chunks):
            logger.info(f"🔄 Processing chunk {i + 1}/{len(chunks)}...")
            result = analyze_transcript(chunk, video_duration=video_duration)

            try:
                data = json.loads(result["content"])
                if "segments" in data:
                    all_segments.extend(data["segments"])
                    logger.info(
                        f"✅ Chunk {i + 1}: "
                        f"{len(data['segments'])} valid segment(s)"
                    )
            except Exception as e:
                logger.warning(f"⚠️  Failed to parse chunk {i + 1}: {e}")
                continue

        if all_segments:
            # Deduplicate: two segments are the same if their start times
            # round to within 10 seconds of each other
            unique_segments = []
            seen_times      = set()

            for seg in sorted(all_segments, key=lambda s: s.get("start_time", 0)):
                time_key = round(seg.get("start_time", 0) / 10) * 10
                if time_key not in seen_times:
                    unique_segments.append(seg)
                    seen_times.add(time_key)

            logger.info(f"📊 Total unique viral segments found: {len(unique_segments)}")
            return {"content": json.dumps({"segments": unique_segments[:10]})}

        logger.warning("⚠️  No valid segments found across all chunks.")
        return {"content": '{"segments": []}'}

    # Short transcript — process directly without chunking
    return analyze_transcript(transcript, video_duration=video_duration)


# ─────────────────────────────────────────────────────────────────────────────
# TESTING
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    test_transcript = """
    [0.0 - 5.0] Welcome to today's video about productivity hacks that actually work.
    [5.0 - 15.0] The first hack is something I call the 2-minute rule. If something takes less than 2 minutes, do it immediately.
    [15.0 - 30.0] This simple rule has transformed my life. I used to procrastinate on small tasks, but now I handle them right away.
    [30.0 - 45.0] The second hack is batching similar tasks together. Instead of checking email 20 times a day, I check it twice.
    [45.0 - 60.0] This has saved me hours every week. I batch my emails, phone calls, and even errands.
    [60.0 - 90.0] The third hack is the Pomodoro Technique. Work for 25 minutes, then take a 5-minute break.
    [90.0 - 120.0] This technique helps me stay focused and avoid burnout. I get more done in less time.
    [120.0 - 150.0] The fourth hack is to eliminate distractions completely. Turn off notifications, close tabs, and focus.
    [150.0 - 180.0] When you eliminate distractions, your productivity skyrockets. I finish in 2 hours what used to take 6.
    [180.0 - 210.0] The fifth and final hack is to review your day every evening. Spend 5 minutes planning tomorrow.
    [210.0 - 240.0] This evening review changed everything for me. I wake up knowing exactly what to do and I never waste morning time figuring out priorities.
    """

    logger.info("🧪 Testing AI Analysis...")
    result = analyze_transcript_with_chunking(test_transcript, video_duration=240.0)

    try:
        data     = json.loads(result["content"])
        segments = data.get("segments", [])
        logger.info(f"✅ Found {len(segments)} publish-ready segment(s):\n")

        for i, seg in enumerate(segments):
            duration = seg["end_time"] - seg["start_time"]
            logger.info(
                f"#{i + 1} [{seg['start_time']:.0f}s – {seg['end_time']:.0f}s] ({duration:.0f}s)\n"
                f"  📌 Title:       {seg.get('title', 'N/A')}\n"
                f"  📝 Description: {seg.get('description', 'N/A')}\n"
                f"  💡 Story Arc:   {seg.get('reason', 'N/A')}\n"
            )
    except Exception as e:
        logger.error(f"❌ Error parsing result: {e}")
        logger.info(f"Raw result: {result}")