#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.


import json
import logging
import os
import re
from dataclasses import asdict
from datetime import datetime, timedelta
from typing import List

import pandas as pd

from models.agent import Agent
from models.policy import Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import get_protected_json, load_env
from utils.selector import Selector
from utils.utils import colorText, get_sanitized_input

logger = logging.getLogger(__name__)


def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
    agents = selectAgents(api)
    history_days = Selector.select_value(
        prompt="Enter how many days of history to pull (1–150): ",
        value_type=int,
        valid_range=(1, 150),
    )

    if not agents or not history_days:
        print(colorText("No agents selected or invalid history range.", "red"))
        return

    historical_date = (datetime.now() - timedelta(days=history_days)).strftime("%Y-%m-%d")
    today = datetime.now().strftime("%Y-%m-%d")

    all_history = []

    for agent in agents:
        try:
            exechistory = api.history_execution(today, historical_date, agent.hostname)
        except Exception as e:
            print(colorText(f"❌ Error retrieving history for {agent.hostname}: {e}", "red"))
            continue

        if isinstance(exechistory, list):
            for block in exechistory:
                record = {
                    "Command": block.get("commandline", "N/A"),
                    "Date": block.get("datetime", "N/A"),
                    "Filename": block.get("filename", "N/A"),
                    "Policy Name": block.get("policyname", "N/A"),
                    "Hostname": block.get("hostname", "N/A"),
                    "Hash": block.get("sha256", "N/A"),
                }
                all_history.append(record)

                if not outputjson:
                    for key, value in record.items():
                        print(colorText(f"{key}: {value}", "green"))
                    print("\n")
        else:
            print(colorText(f"No execution history found for {agent.hostname}.", "yellow"))

    if outputjson:
        print(json.dumps(all_history, indent=2))


def findAllAgents(api):
    # Step 1: Load data from API
    policies = [Policy(**row["data"]) for _, row in api.policy_find_all().iterrows()]
    agents = [Agent(**row["data"]) for _, row in api.agent_find_all().iterrows()]

    # Step 2: Create groupid → groupname map
    groupid_to_name = {policy.groupid: policy.name for policy in policies}

    # Step 3: Enrich agents
    for agent in agents:
        agent.enrich(groupid_to_name)

    return agents

def findAgents(api, return_dataframe):
    agents = selectAgents(api)
    working_dir = load_env("WORKING_DIR")

    if not agents:
        logging.warning("No agents or policies found.")
        print("No agents matched the criteria.")
        return

    # Convert enriched agents to DataFrame
    agent_dicts = [asdict(agent) for agent in agents]
    agent_df = pd.DataFrame(agent_dicts)

    if return_dataframe:
        logging.debug("Returning DataFrame to caller.")
        return agent_df

    # Otherwise, print and optionally export
    print(agent_df)
    logging.debug("Displayed DataFrame to console.")

    user_input = get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ").strip().lower()
    if user_input == 'y':
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        filename = f"agentsearch_{timestamp}.csv"
        file_path = os.path.join(working_dir, filename)

        agent_df.to_csv(file_path, index=False)
        logging.info(f"Exported DataFrame to {file_path}")

        print(
            colorText(
                f"\n✅ Matched devices exported to: {working_dir}\\{filename}",
                "green",
            )
        )
    else:
        logging.debug("User declined to export the DataFrame.")


def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
    print(colorText("🔍 Device Search", "cyan"))
    print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan"))
    print(colorText("When you're done, press Enter twice.\n", "cyan"))
    print(colorText("Example:", "cyan"))
    print(colorText("H00000", "cyan"))
    print(colorText("UTN00000", "cyan"))
    print(colorText("i-hSuperSecretServer", "cyan"))
    print(colorText("u-hVenderBroke\n", "cyan"))

    print(colorText("Paste or type your device names below:", "white"))
    policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]

    device_input_lines = []
    empty_line_count = 0

    # Regex to validate each line
    valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$')

    while True:
        line = get_sanitized_input("")
        stripped_line = line.strip()

        if stripped_line == "":
            empty_line_count += 1
            if empty_line_count == 2:
                break
            continue  # Don't validate empty lines
        else:
            empty_line_count = 0

        # Validate only non-empty lines
        if valid_line_pattern.match(stripped_line):
            device_input_lines.append(stripped_line)
        else:
            print(colorText(f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", "yellow"))
            
    device_names = [name for name in device_input_lines if name]
    if not device_names:
        logger.debug("No device names entered")
        print(colorText("⚠️ No device names entered.", "red"))
        return []

    # Build regex pattern to match hostnames
    pattern = "|".join(map(re.escape, device_names))
    regex = re.compile(pattern, re.IGNORECASE)

    # Fetch agents
    agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()]
    matched_agents = [agent for agent in agents if regex.search(agent.hostname)]
    matched_agents.sort(key=lambda agent: agent.hostname.lower())

    # Show unmatched
    unmatched = [name for name in device_names if not any(regex.search(agent.hostname) for agent in agents)]
    if unmatched:
        logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
        print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))

    if not matched_agents:
        logger.debug("❌ No matching devices found.")
        print(colorText("❌ No matching devices found.", "red"))
    else:
        logger.debug(f"✅ Found {len(matched_agents)} matching device(s).")        
        print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green"))

    # Enrich each agent using its class method
    for agent in matched_agents:
        agent.enrich_with_policies(policies)

    return matched_agents
def moveAgentToRelatedPolicy(
    api: AirlockAPIWrapper,
    agent: Agent,
    mode: str = "audit",
):
    """
    Moves an agent between audit and enforcement policies based on the mode.

    Args:
        api: AirlockAPIWrapper instance.
        agent: Agent object.
        policy_relationship_map: Dict mapping enforcement → audit.
        mode: 'audit' to move to audit, 'enforcement' to move to enforcement.
    """
    policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}")

    if mode == "audit":
        if agent.groupid in policy_relationship_map:
            target_policy = policy_relationship_map[agent.groupid]
        elif agent.groupid in policy_relationship_map.values():
            logger.debug(f"Agent {agent.hostname} is already in an audit group. No action needed.")
            print(f"Agent {agent.hostname} is already in an audit group. No action needed.")
            return
        else:
            logger.warning(f"Error: No corresponding audit policy found for groupid: {agent.groupid}.")
            return

    elif mode == "enforcement":
        inverse_map = {v: k for k, v in policy_relationship_map.items()}
        if agent.groupid in inverse_map:
            target_policy = inverse_map[agent.groupid]
        elif agent.groupid in inverse_map.values():
            logger.info(f"Agent {agent.hostname} is already in an enforcement group. No action needed.")
            return
        else:
            logger.warning(f"Error: No corresponding enforcement policy found for groupid: {agent.groupid}.")
            return

    else:
        logger.error(f"Unknown mode '{mode}'. Use 'audit' or 'enforcement'.")
        return

    api.agent_move(agent.agentid, target_policy)