import http.server
import socketserver
import os
import io
import re
import json
from datetime import datetime
from importlib import resources

from sentipede.template import render_listing
from sentipede.ui import (
    gradient_text,
    colorize_status,
    colored_text,
    print_report_box,
    COLOR_1,
    COLOR_2,
    COLOR_3,
    STATUS_ERROR,
    STATUS_WARNING,
)
from sentipede.api import API


class StyledHandler(http.server.SimpleHTTPRequestHandler):
    server_root_path = None
    api: API = None
    received_filename = None

    def log_message(self, format, *args):
        timestamp = datetime.now().strftime("%H:%M:%S")
        colored_timestamp = colored_text(timestamp, COLOR_1)

        status = str(args[1]) if len(args) > 1 else "-"
        colored_status = colorize_status(int(status) if status.isdigit() else status)

        client_ip, client_port = self.client_address
        method = getattr(self, "command", "?")
        path = getattr(self, "path", "?")

        log_line = (
            f"[{colored_timestamp}] [{client_ip}] {colored_status} "
            f"{colored_text(method, COLOR_2)} {path} "
        )

        is_external = client_ip not in ['127.0.0.1', 'localhost', '::1']
        is_file_download = (
            method == "GET" and
            status == "200" and
            not path.startswith('/templates/') and
            path not in ['/', '/api/interfaces'] and
            not path.endswith('/')
        )

        if is_external and is_file_download:
            log_line += gradient_text("successfully sent")

        if self.received_filename:
            log_line += gradient_text(f"✓ File received: {self.received_filename}")
            self.received_filename = None

        print(log_line)

    def do_GET(self):
        if self.path == '/favicon.ico':
            self.send_response(204)
            self.end_headers()
            return

        if self.path.startswith('/templates/'):
            filename = self.path[len('/templates/'):]
            try:
                ref = resources.files("sentipede.templates").joinpath(filename)
                content = ref.read_bytes()
            except Exception:
                self.send_error(404, "File not found")
                return

            if filename.endswith('.css'):
                content_type = 'text/css'
            elif filename.endswith('.js'):
                content_type = 'application/javascript'
            elif filename.endswith('.json'):
                content_type = 'application/json'
            else:
                content_type = 'text/plain'

            self.send_response(200)
            self.send_header('Content-type', content_type)
            self.send_header('Content-Length', str(len(content)))
            self.end_headers()
            self.wfile.write(content)
            return

        if self.path.startswith('/proxy/'):
            target_url = self.path[len('/proxy/'):]
            if not target_url.startswith('http'):
                target_url = 'http://' + target_url

            print(f"[{colored_text(datetime.now().strftime('%H:%M:%S'), COLOR_1)}] {colored_text('PROXY', COLOR_3)} {target_url}")

            stream = self.api.proxy_stream(target_url)
            _, c_type, c_length, error = next(stream)

            if error:
                print(f"[{colored_text('!', STATUS_ERROR)}] Proxy failed: {error}")
                self.send_error(502, f"Bad Gateway: {error}")
                return

            filename = target_url.split('/')[-1] or 'downloaded_file'
            self.send_response(200)
            self.send_header('Content-Type', c_type)
            if c_length is not None:
                self.send_header('Content-Length', c_length)
            self.send_header('Content-Disposition', f'inline; filename="{filename}"')
            self.end_headers()

            try:
                for chunk in stream:
                    self.wfile.write(chunk)
            except (BrokenPipeError, ConnectionResetError):
                pass
            return

        if self.path == '/api/interfaces':
            interfaces = self.api.get_all_interfaces()
            encoded = json.dumps({'interfaces': interfaces}).encode('utf-8')
            self.send_response(200)
            self.send_header("Content-type", "application/json")
            self.end_headers()
            self.wfile.write(encoded)
        else:
            try:
                super().do_GET()
            except BrokenPipeError:
                pass

    def do_POST(self):
        if self.path.startswith('/upload'):
            try:
                content_type = self.headers.get('content-type', '')
                content_length = int(self.headers.get('Content-Length', 0))

                if 'multipart/form-data' in content_type:
                    boundary = re.search(r'boundary=([^\s;]+)', content_type)
                    if not boundary:
                        self.send_error(400, "Missing boundary in multipart")
                        return

                    boundary_bytes = boundary.group(1).encode()
                    body = self.rfile.read(content_length)

                    # Split on boundary, skip preamble and epilogue
                    parts = body.split(b'--' + boundary_bytes)
                    file_data = None
                    orig_filename = "uploaded_file"

                    for part in parts[1:]:  # skip preamble
                        if part in (b'--\r\n', b'--'):
                            break
                        # Split headers from body
                        if b'\r\n\r\n' not in part:
                            continue
                        headers_raw, _, file_body = part.partition(b'\r\n\r\n')
                        headers_str = headers_raw.decode(errors='replace')

                        # Only handle the 'file' field
                        if 'name="file"' not in headers_str:
                            continue

                        # Extract filename from Content-Disposition
                        fn_match = re.search(r'filename="([^"]+)"', headers_str)
                        if fn_match:
                            orig_filename = fn_match.group(1)

                        # Strip trailing \r\n added by multipart
                        file_data = file_body.rstrip(b'\r\n')
                        break

                    if file_data is not None:
                        filename = self.api.save_upload_raw(
                            data=file_data,
                            filename=orig_filename,
                            dest_dir='.',
                        )
                        self.received_filename = filename
                        self.send_response(200)
                        self.send_header('Content-type', 'text/plain')
                        self.end_headers()
                        self.wfile.write(b'OK')
                    else:
                        self.send_error(400, "No file in request")
                else:
                    filename = self.api.parse_upload_filename_from_url(self.path, default='uploaded_file')
                    data = self.rfile.read(content_length)
                    filename = self.api.save_upload_raw(data=data, filename=filename, dest_dir='.')
                    self.received_filename = filename
                    self.send_response(200)
                    self.send_header('Content-type', 'text/plain')
                    self.end_headers()
                    self.wfile.write(b'OK')

            except Exception as e:
                timestamp = datetime.now().strftime("%H:%M:%S")
                print(f"[{colored_text(timestamp, COLOR_1)}] {colorize_status(500)} {colored_text('ERROR', STATUS_ERROR)} {str(e)}")
                self.send_error(500)
        else:
            self.send_error(404, "Endpoint not found")

    def list_directory(self, path):
        files_data = self.api.list_directory_for_ui(path)
        if files_data is None:
            self.send_error(404, "No permission")
            return None

        server_address = self.server.server_address
        host = self.headers.get('Host', f"{server_address[0]}:{server_address[1]}")
        base_url = f"http://{host}"

        display_path = self.api.get_display_path(path)
        content = render_listing(display_path, files_data, base_url)

        encoded = content.encode('utf-8')
        f = io.BytesIO()
        f.write(encoded)
        f.seek(0)

        self.send_response(200)
        self.send_header("Content-type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(encoded)))
        self.end_headers()
        return f


class WebExfil:
    def __init__(self, directory=".", port=8000):
        self.directory = os.path.abspath(directory)
        self.port = port
        self.api = API(server_root_path=self.directory)

    def run(self):
        os.chdir(self.directory)
        StyledHandler.server_root_path = self.directory
        StyledHandler.api = self.api

        info_data = {
            "WEB UI": f"http://localhost:{self.port}",
            "ROOT":   self.directory,
            "PID":    os.getpid()
        }
        print_report_box("SENTIPEDE SERVER CONFIG", info_data)
        print("\n" + gradient_text("--- [ ACTIVITY LOGS ] ---") + "\n")

        socketserver.TCPServer.allow_reuse_address = True
        with socketserver.TCPServer(("", self.port), StyledHandler) as httpd:
            try:
                httpd.serve_forever()
            except KeyboardInterrupt:
                print(f"\n\n{colorize_status('!')} {colored_text('Stopping server...', STATUS_WARNING)}")
                httpd.shutdown()