import spaces
import gradio as gr
import logging
import os
import tempfile
import pandas as pd
import requests
from bs4 import BeautifulSoup
import torch
import whisper
import subprocess
from pydub import AudioSegment
import fitz
import docx
import yt_dlp
from functools import lru_cache
import gc
import time
from huggingface_hub import login
from unsloth import FastLanguageModel
from transformers import AutoTokenizer

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

# Login to Hugging Face Hub if token is available
HUGGINGFACE_TOKEN = os.environ.get('HUGGINGFACE_TOKEN')
if HUGGINGFACE_TOKEN:
    login(token=HUGGINGFACE_TOKEN)

class ModelManager:
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(ModelManager, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if not self._initialized:
            self.tokenizer = None
            self.model = None
            self.pipeline = None
            self.whisper_model = None
            self._initialized = True
            self.last_used = time.time()
    
    @spaces.GPU()
    def initialize_llm(self):
        """Initialize LLM model with Unsloth optimization"""
        try:
            MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
            
            logger.info("Loading Unsloth-optimized model...")
            self.model, self.tokenizer = FastLanguageModel.from_pretrained(
                model_name = MODEL_NAME,
                max_seq_length = 2048,
                dtype = torch.float16,
                load_in_4bit = True,
                token = HUGGINGFACE_TOKEN,
            )
            
            # Enable LoRA for better ZeroGPU performance
            self.model = FastLanguageModel.get_peft_model(
                self.model,
                r = 16,
                target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                                 "gate_proj", "up_proj", "down_proj"],
                lora_alpha = 16,
                lora_dropout = 0,
                bias = "none",
                use_gradient_checkpointing = True,
                random_state = 3407,
                max_seq_length = 2048,
            )
            
            logger.info("LLM initialized successfully with Unsloth")
            self.last_used = time.time()
            return True
            
        except Exception as e:
            logger.error(f"Error initializing LLM: {str(e)}")
            raise

    @spaces.GPU()
    def initialize_whisper(self):
        """Initialize Whisper model with safety fix"""
        try:
            logger.info("Loading Whisper model...")
            # Load with weights_only=True for security
            self.whisper_model = whisper.load_model(
                "tiny",
                device="cuda" if torch.cuda.is_available() else "cpu",
                download_root="/tmp/whisper",
                weights_only=True  # Security fix
            )
            logger.info("Whisper model initialized successfully")
            self.last_used = time.time()
            return True
        except Exception as e:
            logger.error(f"Error initializing Whisper: {str(e)}")
            raise

    def check_llm_initialized(self):
        """Check if LLM is initialized and initialize if needed"""
        if self.tokenizer is None or self.model is None:
            logger.info("LLM not initialized, initializing...")
            self.initialize_llm()
        self.last_used = time.time()
    
    def check_whisper_initialized(self):
        """Check if Whisper model is initialized and initialize if needed"""
        if self.whisper_model is None:
            logger.info("Whisper model not initialized, initializing...")
            self.initialize_whisper()
        self.last_used = time.time()
    
    def reset_models(self, force=False):
        """Reset models to free memory if they haven't been used recently"""
        current_time = time.time()
        if force or (current_time - self.last_used > 600):  
            try:
                logger.info("Resetting models to free memory...")
                
                if self.model is not None:
                    del self.model
                    
                if self.tokenizer is not None:
                    del self.tokenizer
                    
                if self.whisper_model is not None:
                    del self.whisper_model
                
                self.tokenizer = None
                self.model = None
                self.whisper_model = None
                
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                    torch.cuda.synchronize()
                
                gc.collect()
                logger.info("Models reset successfully")
                
            except Exception as e:
                logger.error(f"Error resetting models: {str(e)}")

model_manager = ModelManager()

@lru_cache(maxsize=32)
def download_social_media_video(url):
    """Download a video from social media."""
    ydl_opts = {
        'format': 'bestaudio/best',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'outtmpl': '%(id)s.%(ext)s',
    }
    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info_dict = ydl.extract_info(url, download=True)
            audio_file = f"{info_dict['id']}.mp3"
        logger.info(f"Video downloaded successfully: {audio_file}")
        return audio_file
    except Exception as e:
        logger.error(f"Error downloading video: {str(e)}")
        raise

def convert_video_to_audio(video_file):
    """Convert a video file to audio using ffmpeg directly."""
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_file:
            output_file = temp_file.name
        
        command = [
            "ffmpeg", 
            "-i", video_file, 
            "-q:a", "0",
            "-map", "a",
            "-vn",
            output_file,
            "-y"
        ]
        
        subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        
        logger.info(f"Video converted to audio: {output_file}")
        return output_file
    except Exception as e:
        logger.error(f"Error converting video: {str(e)}")
        raise

