#!/usr/bin/python3

import json
import datetime
import uvicorn
from fastapi import FastAPI, BackgroundTasks, HTTPException, Depends, Header
from pydantic import BaseModel
import paramiko
import yaml
import concurrent.futures
import re
import subprocess
import logging
import select
import os
import pwd
import pam
import secrets
import sdnotify
from contextlib import asynccontextmanager
import sqlite3
from pathlib import Path
from packaging.version import Version


# --- Configuration ---
SSH_USER = 'root'
SSH_KEY = '/root/.ssh/id_rsa'
NGCP_DATA = '/ngcp-data/ngcp-upgrade'
NGCP_UPGRADE_API_SOCKET = '/run/ngcp-upgrade-api/ngcp-upgrade-api.sock'
STORAGE_FILE = f"{NGCP_DATA}/upgrade.db"
STORAGE = None

CHECK_FILE_NAME = "/usr/share/ngcp-upgrade/templates/check.template"
STAGE1_FILE_NAME = "/usr/share/ngcp-upgrade/templates/first_stage.template"
STAGE2_FILE_NAME = "/usr/share/ngcp-upgrade/templates/second_stage.template"
CONF_FILE_NAME = "/usr/share/ngcp-upgrade/conf/configuration"

SESSION_DB = {}
SESSION_TIMEOUT_MINUTES = 60

RUNNING_SCENARIO = None
STOP_UPGRADE = False


class CreateScenarioRequest(BaseModel):
    release_from: str
    release_to: str


class RunScenarioRequest(BaseModel):
    scenario_name: str
    force: str = "false"


class GetStepRequest(BaseModel):
    scenario_name: str
    step_name: str = "all"
    node_name: str = "localhost"


class LoginSchema(BaseModel):
    username: str
    password: str


@asynccontextmanager
async def lifespan(app: FastAPI):
    n.notify("READY=1")
    yield
    n.notify("STOPPING=1")


n = sdnotify.SystemdNotifier()

app = FastAPI(lifespan=lifespan)
pam_check = pam.pam()


def init_storage():
    root_uid = pwd.getpwnam("root").pw_uid
    root_gid = pwd.getpwnam("root").pw_gid
    Path(NGCP_DATA).mkdir(parents=True, exist_ok=True)
    os.chown(NGCP_DATA, root_uid, root_gid)
    Path(STORAGE_FILE).touch(mode=0o644, exist_ok=True)
    os.chown(STORAGE_FILE, root_uid, root_gid)

    global STORAGE
    STORAGE = sqlite3.connect(STORAGE_FILE, timeout=10,
                              check_same_thread=False)
    STORAGE.execute("""
                CREATE TABLE IF NOT EXISTS ngcp_upgrade
                (scenario_name VARCHAR(255) NOT NULL PRIMARY KEY,
                status VARCHAR(255) NOT NULL,
                scenario JSON NOT NULL)
            """)
    STORAGE.commit()


def read_from_storage(scenario_name):
    cursor = STORAGE.execute(
                "SELECT scenario FROM ngcp_upgrade WHERE scenario_name = ?",
                (scenario_name,)
            )
    row = cursor.fetchone()
    if row:
        try:
            return json.loads(row[0])
        except json.JSONDecodeError:
            return {}

    return {}


def write_to_storage(scenario_name, data):
    STORAGE.execute("INSERT OR REPLACE INTO ngcp_upgrade "
                    "(scenario_name, scenario, status) "
                    " VALUES (?, ?, ?)",
                    (scenario_name, json.dumps(data), data['status']))
    STORAGE.commit()


def read_list_from_storage():
    cursor = STORAGE.execute("SELECT scenario_name, status FROM ngcp_upgrade")
    rows = cursor.fetchall()
    if rows:
        return [{row[0]: row[1]} for row in rows]
    else:
        return []


