"""Variable search functionality."""

import difflib
import html
import re
from pycen.explore.metadata import get_variables
from pycen.explore.theme_config import get_theme

class SearchResults:
    '''Container for search results.'''

    def __init__(self, results):
        """
        Parameters
        ----------
        results : list of dict
            Search results
        """
        self._results = results
        self.variables = [r['code'] for r in results]

    def show(self):
        """Display search results."""
        if not self._results:
            print("No results found.")
            return

        # detect environment
        try:
            from IPython import get_ipython
            ipython = get_ipython()
            if ipython and 'ZMQInteractiveShell' in str(type(ipython)):
                # Jupyter notebook
                from IPython.display import HTML, display
                display(HTML(self._render_html()))
                return
        except Exception:
            pass

        # terminal/script
        self._render_terminal()

    def _render_html(self):
        """Render as collapsible HTML tree."""
        theme = get_theme()
        highlight_vars = set(theme.get("highlight_vars", []))
        colors = theme.get("colors", {})
        cat_color = colors.get("category", "#F7B608")
        cat_text = colors.get("category_text", "#2c5282")
        table_color = colors.get("table", "#e96604")
        var_color = colors.get("var", "#40c3dc")
        hover_color = colors.get("hover", "#f7fafc")
        meta_color = colors.get("meta", "#718096")
        vars_color = colors.get("vars", "#4a5568")
        var_code_color = colors.get("var_code", "#a0aec0")
        why_color = colors.get("why", "#92400e")

        cats = {
            'B01': 'Age and Sex', 'B02': 'Race', 'B03': 'Hispanic Origin', 'B04': 'Ancestry',
            'B05': 'Foreign Born/Citizenship', 'B06': 'Place of Birth', 'B07': 'Migration',
            'B08': 'Commuting', 'B09': 'Children/Household', 'B10': 'Grandparents',
            'B11': 'Household/Family Type', 'B12': 'Marital Status', 'B13': 'Fertility',
            'B14': 'School Enrollment', 'B15': 'Education', 'B16': 'Language',
            'B17': 'Poverty', 'B18': 'Disability', 'B19': 'Income', 'B20': 'Earnings',
            'B21': 'Veteran Status', 'B22': 'Public Assistance', 'B23': 'Employment',
            'B24': 'Industry/Occupation', 'B25': 'Housing', 'B26': 'Group Quarters',
            'B27': 'Health Insurance', 'B28': 'Computer/Internet', 'B29': 'Citizen Voting-Age',
            'C01': 'Age and Sex', 'C02': 'Race', 'C03': 'Hispanic Origin', 'C04': 'Ancestry',
            'C05': 'Foreign Born/Citizenship', 'C06': 'Place of Birth', 'C07': 'Migration',
            'C08': 'Commuting', 'C09': 'Children/Household', 'C10': 'Grandparents',
            'C11': 'Household/Family Type', 'C12': 'Marital Status', 'C13': 'Fertility',
            'C14': 'School Enrollment', 'C15': 'Education', 'C16': 'Language',
            'C17': 'Poverty', 'C18': 'Disability', 'C19': 'Income', 'C20': 'Earnings',
            'C21': 'Veteran Status', 'C22': 'Public Assistance', 'C23': 'Employment',
            'C24': 'Industry/Occupation', 'C25': 'Housing', 'C26': 'Group Quarters',
            'C27': 'Health Insurance', 'C28': 'Computer/Internet', 'C29': 'Citizen Voting-Age'
        }

        categories = {}
        table_concepts = {}
        for r in self._results:
            code = r.get('code', '')
            label = r.get('label') or ''
            concept = r.get('concept') or ''
            # extract table code (everything before underscore)
            table = code.split('_')[0] if '_' in code else code
            # category based on first 3 characters (e.g., B25, C17)
            category = cats.get(table[:3], 'Other')
            categories.setdefault(category, {}).setdefault(table, []).append({
                'code': code,
                'label': label,
                'concept': concept
            })
            if table not in table_concepts or not table_concepts[table]:
                table_concepts[table] = concept

        # extract table codes from highlight variables
        highlight_tables = {v.split('_')[0] if '_' in v else v for v in highlight_vars}
        highlighted_categories = {
            cat for cat, tables in categories.items()
            if any(t in highlight_tables for t in tables.keys())
        }

        html_out = (
            '<style>'
            '.tree{font-family:monospace;font-size:13px}'
            'details{margin:5px 0}'
            'summary{cursor:pointer;padding:3px}'
            f'summary:hover{{background:{hover_color}}}'
            f'.cat>summary{{font-weight:bold;color:{cat_text}}}'
            f'.highlight-cat{{background:{cat_color};padding:2px 4px;border-radius:3px;color:{cat_text}}}'
            f'.meta{{color:{meta_color};font-size:11px;margin-left:20px}}'
            f'.vars{{margin-left:20px;color:{vars_color};line-height:1.6}}'
            f'.highlight-table{{background:{table_color};padding:2px 4px;border-radius:3px;color:{cat_text}}}'
            f'.highlight-var{{background:{var_color};font-weight:500;color:{vars_color}}}'
            f'.why{{color:{why_color};font-size:11px;font-style:italic;margin-left:20px}}'
            '</style><div class="tree">'
        )

        # race/ethnicity iteration suffixes
        race_suffixes = {
            'A': 'White Alone',
            'B': 'Black or African American Alone',
            'C': 'American Indian and Alaska Native Alone',
            'D': 'Asian Alone',
            'E': 'Native Hawaiian and Other Pacific Islander Alone',
            'F': 'Some Other Race Alone',
            'G': 'Two or More Races',
            'H': 'White Alone, Not Hispanic or Latino',
            'I': 'Hispanic or Latino'
        }

        for cat in sorted(categories.keys()):
            cat_tables = categories[cat]
            cat_h = ' class="highlight-cat"' if cat in highlighted_categories else ''
            html_out += f'<details class="cat"><summary{cat_h}>{html.escape(cat)} ({len(cat_tables)} tables)</summary>'
            for table in sorted(cat_tables.keys()):
                rows = sorted(cat_tables[table], key=lambda row: row.get('code', ''))
                first_code = rows[0]['code']

                # detect race/ethnicity iteration
                base_table = table.rstrip('ABCDEFGHI')
                race_suffix = table[len(base_table):] if len(table) > len(base_table) else ''
                race_desc = race_suffixes.get(race_suffix, '')
                suffix = f' ({race_desc})' if race_desc else ''

                has_h = table in highlight_tables
                tbl_h = ' class="highlight-table"' if has_h else ''
                concept = table_concepts.get(table, '')
                html_out += f'<details><summary{tbl_h}>{html.escape(table)}: {html.escape(concept.title())}{suffix}</summary>'
                if has_h:
                    html_out += '<div class="why">contains commonly used variables</div>'
                html_out += f'<div class="meta">Variables: {len(rows)}</div><div class="vars">'
                for row in rows:
                    label = row['label']
                    label_parts = label.split('!!') if label else ['']
                    indent = '&nbsp;' * (max(len(label_parts) - 1, 0) * 2)
                    var_h = ' class="highlight-var"' if row['code'] in highlight_vars else ''
                    label_text = html.escape(label_parts[-1]) if label_parts[-1] else 'N/A'
                    html_out += (
                        f'{indent}&bull; <span{var_h}>{label_text}</span> '
                        f'<span style="color:{var_code_color}">({html.escape(row["code"])})</span><br>'
                    )
                html_out += '</div></details>'
            html_out += '</details>'

        html_out += '</div>'
        return html_out

    def _render_terminal(self):
        """Render as text table."""
        try:
            from rich.console import Console
            from rich.table import Table

            console = Console()
            table = Table(title=f"Search Results ({len(self._results)} variables)")

            table.add_column("Variable", style="cyan")
            table.add_column("Label", style="white")
            table.add_column("Concept", style="dim")

            for r in self._results[:20]:  # limit display
                table.add_row(
                    r['code'],
                    r['label'][:60] + ('...' if len(r['label']) > 60 else ''),
                    r['concept'][:40] + ('...' if len(r['concept']) > 40 else '')
                )

            if len(self._results) > 20:
                table.add_row("...", f"... and {len(self._results) - 20} more", "...")

            console.print(table)
        except ImportError:
            # fallback to simple print
            print(f"\nSearch Results ({len(self._results)} variables):\n")
            for r in self._results[:20]:
                print(f"  {r['code']}: {r['label'][:60]}")
            if len(self._results) > 20:
                print(f"  ... and {len(self._results) - 20} more")

    def select(self, indices):
        """
        Select variables by index.

        Parameters
        ----------
        indices : list of int
            Indices to select

        Returns
        -------
        list of str
            Selected variable codes
        """
        return [self.variables[i] for i in indices]

