import argparse
import sys
from pathlib import Path
from typing import Any, List

import streamlit as st
from sqlalchemy.orm import sessionmaker
from transformers import AutoTokenizer

from trinity.buffer.schema import init_engine
from trinity.common.config import StorageConfig
from trinity.common.experience import Experience
from trinity.common.experience_visualizer import build_experience_token_view
from trinity.utils.log import get_logger


class SQLExperienceViewer:
    def __init__(self, config: StorageConfig) -> None:
        self.logger = get_logger(f"sql_{config.name}", in_ray_actor=True)
        if not config.path:
            raise ValueError("`path` is required for SQL storage type.")
        self.engine, self.table_model_cls = init_engine(
            db_url=config.path,
            table_name=config.name,
            schema_type=config.schema_type,
        )
        self.session = sessionmaker(bind=self.engine)

    def get_experiences(self, offset: int, limit: int = 10) -> List[Experience]:
        self.logger.info("Viewing experiences from offset %s with limit %s.", offset, limit)
        with self.session() as session:
            query = session.query(self.table_model_cls).offset(offset).limit(limit)
            results = query.all()
            exps = [self.table_model_cls.to_experience(row) for row in results]
        return exps

    def total_experiences(self) -> int:
        with self.session() as session:
            count = session.query(self.table_model_cls).count()
        return count

    @staticmethod
    def run_viewer(
        model_path: str, db_url: str, table_name: str, schema_type: str, port: int
    ) -> None:
        """Start the Streamlit viewer.

        Args:
            model_path (str): Path to the tokenizer/model directory.
            db_url (str): Database URL for the experience database.
            table_name (str): Name of the experience table in the database.
            schema_type (str): Schema type of the experience table.
            port (int): Port number to run the Streamlit app on.
        """

        from streamlit.web import cli

        viewer_path = Path(__file__)
        sys.argv = [
            "streamlit",
            "run",
            str(viewer_path.resolve()),
            "--server.port",
            str(port),
            "--server.fileWatcherType",
            "none",
            "--",
            "--db-url",
            db_url,
            "--table",
            table_name,
            "--schema",
            schema_type,
            "--tokenizer",
            model_path,
        ]
        sys.exit(cli.main())


st.set_page_config(page_title="Trinity-RFT Experience Visualizer", layout="wide")


def get_color_for_action_mask(action_mask_value: int) -> str:
    """Return color based on action_mask value"""
    if action_mask_value == 1:
        return "#c8e6c9"
    else:
        return "#ffcdd2"


def render_token_detail_html(html: str) -> None:
    """Render token detail in the adaptive expanded panel style."""
    with st.container(border=True):
        st.markdown("**🔍 Response Tokens Detail:**")
        st.html(html)


def render_experience(exp: Experience, tokenizer: Any) -> None:
    """Render a single experience sequence in Streamlit."""
    token_view = build_experience_token_view(exp, tokenizer)

    # HTML escape function
    def html_escape(text):
        return (
            text.replace("&", "&amp;")
            .replace("<", "&lt;")
            .replace(">", "&gt;")
            .replace('"', "&quot;")
            .replace("'", "&#39;")
        )

    # === Use Streamlit Native Components for Prompt and Response ===
    st.subheader(f"Experience [{exp.eid}]")

    # Prompt section using st.text_area
    st.markdown("**📝 Prompt:**")
    st.code(token_view.prompt_text, language=None, wrap_lines=True, line_numbers=True)

    # Response section using st.text_area
    st.markdown("**💬 Response:**")
    st.code(token_view.response_text, language=None, wrap_lines=True, line_numbers=True)

    # Reward and other info
    st.markdown("**🏆 Reward and Other Info:**")
    reward, info, metrics = st.columns(3)
    reward.metric("**Reward:**", f"{exp.reward or 0.0:.4f}")
    metrics.markdown("**Metrics:**")
    metrics.json(exp.metrics or {}, expanded=False)
    info.markdown("**Info:**")
    info.json(exp.info or {}, expanded=False)

    # Build HTML only for Response Tokens Detail
    html = """
    <style>
        .token-detail-root * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }

        .token-detail-root {
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Roboto', sans-serif;
            padding: 10px;
        }

        .token-detail-root .token-container {
            display: flex;
            flex-wrap: wrap;
            gap: 5px;
            padding: 15px;
            background-color: white;
            border-radius: 5px;
        }

        .token-detail-root .token-box {
            display: inline-flex;
            flex-direction: column;
            align-items: center;
            padding: 8px 12px;
            border-radius: 5px;
            border: 1px solid #ddd;
            min-width: 60px;
            transition: transform 0.2s, box-shadow 0.2s;
        }

        .token-detail-root .token-box:hover {
            transform: scale(1.5);
            box-shadow: 0 4px 8px rgba(0,0,0,0.2);
            z-index: 10;
        }

        .token-detail-root .token-text {
            font-family: 'Courier New', monospace;
            font-size: 14px;
            font-weight: bold;
            margin-bottom: 5px;
            text-align: center;
            word-break: break-all;
            max-width: 100px;
        }

        .token-detail-root .token-logprob {
            font-size: 11px;
            color: #555;
            font-family: 'Courier New', monospace;
            text-align: center;
        }
    </style>
    <div class="token-detail-root">
        <div class="token-container">
    """

    # Add each response token
    for token in token_view.response_tokens:
        bg_color = get_color_for_action_mask(int(token.is_action))

        # Handle special character display
        token_display = token.token_text.replace(" ", "␣").replace("\n", "↵").replace("\t", "⇥")
        token_display = html_escape(token_display)
        logprob_text = f"{token.logprob:.4f}" if token.logprob is not None else "N/A"

        html += f"""
                <div class="token-box" style="background-color: {bg_color};">
                    <div class="token-text">{token_display}</div>
                    <div class="token-logprob">{logprob_text}</div>
                </div>
        """
    html += """
        </div>
    </div>
    """

    render_token_detail_html(html)


