# utils.py

import os
import re
import json
import requests
import tempfile
from bs4 import BeautifulSoup
from typing import List, Literal
from pydantic import BaseModel
from pydub import AudioSegment, effects
from transformers import pipeline
import yt_dlp
import tiktoken
from groq import Groq
import numpy as np
import torch
import random

class DialogueItem(BaseModel):
    speaker: Literal["Jane", "John"]
    text: str

class Dialogue(BaseModel):
    dialogue: List[DialogueItem]

# Initialize Whisper ASR pipeline
asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-tiny.en", device=0 if torch.cuda.is_available() else -1)

def truncate_text(text, max_tokens=2048):
    print("[LOG] Truncating text if needed.")
    tokenizer = tiktoken.get_encoding("cl100k_base")
    tokens = tokenizer.encode(text)
    if len(tokens) > max_tokens:
        print("[LOG] Text too long, truncating.")
        return tokenizer.decode(tokens[:max_tokens])
    return text

def extract_text_from_url(url):
    print("[LOG] Extracting text from URL:", url)
    try:
        response = requests.get(url)
        if response.status_code != 200:
            print(f"[ERROR] Failed to fetch URL: {url} with status code {response.status_code}")
            return ""
        soup = BeautifulSoup(response.text, 'html.parser')
        for script in soup(["script", "style"]):
            script.decompose()
        text = soup.get_text(separator=' ')
        print("[LOG] Text extraction from URL successful.")
        return text
    except Exception as e:
        print(f"[ERROR] Exception during text extraction from URL: {e}")
        return ""

def pitch_shift(audio: AudioSegment, semitones: int) -> AudioSegment:
    """
    Shifts the pitch of an AudioSegment by a given number of semitones.
    Positive semitones shift the pitch up, negative shifts it down.
    """
    print(f"[LOG] Shifting pitch by {semitones} semitones.")
    new_sample_rate = int(audio.frame_rate * (2.0 ** (semitones / 12.0)))
    shifted_audio = audio._spawn(audio.raw_data, overrides={'frame_rate': new_sample_rate})
    return shifted_audio.set_frame_rate(audio.frame_rate)

def is_sufficient(text: str, min_word_count: int = 500) -> bool:
    """
    Determines if the fetched information meets the sufficiency criteria.
    """
    word_count = len(text.split())
    print(f"[DEBUG] Aggregated word count: {word_count}")
    return word_count >= min_word_count

def query_llm_for_additional_info(topic: str, existing_text: str) -> str:
    """
    Queries the Groq API to retrieve additional relevant information from the LLM's knowledge base.
    """
    print("[LOG] Querying LLM for additional information.")
    system_prompt = (
        "You are an AI assistant with extensive knowledge up to 2023-10. "
        "Provide additional relevant information on the following topic based on your knowledge base.\n\n"
        f"Topic: {topic}\n\n"
        f"Existing Information: {existing_text}\n\n"
        "Please add more insightful details, facts, and perspectives to enhance the understanding of the topic."
    )
    groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
    try:
        response = groq_client.chat.completions.create(
            messages=[{"role": "system", "content": system_prompt}],
            model="llama-3.3-70b-versatile",
            max_tokens=1024,
            temperature=0.7
        )
    except Exception as e:
        print("[ERROR] Groq API error during fallback:", e)
        return ""
    additional_info = response.choices[0].message.content.strip()
    print("[DEBUG] Additional information from LLM:")
    print(additional_info)
    return additional_info