def delete_from_storage(scenario_name):
    STORAGE.execute("DELETE FROM ngcp_upgrade WHERE scenario_name = ?",
                    (scenario_name,))
    STORAGE.commit()


def setup_logger(name, log_file, level=logging.INFO):
    handler = logging.FileHandler(log_file)
    formatter = logging.Formatter('%(message)s')
    handler.setFormatter(formatter)

    logger = logging.getLogger(name)
    logger.setLevel(level)
    logger.addHandler(handler)

    return logger


def worker(
        server_ip, command, from_version, to_version,
        scenario_name, step_name, force):
    argument = ''
    chroot_pattern = r'^(.+) (chroot)$'
    chroot_match = re.search(chroot_pattern, command)
    if chroot_match:
        command = chroot_match.group(1)
        argument = chroot_match.group(2)

    os.makedirs(f"{NGCP_DATA}/{scenario_name}/{server_ip}", exist_ok=True)

    logger = setup_logger(f"{step_name}_{server_ip}",
                          f"{NGCP_DATA}/{scenario_name}/"
                          f"{server_ip}/{step_name}")

    wrap_cmd = f"source {CONF_FILE_NAME} && "
    wrap_cmd += f"OLD_VERSION={from_version} "
    wrap_cmd += f"UPGRADE_VERSION={to_version} "
    wrap_cmd += f"FORCE_UPGRADE={force} "
    wrap_cmd += f"LOGFILE=/tmp/logfile "
    wrap_cmd += f"SCENARIO_NAME={scenario_name} "
    wrap_cmd += f"NGCP_UPGRADE=ngcp-upgrade {command} {argument}"

    exit_status = 1
    request_confirmation = False
    main_logger.info(f"Running {command} on {server_ip}")
    if EDITION == 'CE':
        try:
            process = subprocess.Popen(
                wrap_cmd,
                shell=True,
                executable="/bin/bash",
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1
            )
            inputs = [process.stdout, process.stderr]
            while process.poll() is None:  # While process is still running
                readable, _, _ = select.select(inputs, [], [], 0.1)
                for pipe in readable:
                    line = pipe.readline()
                    if line:
                        if pipe is process.stdout:
                            if "REQUEST_CONFIRMATION" in line:
                                request_confirmation = True
                            logger.info(f"STDOUT: {line.strip()}")
                        else:
                            logger.error(f"STDERR: {line.strip()}")

            # Read any remaining data after the process exits
            for line in process.stdout:
                logger.info(f"STDOUT: {line.strip()}")
            for line in process.stderr:
                logger.error(f"STDERR: {line.strip()}")
            exit_status = process.returncode
        except subprocess.FileNotFoundError as e:
            logger.info(e.stdout)
            logger.info(e.stderr)
            exit_status = 1
    else:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        try:
            ssh.connect(server_ip, username=SSH_USER,
                        key_filename=SSH_KEY, timeout=10)
            transport = ssh.get_transport()
            channel = transport.open_session()
            channel.exec_command(wrap_cmd)

            while True:
                # Check if stdout or stderr have data ready to read
                # timeout = 0.1 prevents the CPU from spinning too fast
                read_ready, _, _ = select.select([channel], [], [], 0.1)

                if read_ready:
                    if channel.recv_ready():
                        line = channel.recv(4096).decode('utf-8', 'replace')
                        if line:
                            if "REQUEST_CONFIRMATION" in line:
                                request_confirmation = True
                            logger.info(f"STDOUT: {line.strip()}")

                    if channel.recv_stderr_ready():
                        line = channel.recv_stderr(
                            4096).decode('utf-8', 'replace')
                        if line:
                            logger.error(f"STDERR: {line.strip()}")

                if (channel.exit_status_ready() and not
                        channel.recv_ready() and not
                        channel.recv_stderr_ready()):
                    break

            exit_status = channel.recv_exit_status()
        except Exception as e:
            main_logger.error(
                f"Exception from {server_ip} step {step_name}: {str(e)}")
        finally:
            ssh.close()

    main_logger.info(
        f"Command {command} on {server_ip} with status {exit_status}")

    if request_confirmation:
        # This is impossible exit code
        # so we can differentiate it from some error
        exit_status = 1000
    return exit_status