def parse_args():
    parser = argparse.ArgumentParser(description="Experience Visualizer")
    parser.add_argument(
        "--db-url",
        type=str,
        help="Path to the experience database.",
    )
    parser.add_argument(
        "--table",
        type=str,
        help="Name of the experience table.",
    )
    parser.add_argument(
        "--schema",
        type=str,
        default="experience",
        choices=("experience", "sft"),
        help="Schema type of the experience table.",
    )
    parser.add_argument(
        "--tokenizer",
        type=str,
        required=True,
        help="Path to the tokenizer.",
    )
    return parser.parse_args()


def main():
    args = parse_args()

    # Initialize SQLExperienceViewer
    config = StorageConfig()
    config.name = args.table
    config.path = args.db_url
    config.schema_type = args.schema
    config.storage_type = "sql"
    viewer = SQLExperienceViewer(config)

    st.title("🎯 Trinity-RFT Experience Visualizer")
    if "page" not in st.session_state:
        st.session_state.page = 1

    # Add instructions
    with st.expander("ℹ️ Instructions"):
        st.markdown(
            """
        - **Green background**: action_mask = 1
        - **Red background**: action_mask = 0
        - **Top**: Token text (special characters: space=␣, newline=↵, tab=⇥)
        - **Bottom**: Logprob value of the token
        - Hover over token to zoom in
        """
        )

    # Get total sequence number
    total_seq_num = viewer.total_experiences()

    # Sidebar configuration
    st.sidebar.header("⚙️ Settings")

    # Pagination settings
    experiences_per_page = st.sidebar.slider(
        "Experiences per page", min_value=1, max_value=20, value=5
    )
    # Calculate total pages
    total_pages = (total_seq_num + experiences_per_page - 1) // experiences_per_page

    # Page selection (sidebar)
    current_page = st.sidebar.number_input(
        "Select page",
        min_value=1,
        max_value=max(1, total_pages),
        step=1,
        value=st.session_state.page,
    )
    if current_page != st.session_state.page:
        st.session_state.page = current_page
        st.rerun()

    # Show statistics
    st.sidebar.markdown("---")
    st.sidebar.metric("Total experiences", total_seq_num)
    st.sidebar.metric("Total pages", total_pages)
    st.sidebar.metric("Current page", f"{st.session_state.page}/{total_pages}")

    # Calculate offset
    offset = (st.session_state.page - 1) * experiences_per_page

    # Get experiences for current page
    experiences = viewer.get_experiences(offset, experiences_per_page)

    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)

    # Generate catalog in sidebar
    exp_catalog = []  # [(eid, subheader_text)]
    if experiences:
        for exp in experiences:
            exp_catalog.append(exp.eid)

    if exp_catalog:
        st.sidebar.markdown("---")
        st.sidebar.markdown("**Contents**")
        catalog_md = "\n".join([f"- [ {eid} ](#exp-{eid})" for eid in exp_catalog])
        st.sidebar.markdown(catalog_md, unsafe_allow_html=True)

    if experiences:
        for exp in experiences:
            st.markdown(f'<a name="exp-{exp.eid}"></a>', unsafe_allow_html=True)
            render_experience(exp, tokenizer)
    else:
        st.warning("No experience data found")

    # Pagination navigation
    st.markdown("---")
    col1, col2, col3 = st.columns([1, 2, 1])

    with col1:
        if st.session_state.page > 1:
            if st.button("⬅️ Previous Page"):
                st.session_state.page = st.session_state.page - 1
                st.rerun()

    with col2:
        st.markdown(
            f"<center>Page {st.session_state.page} / {total_pages}</center>", unsafe_allow_html=True
        )
    with col3:
        if st.session_state.page < total_pages:
            if st.button("Next Page ➡️"):
                st.session_state.page = st.session_state.page + 1
                st.rerun()


if __name__ == "__main__":
    main()
