import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import requests
from bs4 import BeautifulSoup
from gtts import gTTS
from moviepy.editor import *
import os
from datetime import datetime
import base64
import logging
import threading
import queue

# ====================== CẤU HÌNH ======================
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Config:
    MODEL_NAME = "vinai/PhoGPT-4B-Chat"  # Model nhẹ cho tiếng Việt
    NEWS_URLS = {
        "all": "https://www.espn.com/soccer/",
        "transfer": "https://www.espn.com/soccer/transfers",
        "match": "https://www.espn.com/soccer/schedule"
    }

# ====================== XỬ LÝ MODEL ======================
class ModelLoader:
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.ready = False
        self.error = None
        self.queue = queue.Queue()

    def load_model(self):
        """Tải model trong thread chính"""
        try:
            logger.info("Đang tải model...")
            self.tokenizer = AutoTokenizer.from_pretrained(Config.MODEL_NAME)
            self.model = AutoModelForCausalLM.from_pretrained(Config.MODEL_NAME)
            self.ready = True
            logger.info("Tải model thành công!")
        except Exception as e:
            self.error = str(e)
            logger.error(f"Lỗi tải model: {e}")
        finally:
            self.queue.put(True)

# ====================== GIAO DIỆN ======================
def setup_ui():
    st.set_page_config(
        page_title="AI Bình Luận Bóng Đá",
        page_icon="⚽",
        layout="wide"
    )
    st.title("⚽ AI Bình Luận Bóng Đá")

# ====================== XỬ LÝ CHÍNH ======================
def main():
    setup_ui()
    
    # Khởi tạo trình tải model
    if 'model_loader' not in st.session_state:
        st.session_state.model_loader = ModelLoader()
        threading.Thread(target=st.session_state.model_loader.load_model).start()

    # Chờ model tải xong
    if not st.session_state.model_loader.ready:
        with st.spinner("Đang tải model AI, vui lòng chờ..."):
            st.session_state.model_loader.queue.get()
        
        if st.session_state.model_loader.error:
            st.error(f"Lỗi: {st.session_state.model_loader.error}")
            return

    # Phần còn lại của ứng dụng
    st.success("Model đã sẵn sàng!")

# ====================== TIỆN ÍCH HỖ TRỢ ======================
def setup_page():
    """Cấu hình trang Streamlit"""
    st.set_page_config(
        page_title="🎤 AI Bình Luận Bóng Đá",
        page_icon="⚽",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    st.markdown("""
    <style>
        .main { background-color: #f8f9fa; }
        .stButton>button {
            background-color: #4CAF50;
            color: white;
            border-radius: 5px;
            padding: 10px 24px;
            margin: 5px 0;
        }
        .news-card {
            background: white;
            border-radius: 10px;
            padding: 15px;
            margin-bottom: 15px;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
        .footer { 
            text-align: center; 
            margin-top: 2rem; 
            color: #666;
            font-size: 0.9em;
        }
        .stProgress > div > div > div {
            background-color: #4CAF50;
        }
    </style>
    """, unsafe_allow_html=True)

# ====================== XỬ LÝ TIN TỨC ======================
class NewsProcessor:
    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
    def fetch_news(self, category="all", limit=5):
        """Lấy tin tức với cơ chế retry tự động"""
        try:
            url = AppConfig.NEWS_SOURCES.get(category, AppConfig.NEWS_SOURCES["all"])
            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"
            }
            
            logger.info(f"Đang lấy tin tức từ {url}")
            response = requests.get(url, headers=headers, timeout=15)
            response.raise_for_status()
            soup = BeautifulSoup(response.text, 'html.parser')
            
            return self._parse_news(soup, url, limit)
        
        except requests.RequestException as e:
            logger.error(f"Lỗi kết nối khi lấy tin tức: {e}")
            return None
        except Exception as e:
            logger.error(f"Lỗi không xác định khi lấy tin tức: {e}")
            return None

    def _parse_news(self, soup, url, limit):
        """Phân tích cú pháp HTML để trích xuất tin tức"""
        news_items = []
        
        try:
            if "espn.com" in url:
                articles = soup.select(".contentItem__contentWrapper")[:limit]
                for article in articles:
                    try:
                        title = article.select_one(".contentItem__title")
                        link = article.find("a", href=True)
                        summary = article.select_one(".contentItem__subhead")
                        
                        if not title or not link:
                            continue
                            
                        news_items.append({
                            "title": title.get_text(strip=True),
                            "summary": summary.get_text(strip=True) if summary else "",
                            "link": "https://www.espn.com" + link["href"]
                        })
                    except Exception as e:
                        logger.warning(f"Bỏ qua bài viết ESPN do lỗi: {e}")
            else:  # VnExpress
                articles = soup.select(".item-news")[:limit]
                for article in articles:
                    try:
                        title = article.select_one("h3.title-news a")
                        summary = article.select_one(".description a")
                        
                        if not title:
                            continue
                            
                        news_items.append({
                            "title": title.get_text(strip=True),
                            "summary": summary.get_text(strip=True) if summary else "",
                            "link": title["href"]
                        })
                    except Exception as e:
                        logger.warning(f"Bỏ qua bài viết VnExpress do lỗi: {e}")
        except Exception as e:
            logger.error(f"Lỗi phân tích tin tức: {e}")
            return None
        
        return news_items if news_items else None