def read_step_log(scenario_name, step_name, node_name):
    file_path = f"{NGCP_DATA}/{scenario_name}/{node_name}/{step_name}"
    if Path(file_path).is_file():
        with open(file_path, 'r') as file:
            return file.read()
    else:
        return (f"There is no log for step {step_name} on node {node_name},"
                " maybe it wasn't run yet")


def lookahead(iterable):
    it = iter(iterable)
    try:
        last = next(it)
    except StopIteration:
        return
    for val in it:
        yield last, False
        last = val
    yield last, True


def execute_scenario(data: RunScenarioRequest):
    scenario_name = data.scenario_name
    force = data.force

    vers_pattern = r'^(mr.+?)_(mr.+?|trunk)_'
    vers_match = re.search(vers_pattern, scenario_name)
    from_version = vers_match.group(1)
    to_version = vers_match.group(2)

    scenario = read_from_storage(scenario_name)
    scenario['status'] = 'running'
    write_to_storage(scenario_name, scenario)

    result = os.makedirs(f"{NGCP_DATA}/{scenario_name}", exist_ok=True)

    global STOP_UPGRADE
    global RUNNING_SCENARIO
    for step_info, last_step in lookahead(scenario['steps']):
        for key in step_info:
            step_name = key
        if step_info[step_name]['status'] == 'success':
            continue
        cmd = step_info[step_name]['command']
        stage = step_info[step_name]['stage']
        executors = {}
        failed_step = False
        request_confirmation = False
        if STOP_UPGRADE:
            main_logger.info(f"Don't run {step_name} as upgrade was stopped")
            scenario['status'] = 'pending'
            write_to_storage(scenario_name, scenario)
            break
        main_logger.info(f"Running step {step_name}...")
        step_info[step_name]['status'] = 'running'
        write_to_storage(scenario_name, scenario)
        with concurrent.futures.ThreadPoolExecutor() as executor:
            for node in step_info[step_name]['nodes'].keys():
                if step_info[step_name]['nodes'][node]['status'] == 'success':
                    continue
                elif step_info[step_name]['nodes'][node]['status'] == \
                        'request_confirmation':
                    step_info[step_name]['nodes'][node]['status'] = 'success'
                    write_to_storage(scenario_name, scenario)
                    continue
                runner = executor.submit(
                    worker, node, cmd, from_version,
                    to_version, scenario_name, step_name, force)
                executors[runner] = node
                step_info[step_name]['nodes'][node]['status'] = 'running'
                write_to_storage(scenario_name, scenario)
            for runner in concurrent.futures.as_completed(executors):
                node = executors[runner]
                try:
                    result = runner.result()
                    if result == 0:
                        main_logger.info(
                            f"Node {node} completed successfully.")
                        step_info[step_name]['nodes'][node]['status'] = \
                            'success'
                        write_to_storage(scenario_name, scenario)
                    elif result == 1000:
                        main_logger.info(
                            f"Node {node} requires attention.")
                        step_info[step_name]['nodes'][node]['status'] = \
                            'request_confirmation'
                        request_confirmation = True
                        write_to_storage(scenario_name, scenario)
                    else:
                        failed_step = True
                        main_logger.info(f"Node {node} failed")
                        step_info[step_name]['nodes'][node]['status'] = 'fail'
                        write_to_storage(scenario_name, scenario)
                except Exception as exc:
                    failed_step = True
                    main_logger.error(
                        f"Node {node} generated an exception: {exc}")
                    step_info[step_name]['nodes'][node]['status'] = 'fail'
                    write_to_storage(scenario_name, scenario)
        if failed_step:
            RUNNING_SCENARIO = None
            main_logger.error(f"Step {step_name} failed")
            step_info[step_name]['status'] = 'fail'
            scenario['status'] = 'fail'
            write_to_storage(scenario_name, scenario)
            break
        elif request_confirmation:
            step_info[step_name]['status'] = 'request_confirmation'
            write_to_storage(scenario_name, scenario)
            if stage == 'check':
                main_logger.info(f"Step {step_name} requires confirmation")
                main_logger.info(f"But this is 'check' stage, so continuing")
            else:
                RUNNING_SCENARIO = None
                scenario['status'] = 'pending'
                write_to_storage(scenario_name, scenario)
                break
        else:
            step_info[step_name]['status'] = 'success'
            write_to_storage(scenario_name, scenario)

        if step_name in ['stop', 'report_check_stop', 'report_upgrade_stop']:
            RUNNING_SCENARIO = None
            scenario['status'] = 'pending'
            if step_name == 'report_upgrade_stop' and last_step:
                scenario['status'] = 'success'
            write_to_storage(scenario_name, scenario)
            break