def preprocess_audio(audio_file):
    """Preprocess the audio file to improve quality."""
    try:
        audio = AudioSegment.from_file(audio_file)
        audio = audio.apply_gain(-audio.dBFS + (-20))
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_file:
            audio.export(temp_file.name, format="mp3")
            logger.info(f"Audio preprocessed: {temp_file.name}")
            return temp_file.name
    except Exception as e:
        logger.error(f"Error preprocessing audio: {str(e)}")
        raise

@spaces.GPU()
def transcribe_audio(file):
    """Transcribe an audio or video file."""
    try:
        model_manager.check_whisper_initialized()
        
        if isinstance(file, str) and file.startswith('http'):
            file_path = download_social_media_video(file)
        elif isinstance(file, str) and file.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
            file_path = convert_video_to_audio(file)
        elif file is not None:
            file_path = preprocess_audio(file.name)
        else:
            return ""

        logger.info(f"Transcribing audio: {file_path}")
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Audio file not found: {file_path}")

        with torch.inference_mode():
            result = model_manager.whisper_model.transcribe(file_path)
                
        transcription = result.get("text", "Error in transcription")
        logger.info(f"Transcription completed: {transcription[:50]}...")
        
        try:
            if os.path.exists(file_path):
                os.remove(file_path)
        except Exception as e:
            logger.warning(f"Could not remove temp file {file_path}: {str(e)}")
            
        return transcription
    except Exception as e:
        logger.error(f"Error transcribing: {str(e)}")
        return f"Error processing the file: {str(e)}"

@lru_cache(maxsize=32)
def read_document(document_path):
    """Read the content of a document."""
    try:
        if document_path.endswith(".pdf"):
            doc = fitz.open(document_path)
            return "\n".join([page.get_text() for page in doc])
        elif document_path.endswith(".docx"):
            doc = docx.Document(document_path)
            return "\n".join([paragraph.text for paragraph in doc.paragraphs])
        elif document_path.endswith((".xlsx", ".xls")):
            return pd.read_excel(document_path).to_string()
        elif document_path.endswith(".csv"):
            return pd.read_csv(document_path).to_string()
        else:
            return "Unsupported file type. Please upload a PDF, DOCX, XLSX or CSV document."
    except Exception as e:
        logger.error(f"Error reading document: {str(e)}")
        return f"Error reading document: {str(e)}"

@lru_cache(maxsize=32)
def read_url(url):
    """Read the content of a URL."""
    if not url or url.strip() == "":
        return ""
        
    try:
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        response = requests.get(url, headers=headers, timeout=15)
        response.raise_for_status()
        soup = BeautifulSoup(response.content, 'html.parser')
        
        for element in soup(["script", "style", "meta", "noscript", "iframe", "header", "footer", "nav"]):
            element.extract()
            
        main_content = soup.find("main") or soup.find("article") or soup.find("div", class_=["content", "main", "article"])
        if main_content:
            text = main_content.get_text(separator='\n', strip=True)
        else:
            text = soup.get_text(separator='\n', strip=True)
            
        lines = [line.strip() for line in text.split('\n') if line.strip()]
        text = '\n'.join(lines)
        
        return text[:10000]
    except Exception as e:
        logger.error(f"Error reading URL: {str(e)}")
        return f"Error reading URL: {str(e)}"

def process_social_content(url):
    """Process social media content."""
    if not url or url.strip() == "":
        return None
        
    try:
        text_content = read_url(url)
        try:
            video_content = transcribe_audio(url)
        except Exception as e:
            logger.error(f"Error processing video content: {str(e)}")
            video_content = None

        return {
            "text": text_content,
            "video": video_content
        }
    except Exception as e:
        logger.error(f"Error processing social content: {str(e)}")
        return None

