import os
import pandas as pd
import streamlit as st

# AI & Database Imports
from langchain_core.documents import Document
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEndpointEmbeddings

# -------------------------------
# 0. CONFIG & STYLING
# -------------------------------
st.set_page_config(
    page_title="IPL 2026 Insight AI",
    page_icon="🏏",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# --- IPL 2026 THEME FONTS & CSS ---
st.markdown("""
    <link href="https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;600;800&family=Plus+Jakarta+Sans:wght@400;600;700&display=swap" rel="stylesheet">
    <style>
        /* --- GLOBAL THEME VARS --- */
        :root {
            --ipl-orange: #FF4B20;
            --ipl-blue: #1A1F35;
            --ipl-dark: #0B0E14;
            --ipl-grey: #C8CDD5;
            --ipl-card-bg: #161B2C;
            --ipl-green: #00D26A;
        }

        /* --- FONTS --- */
        html, body, [class="stApp"] {
            font-family: 'Outfit', sans-serif;
            background-color: var(--ipl-dark);
            color: white;
        }
        
        h1, h2, h3, h4, h5, h6 {
            font-family: 'Plus Jakarta Sans', sans-serif !important;
            font-weight: 800;
            letter-spacing: -0.5px;
            color: white !important;
        }

        /* --- CUSTOM CARD CONTAINERS --- */
        .custom-card {
            background-color: var(--ipl-card-bg);
            border: 1px solid #2D3348;
            border-radius: 16px;
            padding: 20px;
            box-shadow: 0 8px 24px rgba(0,0,0,0.2);
            margin-bottom: 20px;
        }

        /* --- INPUT FIELDS --- */
        .stTextInput > div > div > input {
            background-color: var(--ipl-card-bg);
            border: 2px solid #2D3348;
            color: white;
            border-radius: 12px;
            padding: 15px;
            font-size: 18px;
            transition: all 0.3s;
            height: 50px;
        }
        .stTextInput > div > div > input:focus {
            border-color: var(--ipl-orange);
            outline: none;
            box-shadow: 0 0 15px rgba(255, 75, 32, 0.2);
        }

        /* --- BUTTONS --- */
        .stButton > button {
            background: linear-gradient(135deg, var(--ipl-orange) 0%, #FF2E00 100%);
            border: none;
            border-radius: 12px;
            color: white;
            font-weight: 700;
            padding: 10px 30px;
            text-transform: uppercase;
            letter-spacing: 1px;
            transition: transform 0.2s;
            height: 50px;
            width: 100%;
        }
        .stButton > button:hover {
            transform: scale(1.02);
            box-shadow: 0 0 20px rgba(255, 75, 32, 0.4);
        }

        /* --- CHAT BUBBLES --- */
        .chat-response {
            background: #1E2336;
            border-left: 4px solid var(--ipl-orange);
            border-radius: 0px 12px 12px 12px;
            padding: 20px;
            margin-bottom: 15px;
            color: #E0E3EB;
            line-height: 1.6;
            box-shadow: 0 4px 12px rgba(0,0,0,0.1);
        }
        
        /* --- HIDE DEFAULT STREAMLIT STYLE --- */
        #MainMenu {visibility: hidden;}
        footer {visibility: hidden;}
        header {visibility: hidden;}

        /* Fix alignment gap */
        [data-testid="stHorizontalBlock"] {
            gap: 10px;
        }
        
    </style>
""", unsafe_allow_html=True)


# -------------------------------
# 1. DATA LOADING & PROCESSING
# -------------------------------
@st.cache_data(ttl=3600)
def load_data():
    csv_files = [
        "batting_stats.csv", "bowling_stats.csv", "deliveries.csv",
        "fielding_stats.csv", "matches.csv", "points_table.csv",
        "squads.csv", "venues.csv"
    ]
    dataframes = {}
    
    with st.spinner('Loading Match Data...'):
        for file_name in csv_files:
            if os.path.exists(file_name):
                key_name = file_name.replace(".csv", "")
                try:
                    dataframes[key_name] = pd.read_csv(file_name)
                except Exception as e:
                    print(f"Error: {e}")
        return dataframes


def process_all_tables_into_sentences(all_tables):
    sentences = []
    
    if "matches" in all_tables:
        df = all_tables["matches"]
        for _, row in df.iterrows():
            winner = row.get('match_winner', 'No result')
            venue = row.get('venue', 'Unknown stadium')
            m_id = row.get('match_id', 'N/A')
            text = f"In IPL 2026 Match {m_id} at {venue}, {winner} won the match"
            if 'wb_runs' in row and row['wb_runs'] > 0:
                text += f" by {row['wb_runs']} runs."
            elif 'wb_wickets' in row and row['wb_wickets'] > 0:
                text += f" by {row['wb_wickets']} wickets."
            else:
                text += "."
            sentences.append(text)

    if "batting_stats" in all_tables:
        df = all_tables["batting_stats"]
        for _, row in df.iterrows():
            sentences.append(
                f"Player {row.get('player_name')} scored {row.get('runs')} runs in a match, "
                f"with a strike rate of {row.get('strike_rate')} and hit {row.get('fours', 0)} fours and {row.get('sixes', 0)} sixes."
            )

    if "bowling_stats" in all_tables:
        df = all_tables["bowling_stats"]
        for _, row in df.iterrows():
            sentences.append(
                f"Bowler {row.get('player_name')} took {row.get('wickets')} wickets while conceding {row.get('runs_conceded')} runs, "
                f"with an economy rate of {row.get('economy')}."
            )

    if "points_table" in all_tables:
        df = all_tables["points_table"]
        for _, row in df.iterrows():
            sentences.append(
                f"In the IPL 2026 points table, team {row.get('team_name')} is ranked at position {row.get('position')} "
                f"with {row.get('points')} points, having won {row.get('won')} matches and lost {row.get('lost')}."
            )

    if "squads" in all_tables:
        df = all_tables["squads"]
        for _, row in df.iterrows():
            sentences.append(
                f"Player {row.get('player_name')} plays for team {row.get('team_name')} as a {row.get('role', 'player')} in IPL 2026."
            )

    if "venues" in all_tables:
        df = all_tables["venues"]
        for _, row in df.iterrows():
            sentences.append(
                f"The stadium {row.get('venue')} is located in the city of {row.get('city', 'India')}."
            )

    if "deliveries" in all_tables:
        df = all_tables["deliveries"].head(500)
        for _, row in df.iterrows():
            sentences.append(
                f"In Match {row.get('match_id')}, over {row.get('over_num')}, ball {row.get('ball_num')}: "
                f"Bowler {row.get('bowler')} bowled to batsman {row.get('batsman')}, resulting in {row.get('total_runs')} runs."
            )

    if "fielding_stats" in all_tables:
        df = all_tables["fielding_stats"]
        for _, row in df.iterrows():
            sentences.append(
                f"Fielder {row.get('player_name')} took {row.get('catches', 0)} catches and made {row.get('stumpings', 0)} stumpings."
            )
            
    return sentences


# -------------------------------
# 2. VECTOR DATABASE BUILDER
# -------------------------------
@st.cache_resource
def get_vector_db():
    all_tables = load_data()
    if not all_tables:
        return None
        
    sentences = process_all_tables_into_sentences(all_tables)
    if not sentences:
        return None

    docs = [Document(page_content=s) for s in sentences]

    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        st.error("Missing HF_TOKEN")
        return None

    with st.spinner('Building Vector Index...'):
        embeddings = HuggingFaceEndpointEmbeddings(
            model="sentence-transformers/all-MiniLM-L6-v2",
            task="feature-extraction",
            huggingfacehub_api_token=hf_token
        )
        vector_db = FAISS.from_documents(docs, embeddings)
    
    return vector_db


# -------------------------------
# 3. MAIN APPLICATION
# -------------------------------

# Header
st.markdown("""
<div style="display:flex; justify-content:space-between; align-items:center; margin-bottom:30px;">
    <div>
        <h1 style="margin:0; font-size: 50px;">IPL 2026 🏏</h1>
        <p style="margin:0; color: #8B9CC3; font-size: 18px;">Official Smart Insight Assistant</p>
    </div>
    <div style="text-align:right;">
        <span style="background:#00D26A; color:#000; padding:10px 20px; border-radius:0px; font-weight:800; font-size:16px; letter-spacing:1px;">LIVE SEASON</span>
    </div>
</div>
""", unsafe_allow_html=True)

# Initialize Data
all_tables = load_data()
vector_db = get_vector_db()

# Input Area
st.markdown('<div class="custom-card">', unsafe_allow_html=True)

# Fixed alignment using vertical alignment CSS
col1, col2 = st.columns([4, 1])
with col1:
    query = st.text_input("Ask your question:", placeholder="e.g., Who are the top batsmen for RCB?", label_visibility="collapsed")
with col2:
    # Add some margin-top to align with input
    st.markdown("<div style='margin-top: 8px;'></div>", unsafe_allow_html=True)
    search_btn = st.button("🔍 Analyze")

st.markdown('</div>', unsafe_allow_html=True)

# Results Area
if query and search_btn:
    if vector_db is None:
        st.error("Database Offline.")
    else:
        with st.spinner('Parsing Data...'):
            results = vector_db.similarity_search(query, k=4)
            
        st.subheader(f"🔎 Found {len(results)} Insights")
        
        # Loop through results
        for res in results:
            html_card = f"""
            <div class="chat-response">
                {res.page_content}
            </div>
            """
            st.markdown(html_card, unsafe_allow_html=True)

elif query and not search_btn:
    st.info("Hit the 'Analyze' button to search the database.")

# Dashboard / Data Explorer at bottom
if all_tables:
    with st.expander("IPL 2026 Overall Statics"):
        tab1, tab2, tab3 = st.tabs(["Points Table", "Batting", "Bowling"])
        with tab1:
            st.dataframe(all_tables.get("points_table"), use_container_width=True)
        with tab2:
            st.dataframe(all_tables.get("batting_stats"), use_container_width=True)
        with tab3:
            st.dataframe(all_tables.get("bowling_stats"), use_container_width=True)