def assign_node_steps(steps, nodes, payload):
    for step_info in steps:
        step = step_info['step']
        s_nodes = step_info['nodes']
        stage = step_info['stage']
        nodes_struct = {}
        if s_nodes == 'default':
            for node in nodes:
                nodes_struct[node] = {'status': 'pending'}
        else:
            for node in nodes:
                if node in s_nodes:
                    nodes_struct[node] = {'status': 'pending'}

        command = f"/usr/share/ngcp-upgrade/steps/{step}"
        step_name = re.sub(r'^.+/', '', step)
        step_struct = {step_name: {'status': 'pending',
                                   'command': command, 'nodes': nodes_struct,
                                   'stage': stage}}
        payload['steps'].append(step_struct)


async def verify_session(x_session_token: str = Header(None)):
    if not x_session_token or x_session_token not in SESSION_DB:
        raise HTTPException(status_code=401, detail="Unauthorized")
    session = SESSION_DB[x_session_token]
    if datetime.datetime.now() > session["expires"]:
        del SESSION_DB[x_session_token]
        raise HTTPException(status_code=401, detail="Session expired")

    session["expires"] = datetime.datetime.now() + \
        datetime.timedelta(minutes=SESSION_TIMEOUT_MINUTES)


def add_temp_step(from_rel: str, to_rel: str, cross_rels: str) -> bool:
    crossing = False
    from_ver = Version(from_rel.replace("mr", ""))
    to_ver = Version(to_rel.replace("mr", ""))
    releases = cross_rels.split(",")
    for release in releases:
        cross_ver = Version(release.replace("mr", ""))
        if from_ver < cross_ver < to_ver:
            crossing = True
            break

    return crossing


@app.post("/login")
async def login(request: LoginSchema):
    if pam_check.authenticate(request.username, request.password):
        token = secrets.token_hex(32)
        expires = datetime.datetime.now() + \
            datetime.timedelta(minutes=SESSION_TIMEOUT_MINUTES)
        SESSION_DB[token] = {
            "user": request.username,
            "expires": expires
        }
        return {"token": token}
    else:
        raise HTTPException(
            status_code=403,
            detail="Invalid system credentials"
        )