def research_topic(topic: str) -> str:
    # News sources
    sources = {
        "BBC": "https://feeds.bbci.co.uk/news/rss.xml",
        "CNN": "http://rss.cnn.com/rss/edition.rss",
        "Associated Press": "https://apnews.com/apf-topnews",
        "NDTV": "https://www.ndtv.com/rss/top-stories",
        "Times of India": "https://timesofindia.indiatimes.com/rssfeeds/296589292.cms",
        "The Hindu": "https://www.thehindu.com/news/national/kerala/rssfeed.xml",
        "Economic Times": "https://economictimes.indiatimes.com/rssfeeds/1977021501.cms",
        "Google News - Custom": f"https://news.google.com/rss/search?q={requests.utils.quote(topic)}&hl=en-IN&gl=IN&ceid=IN:en",
    }

    summary_parts = []

    # Wikipedia summary
    wiki_summary = fetch_wikipedia_summary(topic)
    if wiki_summary:
        summary_parts.append(f"From Wikipedia: {wiki_summary}")

    for name, url in sources.items():
        try:
            items = fetch_rss_feed(url)
            if not items:
                continue
            title, desc, link = find_relevant_article(items, topic, min_match=2)
            if link:
                article_text = fetch_article_text(link)
                if article_text:
                    summary_parts.append(f"From {name}: {article_text}")
                else:
                    summary_parts.append(f"From {name}: {title} - {desc}")
        except Exception as e:
            print(f"[ERROR] Error fetching from {name} RSS feed:", e)
            continue

    aggregated_info = " ".join(summary_parts)
    print("[DEBUG] Aggregated info from primary sources:")
    print(aggregated_info)

    if not is_sufficient(aggregated_info):
        print("[LOG] Insufficient info from primary sources. Fallback to LLM.")
        additional_info = query_llm_for_additional_info(topic, aggregated_info)
        if additional_info:
            aggregated_info += " " + additional_info
        else:
            print("[ERROR] Failed to retrieve additional info from LLM.")

    if not aggregated_info:
        return f"Sorry, I couldn't find recent information on '{topic}'."

    return aggregated_info

def fetch_wikipedia_summary(topic: str) -> str:
    print("[LOG] Fetching Wikipedia summary for:", topic)
    try:
        search_url = f"https://en.wikipedia.org/w/api.php?action=opensearch&search={requests.utils.quote(topic)}&limit=1&namespace=0&format=json"
        resp = requests.get(search_url)
        if resp.status_code != 200:
            print(f"[ERROR] Failed to fetch Wikipedia search results for topic: {topic}")
            return ""
        data = resp.json()
        if len(data) > 1 and data[1]:
            title = data[1][0]
            summary_url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{requests.utils.quote(title)}"
            s_resp = requests.get(summary_url)
            if s_resp.status_code == 200:
                s_data = s_resp.json()
                if "extract" in s_data:
                    print("[LOG] Wikipedia summary fetched successfully.")
                    return s_data["extract"]
        return ""
    except Exception as e:
        print(f"[ERROR] Exception during Wikipedia summary fetch: {e}")
        return ""

def fetch_rss_feed(feed_url: str) -> list:
    print("[LOG] Fetching RSS feed:", feed_url)
    try:
        resp = requests.get(feed_url)
        if resp.status_code != 200:
            print(f"[ERROR] Failed to fetch RSS feed: {feed_url}")
            return []
        soup = BeautifulSoup(resp.content, "html.parser")
        items = soup.find_all("item")
        return items
    except Exception as e:
        print(f"[ERROR] Exception fetching RSS feed {feed_url}: {e}")
        return []

def find_relevant_article(items, topic: str, min_match=2) -> tuple:
    """
    Searches for relevant articles based on topic keywords.
    """
    print("[LOG] Finding relevant articles...")
    keywords = re.findall(r'\w+', topic.lower())
    for item in items:
        title = item.find("title").get_text().strip() if item.find("title") else ""
        description = item.find("description").get_text().strip() if item.find("description") else ""
        text = f"{title.lower()} {description.lower()}"
        matches = sum(1 for kw in keywords if kw in text)
        if matches >= min_match:
            link = item.find("link").get_text().strip() if item.find("link") else ""
            print(f"[LOG] Relevant article found: {title}")
            return title, description, link
    return None, None, None