# ====================== XỬ LÝ AI ======================
class AIModelLoader:
    def __init__(self):
        self.model_loaded = False
        self.model = None
        self.tokenizer = None
        self.error = None
        self.loading_queue = queue.Queue()

    def load_in_main_thread(self):
        """Tải model trong thread chính"""
        try:
            logger.info("Đang tải model AI...")
            self.tokenizer = AutoTokenizer.from_pretrained(AppConfig.MODEL_NAME)
            self.model = AutoModelForCausalLM.from_pretrained(AppConfig.MODEL_NAME)
            self.model_loaded = True
            logger.info("Tải model thành công!")
        except Exception as e:
            self.error = str(e)
            logger.error(f"Lỗi tải model: {e}")
        finally:
            self.loading_queue.put(True)

    def get_model_pipeline(self):
        """Tạo pipeline sau khi model đã tải"""
        if not self.model_loaded:
            return None
            
        return pipeline(
            "text-generation",
            model=self.model,
            tokenizer=self.tokenizer,
            device="cpu"
        )

class AICommentator:
    def __init__(self, model_pipeline):
        self.model = model_pipeline
        
    def generate_commentary(self, news_item, style="Triết học"):
        """Tạo bình luận sử dụng AI"""
        try:
            prompt = f"""
            Bạn là chuyên gia bình luận bóng đá. Hãy phân tích tin sau theo phong cách {style}:
            
            **Tiêu đề:** {news_item['title']}
            **Nội dung:** {news_item['summary']}
            
            Yêu cầu:
            - Dài 100-150 từ
            - Kết hợp phân tích kỹ thuật
            - Rút ra bài học cuộc sống
            - Giọng văn {style}
            - Viết bằng tiếng Việt tự nhiên
            """
            
            result = self.model(
                prompt,
                max_length=400,
                num_return_sequences=1,
                temperature=0.7,
                do_sample=True,
                top_k=50,
                top_p=0.95,
                early_stopping=True
            )
            
            return result[0]['generated_text']
        except Exception as e:
            logger.error(f"Lỗi tạo bình luận: {e}")
            return None

# ====================== XỬ LÝ VIDEO ======================
class VideoProducer:
    def __init__(self):
        self.output_dir = "output_videos"
        os.makedirs(self.output_dir, exist_ok=True)
    
    def create_scene(self, text, scene_num, bg_color=None):
        """Tạo một phân cảnh video"""
        try:
            if bg_color is None:
                bg_color = random.choice(AppConfig.BACKGROUND_COLORS)
            
            # Tạo audio từ text
            audio_path = f"{self.output_dir}/scene_{scene_num}_audio.mp3"
            tts = gTTS(text=text, lang='vi', slow=False)
            tts.save(audio_path)
            
            # Tạo video
            clip = ColorClip(AppConfig.VIDEO_RESOLUTION, color=bg_color, duration=15)
            txt_clip = TextClip(
                text,
                fontsize=28,
                color='white',
                font=AppConfig.FONT_PATH,
                size=(1100, 600),
                method='caption',
                align='center'
            ).set_position('center').set_duration(15)
            
            video = CompositeVideoClip([clip, txt_clip])
            video = video.set_audio(AudioFileClip(audio_path))
            
            scene_path = f"{self.output_dir}/scene_{scene_num}.mp4"
            video.write_videofile(scene_path, fps=24, codec='libx264', threads=4)
            
            return scene_path
        except Exception as e:
            logger.error(f"Lỗi tạo phân cảnh: {e}")
            return None
    
    def combine_scenes(self, scene_paths, output_name="final_video.mp4"):
        """Ghép các phân cảnh thành video hoàn chỉnh"""
        try:
            valid_scenes = [p for p in scene_paths if p is not None]
            if not valid_scenes:
                return None
                
            clips = [VideoFileClip(scene) for scene in valid_scenes]
            final_clip = concatenate_videoclips(clips)
            
            # Thêm nhạc nền nếu có
            if os.path.exists("background_music.mp3"):
                music = AudioFileClip("background_music.mp3").volumex(0.2)
                if music.duration > final_clip.duration:
                    music = music.subclip(0, final_clip.duration)
                final_audio = CompositeAudioClip([final_clip.audio, music])
                final_clip = final_clip.set_audio(final_audio)
            
            output_path = f"{self.output_dir}/{output_name}"
            final_clip.write_videofile(output_path, codec='libx264', audio_codec='aac')
            return output_path
        except Exception as e:
            logger.error(f"Lỗi ghép video: {e}")
            return None