@spaces.GPU()
def generate_news(instructions, facts, size, tone, *args):
    """Generate a news article based on provided data"""
    try:
        if isinstance(size, float):
            size = int(size)
        elif not isinstance(size, int):
            size = 250
            
        model_manager.check_llm_initialized()
        
        knowledge_base = {
            "instructions": instructions or "",
            "facts": facts or "",
            "document_content": [],
            "audio_data": [],
            "url_content": [],
            "social_content": []
        }

        num_audios = 5 * 3
        num_social_urls = 3 * 3
        num_urls = 5

        args = list(args)
        
        while len(args) < (num_audios + num_social_urls + num_urls + 5):
            args.append("")
            
        audios = args[:num_audios]
        social_urls = args[num_audios:num_audios+num_social_urls]
        urls = args[num_audios+num_social_urls:num_audios+num_social_urls+num_urls]
        documents = args[num_audios+num_social_urls+num_urls:]

        logger.info("Processing URLs...")
        for url in urls:
            if url and isinstance(url, str) and url.strip():
                content = read_url(url)
                if content and not content.startswith("Error"):
                    knowledge_base["url_content"].append(content)

        logger.info("Processing documents...")
        for document in documents:
            if document and hasattr(document, 'name'):
                content = read_document(document.name)
                if content and not content.startswith("Error"):
                    knowledge_base["document_content"].append(content)

        logger.info("Processing audio/video files...")
        for i in range(0, len(audios), 3):
            if i+2 < len(audios):
                audio_file, name, position = audios[i:i+3]
                if audio_file and hasattr(audio_file, 'name'):
                    knowledge_base["audio_data"].append({
                        "audio": audio_file,
                        "name": name or "Unknown",
                        "position": position or "Not specified"
                    })

        logger.info("Processing social media content...")
        for i in range(0, len(social_urls), 3):
            if i+2 < len(social_urls):
                social_url, social_name, social_context = social_urls[i:i+3]
                if social_url and isinstance(social_url, str) and social_url.strip():
                    social_content = process_social_content(social_url)
                    if social_content:
                        knowledge_base["social_content"].append({
                            "url": social_url,
                            "name": social_name or "Unknown",
                            "context": social_context or "Not specified",
                            "text": social_content.get("text", ""),
                            "video": social_content.get("video", "")
                        })

        transcriptions_text = ""
        raw_transcriptions = ""

        logger.info("Transcribing audio...")
        for idx, data in enumerate(knowledge_base["audio_data"]):
            if data["audio"] is not None:
                transcription = transcribe_audio(data["audio"])
                if transcription and not transcription.startswith("Error"):
                    transcriptions_text += f'"{transcription}" - {data["name"]}, {data["position"]}\n\n'
                    raw_transcriptions += f'[Audio/Video {idx + 1}]: "{transcription}" - {data["name"]}, {data["position"]}\n\n'

        for idx, data in enumerate(knowledge_base["social_content"]):
            if data["text"] and not str(data["text"]).startswith("Error"):
                text_excerpt = data["text"][:500] + "..." if len(data["text"]) > 500 else data["text"]
                social_text = f'[Social media {idx+1} - text]: "{text_excerpt}" - {data["name"]}, {data["context"]}\n\n'
                transcriptions_text += social_text
                raw_transcriptions += social_text
                
            if data["video"] and not str(data["video"]).startswith("Error"):
                video_transcription = f'[Social media {idx+1} - video]: "{data["video"]}" - {data["name"]}, {data["context"]}\n\n'
                transcriptions_text += video_transcription
                raw_transcriptions += video_transcription

        document_summaries = []
        for idx, doc in enumerate(knowledge_base["document_content"]):
            if len(doc) > 1000:
                doc_excerpt = doc[:1000] + "... [document continues]"
            else:
                doc_excerpt = doc
            document_summaries.append(f"[Document {idx+1}]: {doc_excerpt}")
        
        document_content = "\n\n".join(document_summaries)
        
        url_summaries = []
        for idx, url_content in enumerate(knowledge_base["url_content"]):
            if len(url_content) > 1000:
                url_excerpt = url_content[:1000] + "... [content continues]"
            else:
                url_excerpt = url_content
            url_summaries.append(f"[URL {idx+1}]: {url_excerpt}")
            
        url_content = "\n\n".join(url_summaries)

        prompt = f"""<s>[INST] You are a professional news writer. Write a news article based on the following information:

Instructions: {knowledge_base["instructions"]}

Facts: {knowledge_base["facts"]}

Additional content from documents: 
{document_content}

Additional content from URLs: 
{url_content}

Use these transcriptions as direct and indirect quotes:
{transcriptions_text}

Follow these requirements:
- Write a title
- Write a 15-word hook that complements the title
- Write the body with approximately {size} words
- Use a {tone} tone
- Answer the 5 Ws (Who, What, When, Where, Why) in the first paragraph
- Use at least 80% direct quotes (in quotation marks)
- Use proper journalistic style
- Do not invent information
- Be rigorous with the provided facts [/INST]"""

        try:
            logger.info("Generating news article...")
            
            max_length = min(len(prompt.split()) + size * 2, 2048)
            
            inputs = model_manager.tokenizer(
                prompt,
                return_tensors = "pt",
                padding = True,
                truncation = True,
                max_length = 2048,
            ).to("cuda")

            outputs = model_manager.model.generate(
                **inputs,
                max_new_tokens = size + 100,
                temperature = 0.7,
                do_sample = True,
                pad_token_id = model_manager.tokenizer.eos_token_id,
            )
            
            generated_text = model_manager.tokenizer.decode(outputs[0], skip_special_tokens = True)
            
            if "[/INST]" in generated_text:
                news_article = generated_text.split("[/INST]")[1].strip()
            else:
                prompt_fragment = " ".join(prompt.split()[:50])
                if prompt_fragment in generated_text:
                    news_article = generated_text[generated_text.find(prompt_fragment) + len(prompt_fragment):].strip()
                else:
                    news_article = generated_text
            
            logger.info(f"News generation completed: {len(news_article)} chars")
            
        except Exception as gen_error:
            logger.error(f"Error in text generation: {str(gen_error)}")
            raise
        
        return news_article, raw_transcriptions

    except Exception as e:
        logger.error(f"Error generating news: {str(e)}")
        try:
            model_manager.reset_models(force=True)
        except Exception as reset_error:
            logger.error(f"Failed to reset models: {str(reset_error)}")
        return f"Error generating news: {str(e)}", "Error processing transcriptions."