def fetch_article_text(link: str) -> str:
    print("[LOG] Fetching article text from:", link)
    if not link:
        return ""
    try:
        resp = requests.get(link)
        if resp.status_code != 200:
            print(f"[ERROR] Failed to fetch article with status {resp.status_code}")
            return ""
        soup = BeautifulSoup(resp.text, 'html.parser')
        paragraphs = soup.find_all("p")
        text = " ".join(p.get_text() for p in paragraphs[:5])
        return text.strip()
    except Exception as e:
        print(f"[ERROR] Error fetching article text: {e}")
        return ""

def generate_script(system_prompt: str, input_text: str, tone: str, target_length: str):
    print("[LOG] Generating script with tone:", tone, "and length:", target_length)
    groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

    length_mapping = {
        "1-3 Mins": (200, 450),
        "3-5 Mins": (450, 750),
        "5-10 Mins": (750, 1500),
        "10-20 Mins": (1500, 3000)
    }
    min_words, max_words = length_mapping.get(target_length, (200, 450))

    tone_description = {
        "Humorous": "funny and exciting, makes people chuckle",
        "Formal": "business-like, well-structured, professional",
        "Casual": "like a conversation between close friends, relaxed and informal",
        "Youthful": "like how teenagers might chat, energetic and lively"
    }
    chosen_tone = tone_description.get(tone, "casual")

    prompt = (
        f"{system_prompt}\n"
        f"TONE: {chosen_tone}\n"
        f"TARGET LENGTH: {target_length} ({min_words}-{max_words} words)\n"
        f"INPUT TEXT: {input_text}\n\n"
        "Please provide the output in the following JSON format without any additional text:\n\n"
        "{\n"
        '    "dialogue": [\n'
        '        {\n'
        '            "speaker": "Jane",\n'
        '            "text": "..." \n'
        '        },\n'
        '        {\n'
        '            "speaker": "John",\n'
        '            "text": "..." \n'
        '        }\n'
        "    ]\n"
        "}"
    )

    try:
        response = groq_client.chat.completions.create(
            messages=[{"role": "system", "content": prompt}],
            model="llama-3.3-70b-versatile",
            max_tokens=2048,
            temperature=0.7
        )
    except Exception as e:
        raise ValueError(f"Error communicating with Groq API: {str(e)}")

    raw_content = response.choices[0].message.content.strip()
    start_index = raw_content.find('{')
    end_index = raw_content.rfind('}')
    if start_index == -1 or end_index == -1:
        raise ValueError("Failed to parse dialogue: No JSON found.")

    json_str = raw_content[start_index:end_index+1].strip()
    data = json.loads(json_str)
    return Dialogue(**data)

# --------------------------------------------------------------
# TTS Preprocessing to handle decimals, hyphens, short thinking pauses, etc.
# --------------------------------------------------------------
def _preprocess_text_for_tts(text: str) -> str:
    """
    1) Convert decimals to spelled-out words ("3.14" -> "three point one four").
    2) Replace hyphens with spaces (so TTS doesn't say 'dash').
    3) Insert filler words or '...' for natural-sounding pauses at significant points.
    """

    # 1) Convert decimals
    def convert_decimal(m):
        number_str = m.group()  # e.g. "3.14"
        parts = number_str.split('.')
        whole_part = _spell_digits(parts[0])  # "three"
        decimal_part = " ".join(_spell_digits(d) for d in parts[1])
        return f"{whole_part} point {decimal_part}"

    text = re.sub(r"\d+\.\d+", convert_decimal, text)

    # 2) Replace hyphens with spaces
    #    e.g. "mother-in-law" -> "mother in law"
    text = re.sub(r"-", " ", text)

    # 3) Insert natural-sounding short pauses:
    #    a) After exclamation points or question marks, add "..." with small chance
    #    b) Random small "thinking" filler for major statements

    # Step 3a: Exclamations / questions
    text = re.sub(r"(!+)", r"\1...", text)  # e.g. "Wow!" -> "Wow!..."
    text = re.sub(r"(\?+)", r"\1...", text) # e.g. "Really?" -> "Really?..."

    # Step 3b: Insert small breaks for "thinking"
    # We'll define some keywords that might indicate a "significant point."
    # e.g. "important", "significant", "crucial", "point", "topic" 
    # Then we insert '..., hmm,' or '..., well,' afterwards with a small chance.
    def insert_thinking_pause(m):
        word = m.group(1)
        if random.random() < 0.5:
            return f"{word}..., hmm,"
        else:
            return f"{word}..., well,"

    keywords_pattern = r"\b(important|significant|crucial|point|topic)\b"
    text = re.sub(keywords_pattern, insert_thinking_pause, text, flags=re.IGNORECASE)

    return text.strip()