@app.post("/create_scenario", dependencies=[Depends(verify_session)])
async def create_scenario(request: CreateScenarioRequest):
    from_rel = request.release_from
    to_rel = request.release_to

    os.makedirs("/ngcp-data/ngcp-upgrade", exist_ok=True)

    current_time = datetime.datetime.now().strftime("%Y-%m-%d_%H_%M_%S")

    steps = []
    # This is a special step to install new version of ngcp-upgrade
    # package to all the systems
    steps.append(
        {'step': 'run_prepare_upgrade', 'nodes': 'default', 'stage': 'check'})
    second_stage_steps = []
    for template_file in CHECK_FILE_NAME, STAGE1_FILE_NAME, STAGE2_FILE_NAME:
        stage = 'check'
        if template_file == STAGE1_FILE_NAME:
            stage = 'first_stage'
        elif template_file == STAGE2_FILE_NAME:
            stage = 'second_stage'
        try:
            with open(template_file, 'r') as file:
                for line in file:
                    if line.startswith("#"):
                        continue
                    commands = line.split("|")
                    step = commands[0].strip()
                    try:
                        cross_rel = commands[1].strip()
                    except IndexError:
                        cross_rel = None
                    try:
                        edition_specific = commands[2].strip()
                    except IndexError:
                        edition_specific = None
                    try:
                        nodes = 'default'
                        templ_nodes = tuple(
                            item.strip() for item in commands[3].split(','))
                        matched_nodes = [
                            node for node in network_schema['hosts'] if
                            node.startswith(templ_nodes)]
                        nodes = ",".join(matched_nodes)
                    except IndexError:
                        nodes = 'default'

                    if not step:
                        continue
                    if cross_rel:
                        if not add_temp_step(from_rel, to_rel, cross_rel):
                            continue
                    if edition_specific:
                        if not edition_specific == EDITION:
                            continue

                    if template_file == STAGE2_FILE_NAME:
                        second_stage_steps.append(
                            {'step': step, 'nodes': nodes, 'stage': stage})
                    else:
                        steps.append(
                            {'step': step, 'nodes': nodes, 'stage': stage})
        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Can't read file {template_file}: {e}"
            )

    scenario_key = f"{from_rel}_{to_rel}_{current_time}"
    payload = {
        'steps': [],
        'status': 'pending'
    }

    if EDITION == 'CE':
        assign_node_steps(steps, ['localhost'], payload)
        assign_node_steps(second_stage_steps, ['localhost'], payload)
    elif EDITION == 'PRO':
        assign_node_steps(steps, network_schema['hosts'], payload)
        assign_node_steps(second_stage_steps, ['sp1'], payload)
        assign_node_steps(second_stage_steps, ['sp2'], payload)
    elif EDITION == 'CARRIER':
        assign_node_steps(steps, network_schema['hosts'], payload)
        assign_node_steps(second_stage_steps, ['web01a'], payload)
        assign_node_steps(second_stage_steps, ['db01a'], payload)
        a_side = []
        for node in network_schema['hosts']:
            if (node.endswith('a') and node not in ('web01a', 'db01a')):
                a_side.append(node)
        assign_node_steps(second_stage_steps, a_side, payload)
        b_side = []
        for node in network_schema['hosts']:
            if node.endswith('b'):
                b_side.append(node)
        assign_node_steps(second_stage_steps, b_side, payload)

    # Special processing of the step - specify which version of package
    # should be installed
    payload['steps'][0]['run_prepare_upgrade']['command'] = (
        f"ngcp-prepare-upgrade {to_rel}"
    )

    try:
        write_to_storage(scenario_key, payload)
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Can't write scenario to storage: {e}"
        )

    return {"scenario_name": scenario_key}


@app.post("/run_scenario", dependencies=[Depends(verify_session)])
async def run_scenario(request: RunScenarioRequest,
                       background_tasks: BackgroundTasks):
    scenario_name = request.scenario_name
    scenario = read_from_storage(scenario_name)

    if scenario is None or not scenario:
        main_logger.error(f"Scenario '{scenario_name}' does not exist.")
        raise HTTPException(
            status_code=500,
            detail={f"Scenario {scenario_name} does not exists."}
        )

    if scenario['status'] == 'running':
        raise HTTPException(
            status_code=500,
            detail=f"Scenario {request.scenario_name} already running"
        )

    global RUNNING_SCENARIO
    if RUNNING_SCENARIO:
        raise HTTPException(
            status_code=500,
            detail=f"Scenario {RUNNING_SCENARIO} already running"
        )

    RUNNING_SCENARIO = scenario_name
    global STOP_UPGRADE
    STOP_UPGRADE = False

    background_tasks.add_task(execute_scenario, request)

    return {"scenario": scenario_name, "status": 'running'}


