import asyncio
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 services.TaskQueue import AsyncTaskQueue, run_sync_task_in_thread
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__)

async def devicehistory(api: AirlockAPIWrapper, outputjson: bool):
    agents = await selectAgents(api)
    history_days = await 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 = []

    async def fetch_history(agent):
        try:
            exechistory = await api.history_execution(today, historical_date, agent.hostname)
            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"))
        except Exception as e:
            print(colorText(f"❌ Error retrieving history for {agent.hostname}: {e}", "red"))

    await asyncio.gather(*(fetch_history(agent) for agent in agents))

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

async def findAllAgents(api: AirlockAPIWrapper):
    policies_df = await api.policy_find_all()
    agents_df = await api.agent_find_all()

    policies = [Policy(**row["data"]) for _, row in policies_df.iterrows()]
    agents = [Agent(**row["data"]) for _, row in agents_df.iterrows()]

    groupid_to_name = {policy.groupid: policy.name for policy in policies}

    queue = AsyncTaskQueue()
    await queue.start_workers()

    async def enrich_agent(agent):
        agent.enrich(groupid_to_name)

    for agent in agents:
        await queue.enqueue(f"enrich_{agent.hostname}", enrich_agent, agent)

    await asyncio.sleep(1)
    await queue.stop_workers()

    return agents

async def findAgents(api: AirlockAPIWrapper, return_dataframe: bool):
    agents = await selectAgents(api)
    working_dir = await load_env("WORKING_DIR")

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

    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

    print(agent_df)
    logging.debug("Displayed DataFrame to console.")

    user_input = await get_sanitized_input("\nWould you like to export the results to a CSV file? (y/n): ")
    user_input = user_input.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)

        await run_sync_task_in_thread(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.")

async 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_df = await api.policy_find_all()
    policies = [Policy(**row.to_dict()) for _, row in policies_df.iterrows()]

    device_input_lines = []
    empty_line_count = 0
    valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$')

    while True:
        line = await get_sanitized_input("")
        stripped_line = line.strip()
        if stripped_line == "":
            empty_line_count += 1
            if empty_line_count == 2:
                break
            continue
        else:
            empty_line_count = 0
            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 []

    pattern = "\n".join(map(re.escape, device_names))
    regex = re.compile(pattern, re.IGNORECASE)

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

    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"))

    for agent in matched_agents:
        agent.enrich_with_policies(policies)

    return matched_agents

async def moveAgentToRelatedPolicy(api: AirlockAPIWrapper, agent: Agent, mode: str = "audit"):
    policy_relationship_map = await get_protected_json("POLICY_MAP_ENF_AUD", "{}")
    #TODO - Have this return the policy name it was moved to instead of the groupid
    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

    result = await api.agent_move(agent.agentid, target_policy)
    if result == {'error': 'Success'}: logger.info(f"{agent.hostname} has been moved to {target_policy}")
    else: logger.debug(result)