import sqlite3
import textfsm
from typing import Dict, List, Tuple, Optional
import io
import time
import click
from multiprocessing import Process, Queue
import multiprocessing
import sys
import threading
from contextlib import contextmanager


class ThreadSafeConnection:
    """Thread-local storage for SQLite connections"""

    def __init__(self, db_path: str, verbose: bool = False):
        self.db_path = db_path
        self.verbose = verbose
        self._local = threading.local()

    @contextmanager
    def get_connection(self):
        """Get a thread-local connection"""
        if not hasattr(self._local, 'connection'):
            self._local.connection = sqlite3.connect(self.db_path)
            self._local.connection.row_factory = sqlite3.Row
            if self.verbose:
                click.echo(f"Created new connection in thread {threading.get_ident()}")

        try:
            yield self._local.connection
        except Exception as e:
            if hasattr(self._local, 'connection'):
                self._local.connection.close()
                delattr(self._local, 'connection')
            raise e

    def close_all(self):
        """Close connection if it exists for current thread"""
        if hasattr(self._local, 'connection'):
            self._local.connection.close()
            delattr(self._local, 'connection')


class TextFSMAutoEngine:
    def __init__(self, db_path: str, verbose: bool = False):
        self.db_path = db_path
        self.verbose = verbose
        self.connection_manager = ThreadSafeConnection(db_path, verbose)

    # Vendor tokens — used to separate platform prefix from command keywords
    _VENDOR_TOKENS = frozenset({
        'arista', 'eos', 'cisco', 'ios', 'iosxe', 'iosxr', 'nxos',
        'juniper', 'junos', 'paloalto', 'panos', 'fortinet', 'fortios',
        'linux', 'vyos', 'hp', 'comware', 'huawei', 'vrp',
    })

    def _split_filter_terms(self, filter_string: str) -> Tuple[List[str], List[str]]:
        """
        Split a filter string into vendor prefix terms and command terms.

        Returns (vendor_terms, command_terms) where each term is >2 chars.
        Vendor tokens are consumed greedily from the front of the string.
        Command terms are stem-normalized: trailing 's' is stripped so that
        LIKE '%neighbor%' matches both "neighbors" and "neighbor_interface".

        Example:
          "juniper_junos_show_lldp_neighbors_interface_ge-0/0/0"
          → vendor: ["juniper", "junos"]
          → command: ["show", "lldp", "neighbor", "interface", "ge"]
        """
        raw_terms = filter_string.replace('-', '_').split('_')
        vendor_terms = []
        command_terms = []
        past_vendor = False

        for term in raw_terms:
            if not term or len(term) <= 2:
                continue
            if not past_vendor and term.lower() in self._VENDOR_TOKENS:
                vendor_terms.append(term)
            else:
                past_vendor = True
                # Strip trailing 's' for naive plural normalization.
                # "neighbors" → "neighbor" so LIKE catches both singular
                # and plural template names. Min length 4 avoids mangling
                # short words like "dns", "bgp", "ios".
                normalized = term
                if len(term) > 4 and term.lower().endswith('s'):
                    normalized = term[:-1]
                command_terms.append(normalized)

        return vendor_terms, command_terms

    def _calculate_template_score(
            self,
            parsed_data: List[Dict],
            template: sqlite3.Row,
            raw_output: str
    ) -> float:
        """
        Score template match quality (0-100 scale).

        Factors:
        - Record count (0-30 pts): Did the template find data?
        - Field richness (0-30 pts): How many fields per record?
        - Population rate (0-25 pts): Are fields actually filled?
        - Consistency (0-15 pts): Uniform data across records?
        """
        if not parsed_data:
            return 0.0

        num_records = len(parsed_data)
        num_fields = len(parsed_data[0].keys()) if parsed_data else 0
        is_version_cmd = 'version' in template['cli_command'].lower()

        # === Factor 1: Record Count (0-30 points) ===
        if is_version_cmd:
            record_score = 30.0 if num_records == 1 else max(0, 15 - (num_records - 1) * 5)
        else:
            if num_records >= 10:
                record_score = 30.0
            elif num_records >= 3:
                record_score = 20.0 + (num_records - 3) * (10.0 / 7.0)
            else:
                record_score = num_records * 10.0

        # === Factor 2: Field Richness (0-30 points) ===
        if num_fields >= 10:
            field_score = 30.0
        elif num_fields >= 6:
            field_score = 20.0 + (num_fields - 6) * 2.5
        elif num_fields >= 3:
            field_score = 10.0 + (num_fields - 3) * (10.0 / 3.0)
        else:
            field_score = num_fields * 5.0

        # === Factor 3: Population Rate (0-25 points) ===
        total_cells = num_records * num_fields
        populated_cells = 0

        for record in parsed_data:
            for value in record.values():
                if value is not None and str(value).strip():
                    populated_cells += 1

        population_rate = populated_cells / total_cells if total_cells > 0 else 0
        population_score = population_rate * 25.0

        # === Factor 4: Consistency (0-15 points) ===
        if num_records > 1:
            field_fill_counts = {key: 0 for key in parsed_data[0].keys()}

            for record in parsed_data:
                for key, value in record.items():
                    if value is not None and str(value).strip():
                        field_fill_counts[key] += 1

            consistent_fields = sum(
                1 for count in field_fill_counts.values()
                if count == 0 or count == num_records
            )
            consistency_rate = consistent_fields / num_fields if num_fields > 0 else 0
            consistency_score = consistency_rate * 15.0
        else:
            consistency_score = 15.0

        total_score = record_score + field_score + population_score + consistency_score

        if self.verbose:
            click.echo(f"    Scoring: records={record_score:.1f}, fields={field_score:.1f}, "
                       f"population={population_score:.1f}, consistency={consistency_score:.1f} "
                       f"-> {total_score:.1f}")

        return total_score

    def find_best_template(self, device_output: str, filter_string: Optional[str] = None) -> Tuple[
        Optional[str], Optional[List[Dict]], float, List[Tuple[str, float, int]]]:
        """
        Try filtered templates against the output and return the best match
        plus all non-zero scores.

        Uses progressive filter relaxation: starts with up to 3 command
        keywords (plus all vendor tokens), then drops one keyword at a time
        until templates are found or only vendor terms remain.

        This keeps the SQL query fast (narrow candidate set) while preventing
        over-filtering on arguments, plurals, or interface names that don't
        appear in template names.

        Example progression for "juniper_junos_show_lldp_neighbors_interface_ge-0/0/0":
          Pass 1: vendor=[juniper,junos] + cmd=[show,lldp,neighbors] → 0 hits
          Pass 2: vendor=[juniper,junos] + cmd=[show,lldp]           → 4 hits → score them
        """
        best_template = None
        best_parsed_output = None
        best_score = 0
        all_scores = []

        with self.connection_manager.get_connection() as conn:
            templates = []

            if filter_string:
                vendor_terms, command_terms = self._split_filter_terms(filter_string)

                # Cap initial attempt at 3 command terms — beyond that is
                # almost always arguments (interface names, IPs, VRFs)
                max_cmd = min(len(command_terms), 3)

                # Progressive relaxation: try N cmd terms, then N-1, ...
                # until we get candidates or exhaust down to vendor-only
                for n_cmd in range(max_cmd, -1, -1):
                    attempt_terms = vendor_terms + command_terms[:n_cmd]
                    templates = self.get_filtered_templates(conn, terms=attempt_terms)

                    if self.verbose:
                        click.echo(
                            f"Filter pass ({n_cmd} cmd terms): "
                            f"{attempt_terms} → {len(templates)} candidates"
                        )

                    if templates:
                        break

                # If still nothing (bad vendor hint?), fall through unfiltered
                if not templates:
                    if self.verbose:
                        click.echo("All filter passes returned 0 — trying unfiltered")
                    templates = self.get_filtered_templates(conn, terms=None)
            else:
                templates = self.get_filtered_templates(conn, terms=None)

            total_templates = len(templates)

            if self.verbose:
                click.echo(f"Scoring {total_templates} candidate templates")

            for idx, template in enumerate(templates, 1):
                if self.verbose:
                    percentage = (idx / total_templates) * 100
                    click.echo(f"\nTemplate {idx}/{total_templates} ({percentage:.1f}%): {template['cli_command']}")

                try:
                    textfsm_template = textfsm.TextFSM(io.StringIO(template['textfsm_content']))
                    parsed = textfsm_template.ParseText(device_output)
                    parsed_dicts = [dict(zip(textfsm_template.header, row)) for row in parsed]
                    score = self._calculate_template_score(parsed_dicts, template, device_output)

                    if self.verbose:
                        click.echo(f" -> Score={score:.2f}, Records={len(parsed_dicts)}")

                    if score > 0:
                        all_scores.append((template['cli_command'], score, len(parsed_dicts)))

                    if score > best_score:
                        best_score = score
                        best_template = template['cli_command']
                        best_parsed_output = parsed_dicts
                        if self.verbose:
                            click.echo(click.style("  New best match!", fg='green'))

                except Exception as e:
                    if self.verbose:
                        click.echo(f" -> Failed to parse: {str(e)}")
                    continue

        all_scores.sort(key=lambda x: x[1], reverse=True)
        return best_template, best_parsed_output, best_score, all_scores

    def get_filtered_templates(
        self,
        connection: sqlite3.Connection,
        filter_string: Optional[str] = None,
        terms: Optional[List[str]] = None,
    ):
        """
        Get templates from database, optionally filtered by LIKE terms.

        Accepts either:
          - terms: pre-split list of terms (used by find_best_template)
          - filter_string: raw string to split (legacy / direct callers)
        """
        cursor = connection.cursor()

        # Resolve terms from either parameter
        if terms is None and filter_string:
            raw = filter_string.replace('-', '_').split('_')
            terms = [t for t in raw if t and len(t) > 2]

        if terms:
            query = "SELECT * FROM templates WHERE 1=1"
            params = []
            for term in terms:
                if term and len(term) > 2:
                    query += " AND cli_command LIKE ?"
                    params.append(f"%{term}%")
            cursor.execute(query, params)
        else:
            cursor.execute("SELECT * FROM templates")
        return cursor.fetchall()

    def __del__(self):
        """Clean up connections on deletion"""
        self.connection_manager.close_all()


# Example usage
if __name__ == '__main__':
    multiprocessing.freeze_support()


    # Example of using the engine in multiple threads
    def worker(engine, output, filter_str):
        result = engine.find_best_template(output, filter_str)
        print(f"Thread {threading.get_ident()}: Found template: {result[0]}")


    engine = TextFSMAutoEngine("./secure_cartography/tfsm_templates.db", verbose=True)
    threads = []

    # Create multiple threads
    for i in range(3):
        t = threading.Thread(target=worker, args=(engine, "sample output", "show version"))
        threads.append(t)
        t.start()

    # Wait for all threads to complete
    for t in threads:
        t.join()