# ====================== GIAO DIỆN NGƯỜI DÙNG ======================
def main():
    setup_page()
    st.title("⚽ AI Bình Luận Bóng Đá Chuyên Nghiệp")
    
    # Khởi tạo trình tải model
    if 'model_loader' not in st.session_state:
        st.session_state.model_loader = AIModelLoader()
        loader_thread = Thread(target=st.session_state.model_loader.load_in_main_thread)
        loader_thread.start()

    # Hiển thị trạng thái tải model
    if not st.session_state.model_loader.model_loaded:
        with st.spinner("🔄 Đang tải model AI, vui lòng chờ (có thể mất vài phút)..."):
            st.session_state.model_loader.loading_queue.get()
            
        if st.session_state.model_loader.error:
            st.error(f"Lỗi tải model: {st.session_state.model_loader.error}")
            return

    # Khởi tạo các thành phần
    model_pipeline = st.session_state.model_loader.get_model_pipeline()
    if not model_pipeline:
        st.error("Không thể khởi tạo pipeline AI")
        return
        
    news_processor = NewsProcessor()
    ai_commentator = AICommentator(model_pipeline)
    video_producer = VideoProducer()
    
    # Sidebar cài đặt
    with st.sidebar:
        st.header("⚙️ Cài Đặt")
        category = st.selectbox(
            "Chuyên mục tin tức:",
            list(AppConfig.NEWS_SOURCES.keys()),
            index=0
        )
        style = st.selectbox(
            "Phong cách bình luận:",
            ["Chuyên sâu", "Truyền cảm hứng", "Phân tích kỹ thuật", "Bình luận vui"],
            index=0
        )
        num_news = st.slider(
            "Số lượng tin:",
            1, AppConfig.MAX_NEWS_ITEMS, 3
        )
    
    # Nút tạo video
    if st.button("🎬 Tạo Video Bình Luận"):
        with st.spinner("📡 Đang lấy tin tức mới nhất..."):
            news_items = news_processor.fetch_news(category, num_news)
        
        if not news_items:
            st.error("Không tìm thấy tin tức phù hợp. Vui lòng thử chuyên mục khác!")
        else:
            st.success(f"✅ Đã lấy {len(news_items)} tin tức mới nhất!")
            
            # Hiển thị tin tức
            st.subheader("📰 Tin Tức Được Chọn")
            cols = st.columns(2)
            for i, item in enumerate(news_items):
                with cols[i % 2]:
                    with st.expander(f"{i+1}. {item['title']}", expanded=True):
                        st.write(item['summary'])
                        st.markdown(f"[Đọc chi tiết]({item['link']})")
            
            # Tạo bình luận
            st.subheader("🧠 AI Đang Viết Bình Luận...")
            commentaries = []
            progress_bar = st.progress(0)
            
            for i, item in enumerate(news_items):
                progress = int((i+1)/len(news_items)*100)
                progress_bar.progress(progress, f"Đang xử lý tin {i+1}/{len(news_items)}...")
                
                comment = ai_commentator.generate_commentary(item, style)
                if comment:
                    commentaries.append({
                        "title": item['title'],
                        "content": comment,
                        "original": item
                    })
            
            progress_bar.empty()
            
            # Hiển thị bình luận
            if not commentaries:
                st.error("Không thể tạo bình luận. Vui lòng thử lại!")
                return
                
            st.success("📝 Đã tạo xong bình luận!")
            for i, comm in enumerate(commentaries):
                with st.expander(f"🎤 Bình luận {i+1}: {comm['title']}", expanded=(i==0)):
                    st.write(comm['content'])
            
            # Tạo video
            if st.button("🎥 Tạo Video Từ Bình Luận", key="generate_video"):
                with st.spinner("📽️ Đang sản xuất video (có thể mất vài phút)..."):
                    scene_paths = []
                    for i, comm in enumerate(commentaries):
                        scene_path = video_producer.create_scene(
                            f"{comm['title']}\n\n{comm['content']}", 
                            i
                        )
                        scene_paths.append(scene_path)
                    
                    final_video = video_producer.combine_scenes(
                        scene_paths,
                        f"binh_luan_{datetime.now().strftime('%Y%m%d_%H%M')}.mp4"
                    )
                
                if final_video:
                    st.success("🎉 Video đã sẵn sàng!")
                    st.video(final_video)
                    
                    # Nút tải về
                    with open(final_video, "rb") as f:
                        st.download_button(
                            "⬇️ Tải Video Về Máy",
                            f,
                            file_name=os.path.basename(final_video),
                            mime="video/mp4"
                        )
                else:
                    st.error("Không thể tạo video. Vui lòng thử lại!")

    # Footer
    st.markdown("---")
    st.markdown("""
    <div class="footer">
        <p>Ứng dụng sử dụng AI để tạo bình luận bóng đá tự động</p>
        <p>© 2023 - Phát triển bởi AI Sports Commentary Team</p>
    </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()