################################################################################
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
import json
import os
import shutil
import sysconfig
import tempfile
from pathlib import Path

from pyflink.common import Configuration, WatermarkStrategy
from pyflink.datastream import (
    KeySelector,
    RuntimeExecutionMode,
    StreamExecutionEnvironment,
)
from pyflink.datastream.connectors.file_system import (
    FileSource,
    StreamFormat,
)

from flink_agents.api.agents.agent import Agent
from flink_agents.api.decorators import action
from flink_agents.api.events.event import Event, InputEvent, OutputEvent
from flink_agents.api.execution_environment import AgentsExecutionEnvironment
from flink_agents.api.runner_context import RunnerContext

os.environ["PYTHONPATH"] = sysconfig.get_paths()["purelib"]


class InputKeySelector(KeySelector):
    """Key selector for input data."""

    def get_key(self, value: dict) -> int:
        """Extract key from input data."""
        return value.get("id", 0)


class PythonEventLoggingAgent(Agent):
    """Agent for testing Python event logging."""

    @action(InputEvent.EVENT_TYPE)
    @staticmethod
    def process_input(event: Event, ctx: RunnerContext) -> None:
        """Process input event and send an output event."""
        input_data = InputEvent.from_event(event).input
        ctx.send_event(
            OutputEvent(output={"processed_review": f"{input_data['review']}"})
        )


def test_python_event_logging(tmp_path: Path) -> None:
    """Test event logs are written to configured directory with expected content."""
    event_log_dir = tmp_path / "event_log"
    default_log_dir = Path(tempfile.gettempdir()) / "flink-agents"
    shutil.rmtree(default_log_dir, ignore_errors=True)

    config = Configuration()
    env = StreamExecutionEnvironment.get_execution_environment(config)
    env.set_runtime_mode(RuntimeExecutionMode.STREAMING)
    env.set_parallelism(1)

    # Create agent environment
    agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env)
    agents_env.get_config().set_str("baseLogDir", str(event_log_dir))

    # Set up input source
    current_dir = Path(__file__).parent
    input_datastream = env.from_source(
        source=FileSource.for_record_stream_format(
            StreamFormat.text_line_format(),
            f"file:///{current_dir}/../resources/input/input_data.txt",
        ).build(),
        watermark_strategy=WatermarkStrategy.no_watermarks(),
        source_name="python_event_logging_test",
    )

    # Parse input
    deserialize_datastream = input_datastream.map(lambda x: json.loads(x))

    # Apply agent
    agents_env.from_datastream(
        input=deserialize_datastream, key_selector=InputKeySelector()
    ).apply(PythonEventLoggingAgent()).to_datastream()

    # Execute the job
    agents_env.execute()

    # Check that log files were created in configured directory
    log_files = list(event_log_dir.glob("events-*.log"))

    # At least one log file should exist
    assert len(log_files) > 0, f"Event log files should be created in {event_log_dir}"

    # Check that log files contain structured event content
    record = None
    record_line = None
    has_processed_review = False
    for log_file in log_files:
        with log_file.open(encoding="utf-8") as handle:
            for line in handle:
                if line.strip():
                    record = json.loads(line)
                    record_line = line
                    event_payload = record.get("event", {})
                    if "processed_review" in json.dumps(event_payload):
                        has_processed_review = True
                        break
        if record is not None and has_processed_review:
            break

    assert record is not None, "Event log file is empty."
    assert record_line is not None, "Event log file is empty."
    assert "timestamp" in record
    assert "logLevel" in record
    assert "eventType" in record
    assert "event" in record
    assert "eventType" in record["event"]
    assert has_processed_review, "Log should contain processed review content"

    event_type_idx = record_line.find('"eventType"')
    id_idx = record_line.find('"id"')
    attributes_idx = record_line.find('"attributes"')
    assert event_type_idx != -1
    assert id_idx != -1
    assert attributes_idx != -1
    assert event_type_idx < id_idx < attributes_idx