def search(query, year, dataset):
    '''
    Search for variables using keywords or phrases.

    Parameters
    ----------
    query : str
        Search query (keywords or natural language)
    year : int
        Census year (required)
    dataset : str
        Dataset (e.g., 'acs5', 'acs1', 'dec_pl', 'dec_sf1')
    Returns
    -------
    SearchResults
        Search results object

    Examples
    --------
    >>> results = search('median income', year=2021, dataset='acs5')
    >>> results.show()
    >>> vars = results.select([0, 1, 2])
    '''
    if year is None:
        raise ValueError("year is required")

    return _keyword_search(query, year, dataset)

def browse(year, dataset):
    '''
    Browse all variables for a dataset/year.

    Parameters
    ----------
    year : int
        Census year
    dataset : str
        Dataset (e.g., 'acs5', 'acs1', 'dec_pl', 'dec_sf1')
    '''
    variables = get_variables(year, dataset)
    results = [
        {'code': code, 'label': info.get('label') or '', 'concept': info.get('concept') or ''}
        for code, info in variables.items()
    ]
    return SearchResults(results)

_SYNONYMS = {
    'wfh': ['work from home', 'worked from home', 'work at home', 'telework', 'telecommute'],
    'remote': ['work from home', 'telework', 'telecommute'],
    'vmt': ['vehicle miles traveled', 'vehicle miles travelled', 'vehicle miles', 'commute', 'travel time', 'transportation'],
    'vkt': ['vehicle kilometers traveled', 'vehicle kilometres traveled', 'vehicle kilometers', 'commute', 'travel time', 'transportation'],
}