def create_demo():
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 📰 NewsIA - AI News Generator")
        gr.Markdown("Create professional news articles from multiple sources.")
        
        with gr.Row():
            with gr.Column(scale=2):
                instructions = gr.Textbox(
                    label="News Instructions",
                    placeholder="Enter specific instructions for news generation",
                    lines=2
                )
                facts = gr.Textbox(
                    label="Key Facts",
                    placeholder="Describe the most important facts to include",
                    lines=4
                )
                
                with gr.Row():
                    size = gr.Slider(
                        label="Approximate Length (words)",
                        minimum=100,
                        maximum=500,
                        value=250,
                        step=50
                    )
                    tone = gr.Dropdown(
                        label="News Tone",
                        choices=["serious", "neutral", "funny", "formal", "informal", "urgent"],
                        value="neutral"
                    )

            with gr.Column(scale=3):
                inputs_list = []
                inputs_list.extend([instructions, facts, size, tone])

                with gr.Tabs():
                    with gr.TabItem("📝 Documents"):
                        documents = []
                        for i in range(1, 6):
                            doc = gr.File(
                                label=f"Document {i}",
                                file_types=["pdf", "docx", "xlsx", "csv"],
                                file_count="single"
                            )
                            documents.append(doc)
                            inputs_list.append(doc)

                    with gr.TabItem("🔊 Audio/Video"):
                        for i in range(1, 6):
                            with gr.Group():
                                gr.Markdown(f"**Source {i}**")
                                file = gr.File(
                                    label=f"Audio/Video {i}",
                                    file_types=["audio", "video"]
                                )
                                with gr.Row():
                                    name = gr.Textbox(
                                        label="Name",
                                        placeholder="Interviewee name"
                                    )
                                    position = gr.Textbox(
                                        label="Position/Role",
                                        placeholder="Position or role"
                                    )
                                inputs_list.extend([file, name, position])

                    with gr.TabItem("🌐 URLs"):
                        for i in range(1, 6):
                            url = gr.Textbox(
                                label=f"URL {i}",
                                placeholder="https://..."
                            )
                            inputs_list.append(url)

                    with gr.TabItem("📱 Social Media"):
                        for i in range(1, 4):
                            with gr.Group():
                                gr.Markdown(f"**Social Media {i}**")
                                social_url = gr.Textbox(
                                    label="URL",
                                    placeholder="https://..."
                                )
                                with gr.Row():
                                    social_name = gr.Textbox(
                                        label="Account/Name",
                                        placeholder="Account or person name"
                                    )
                                    social_context = gr.Textbox(
                                        label="Context",
                                        placeholder="Relevant context"
                                    )
                                inputs_list.extend([social_url, social_name, social_context])

        with gr.Row():
            generate_btn = gr.Button("✨ Generate News", variant="primary")
            reset_btn = gr.Button("🔄 Clear All")

        with gr.Tabs():
            with gr.TabItem("📄 Generated News"):
                news_output = gr.Textbox(
                    label="News Draft",
                    lines=15,
                    show_copy_button=True
                )
            
            with gr.TabItem("🎙️ Transcriptions"):
                transcriptions_output = gr.Textbox(
                    label="Source Transcriptions",
                    lines=10,
                    show_copy_button=True
                )

        generate_btn.click(
            fn=generate_news,
            inputs=inputs_list,
            outputs=[news_output, transcriptions_output]
        )
        
        def reset_all():
            return [None]*len(inputs_list) + ["", ""]
        
        reset_btn.click(
            fn=reset_all,
            inputs=None,
            outputs=inputs_list + [news_output, transcriptions_output]
        )

    return demo

if __name__ == "__main__":
    try:
        model_manager.initialize_whisper()
    except Exception as e:
        logger.warning(f"Initial whisper model loading failed: {str(e)}")
    
    demo = create_demo()
    demo.queue(max_size=5)
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860
    )