def _run_event_logging_pipeline(
    tmp_path: Path, config_overrides: dict[str, str] | None = None
) -> Path:
    """Run the event logging pipeline and return the event log directory.

    Args:
        tmp_path: Temporary directory for log output.
        config_overrides: Optional dict of config key-value pairs to set.

    Returns:
        The event log directory path.
    """
    event_log_dir = tmp_path / "event_log"
    default_log_dir = Path(tempfile.gettempdir()) / "flink-agents"
    shutil.rmtree(default_log_dir, ignore_errors=True)

    config = Configuration()
    env = StreamExecutionEnvironment.get_execution_environment(config)
    env.set_runtime_mode(RuntimeExecutionMode.STREAMING)
    env.set_parallelism(1)

    agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env)
    agents_env.get_config().set_str("baseLogDir", str(event_log_dir))

    if config_overrides:
        for key, value in config_overrides.items():
            agents_env.get_config().set_str(key, value)

    current_dir = Path(__file__).parent
    input_datastream = env.from_source(
        source=FileSource.for_record_stream_format(
            StreamFormat.text_line_format(),
            f"file:///{current_dir}/../resources/input/input_data.txt",
        ).build(),
        watermark_strategy=WatermarkStrategy.no_watermarks(),
        source_name="python_event_logging_test",
    )

    deserialize_datastream = input_datastream.map(lambda x: json.loads(x))

    agents_env.from_datastream(
        input=deserialize_datastream, key_selector=InputKeySelector()
    ).apply(PythonEventLoggingAgent()).to_datastream()

    agents_env.execute()
    return event_log_dir


def _read_log_records(event_log_dir: Path) -> list[dict]:
    """Read all JSON records from event log files.

    Args:
        event_log_dir: Directory containing event log files.

    Returns:
        List of parsed JSON records.
    """
    records: list[dict] = []
    for log_file in event_log_dir.glob("events-*.log"):
        with log_file.open(encoding="utf-8") as handle:
            records.extend(json.loads(line) for line in handle if line.strip())
    return records


def test_event_log_verbose_level(tmp_path: Path) -> None:
    """Test that VERBOSE log level writes events without truncation."""
    event_log_dir = _run_event_logging_pipeline(
        tmp_path, config_overrides={"event-log.level": "VERBOSE"}
    )

    records = _read_log_records(event_log_dir)
    assert len(records) > 0, "VERBOSE level should produce event log records"

    for record in records:
        assert record.get("logLevel") == "VERBOSE", (
            f"Expected logLevel VERBOSE, got {record.get('logLevel')}"
        )

    # VERBOSE should not truncate content (no truncatedString wrappers)
    raw_content = json.dumps(records)
    assert "truncatedString" not in raw_content, (
        "VERBOSE level should not truncate any content"
    )


def test_event_log_off_level(tmp_path: Path) -> None:
    """Test that OFF log level suppresses all event logging."""
    event_log_dir = _run_event_logging_pipeline(
        tmp_path, config_overrides={"event-log.level": "OFF"}
    )

    records = _read_log_records(event_log_dir)
    assert len(records) == 0, (
        f"OFF level should not produce any event log records, but found {len(records)}"
    )


def test_event_log_standard_truncation(tmp_path: Path) -> None:
    """Test that STANDARD level truncates strings exceeding max-string-length."""
    event_log_dir = _run_event_logging_pipeline(
        tmp_path,
        config_overrides={
            "event-log.level": "STANDARD",
            "event-log.standard.max-string-length": "10",
        },
    )

    records = _read_log_records(event_log_dir)
    assert len(records) > 0, "STANDARD level should produce event log records"

    for record in records:
        assert record.get("logLevel") == "STANDARD", (
            f"Expected logLevel STANDARD, got {record.get('logLevel')}"
        )

    # With max-string-length=10, any string longer than 10 chars should be truncated
    raw_content = json.dumps(records)
    assert "truncatedString" in raw_content, (
        "STANDARD level with max-string-length=10 should truncate long strings"
    )