@app.post("/get_scenario_status", dependencies=[Depends(verify_session)])
async def get_scenario_status(request: RunScenarioRequest):
    scenario_name = request.scenario_name

    scenario = read_from_storage(scenario_name)
    if scenario is None:
        main_logger.error(f"Scenario '{scenario_name}' does not exist.")
        return {f"Scenario {scenario_name} does not exists."}

    return {"scenario": scenario_name, "status": scenario['status']}


@app.post("/get_scenarios_list", dependencies=[Depends(verify_session)])
async def get_scenarios_list():
    scenarios = read_list_from_storage()

    return {'scenarios': scenarios}


@app.post("/delete_scenario", dependencies=[Depends(verify_session)])
async def delete_scenario(request: RunScenarioRequest):
    scenario_name = request.scenario_name
    delete_from_storage(scenario_name)

    return {"scenario": scenario_name, "status": 'deleted'}


@app.post("/get_step_status", dependencies=[Depends(verify_session)])
async def get_step_status(request: GetStepRequest):
    scenario_name = request.scenario_name
    step_name = request.step_name

    scenario = read_from_storage(scenario_name)
    if scenario is None:
        main_logger.error(f"Scenario '{scenario_name}' does not exist.")
        return {f"Scenario {scenario_name} does not exists."}

    if step_name == "all":
        output = scenario.get("steps")
    else:
        steps = scenario.get("steps", [])
        step_dict = next((step for step in steps if step_name in step), None)
        if not step_dict:
            err = f"Step '{step_name}' not found "
            err.join(f"in scenario '{scenario_name}'.")
            main_logger.error(err)
            return {err}
        output = step_dict

    return output


@app.post("/stop", dependencies=[Depends(verify_session)])
async def stop():
    global STOP_UPGRADE
    STOP_UPGRADE = True

    global RUNNING_SCENARIO
    if not RUNNING_SCENARIO:
        return {"detail": "No scenario is running"}

    scenario_name = RUNNING_SCENARIO
    RUNNING_SCENARIO = None

    return {"detail": f"Scenario {scenario_name} stopped"}


@app.post("/get_step_log", dependencies=[Depends(verify_session)])
async def get_step_log(request: GetStepRequest):
    scenario_name = request.scenario_name
    step_name = request.step_name
    node_name = request.node_name
    if EDITION == 'PRO' and node_name == 'localhost':
        node_name = 'sp1'
    elif EDITION == 'CARRIER' and node_name == 'localhost':
        node_name = 'web01a'

    if step_name == 'all':
        main_logger.error(f"Step in /get_step_log does not defined.")
        return {"detail": f"You need to specify step_name in /get_step_log"}

    scenario = read_from_storage(scenario_name)
    if scenario is None:
        main_logger.error(f"Scenario '{scenario_name}' does not exist.")
        return {f"Scenario {scenario_name} does not exists."}

    data = read_step_log(scenario_name, step_name, node_name)

    return {"detail": data}


if __name__ == "__main__":

    main_logger = setup_logger(
        "ngcp-upgrade-api",
        "/var/log/ngcp/ngcp-upgrade-api.log")

    with open('/etc/ngcp-config/network.yml', 'r') as file:
        network_schema = yaml.safe_load(file)
    EDITION = 'CE'
    if 'self' not in network_schema["hosts"]:
        EDITION = 'PRO'
        if 'web01a' in network_schema["hosts"]:
            EDITION = 'CARRIER'

    init_storage()

    uvicorn.run(app, uds=NGCP_UPGRADE_API_SOCKET)