def _spell_digits(d: str) -> str:
    """
    Convert each digit '3' -> 'three', '5' -> 'five', etc.
    """
    digit_map = {
        '0': 'zero', '1': 'one', '2': 'two', '3': 'three',
        '4': 'four','5': 'five','6': 'six','7': 'seven',
        '8': 'eight','9': 'nine'
    }
    return " ".join(digit_map[ch] for ch in d if ch in digit_map)

def generate_audio_mp3(text: str, speaker: str) -> str:
    """
    Main TTS function, calls Deepgram with preprocessed text.
    Returns path to a temporary MP3 file.
    """
    try:
        print(f"[LOG] Generating audio for speaker: {speaker}")

        # Preprocess text (decimal/hyphen/pause insertion)
        processed_text = _preprocess_text_for_tts(text)

        # Define Deepgram API endpoint
        deepgram_api_url = "https://api.deepgram.com/v1/speak"
        params = {
            "model": "aura-asteria-en",  # default female
        }
        if speaker == "John":
            params["model"] = "aura-perseus-en"

        headers = {
            "Accept": "audio/mpeg",
            "Content-Type": "application/json",
            "Authorization": f"Token {os.environ.get('DEEPGRAM_API_KEY')}"
        }
        body = {
            "text": processed_text
        }

        print("[LOG] Sending TTS request to Deepgram...")
        response = requests.post(deepgram_api_url, params=params, headers=headers, json=body, stream=True)
        if response.status_code != 200:
            raise ValueError(f"Deepgram TTS error: {response.status_code}, {response.text}")

        content_type = response.headers.get('Content-Type', '')
        if 'audio/mpeg' not in content_type:
            raise ValueError("Unexpected Content-Type from Deepgram.")

        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as mp3_file:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    mp3_file.write(chunk)
            mp3_path = mp3_file.name

        # Normalize volume
        audio_seg = AudioSegment.from_file(mp3_path, format="mp3")
        audio_seg = effects.normalize(audio_seg)

        final_mp3_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3").name
        audio_seg.export(final_mp3_path, format="mp3")

        if os.path.exists(mp3_path):
            os.remove(mp3_path)

        return final_mp3_path
    except Exception as e:
        print("[ERROR] Error generating audio:", e)
        raise ValueError(f"Error generating audio: {str(e)}")

def transcribe_youtube_video(video_url: str) -> str:
    """
    Downloads and transcribes the audio from a YouTube video using Whisper (pipeline).
    """
    print("[LOG] Transcribing YouTube video:", video_url)
    fd, audio_file = tempfile.mkstemp(suffix=".wav")
    os.close(fd)

    ydl_opts = {
        'format': 'bestaudio/best',
        'outtmpl': audio_file,
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'wav',
            'preferredquality': '192'
        }],
        'quiet': True,
        'no_warnings': True,
    }

    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([video_url])
    except yt_dlp.utils.DownloadError as e:
        raise ValueError(f"Error downloading YouTube video: {str(e)}")

    print("[LOG] Audio downloaded at:", audio_file)
    try:
        result = asr_pipeline(audio_file)
        transcript = result["text"]
        print("[LOG] Transcription completed.")
        return transcript.strip()
    except Exception as e:
        raise ValueError(f"Error transcribing YouTube video: {str(e)}")
    finally:
        if os.path.exists(audio_file):
            os.remove(audio_file)
            print(f"[LOG] Removed temporary audio file: {audio_file}")