def _expand_tokens(query):
    '''
    Expand query into search tokens with synonym expansion.

    Splits query into tokens and expands abbreviations/synonyms.
    Example: "wfh" --> ["wfh", "work", "from", "home", "telework", "telecommute"]
    '''
    tokens = [t for t in re.split(r'[^A-Za-z0-9]+', query.lower()) if t]
    expanded = set(tokens)
    for tok in tokens:
        for phrase in _SYNONYMS.get(tok, []):
            expanded.update([t for t in re.split(r'[^A-Za-z0-9]+', phrase.lower()) if t])
    return list(expanded) if expanded else [query.lower()]

def _keyword_search(query, year, dataset):
    '''
    Search Census variables using two-phase matching.

    Phase 1 (Exact): Substring match on variable labels/concepts
    Phase 2 (Fuzzy): If no exact matches, use similarity matching (84% threshold)
                     with difflib.SequenceMatcher for typo tolerance

    Returns SearchResults object with matching variables.
    '''
    if year is None:
        raise ValueError("year is required")
    if not dataset:
        raise ValueError("dataset is required")

    tokens = _expand_tokens(query)
    variables = get_variables(year, dataset)

    # Phase 1: Exact substring matching
    results = []
    for code, info in variables.items():
        label = (info.get('label') or '').lower()
        concept = (info.get('concept') or '').lower()
        if any(tok in label or tok in concept for tok in tokens):
            results.append({
                'code': code,
                'label': info.get('label') or '',
                'concept': info.get('concept') or ''
            })

    # Phase 2: Fuzzy matching fallback (if no exact matches)
    if not results:
        for code, info in variables.items():
            label = (info.get('label') or '')
            concept = (info.get('concept') or '')
            words = re.split(r'[^A-Za-z0-9]+', f"{label} {concept}".lower())
            matched = False
            for tok in tokens:
                if len(tok) < 4:  # Skip short tokens for fuzzy matching
                    continue
                for w in words:
                    if not w:
                        continue
                    if difflib.SequenceMatcher(None, tok, w).ratio() >= 0.84:
                        matched = True
                        break
                if matched:
                    break
            if matched:
                results.append({
                    'code': code,
                    'label': info.get('label') or '',
                    'concept': info.get('concept') or ''
                })

    return SearchResults(results)
