<?php

// Chains/Chain.php

namespace Kambo\Langchain\Chains;

use Kambo\Langchain\Callbacks\BaseCallbackManager;
use Kambo\Langchain\Callbacks\StdOutCallbackHandler;
use Kambo\Langchain\Memory\BaseMemory;
use Kambo\Langchain\Callbacks\CallbackManager;
use Kambo\Langchain\Exceptions\NotImplemented;
use Kambo\Langchain\Exceptions\LogicException;
use InvalidArgumentException;
use Exception;

use function array_diff;
use function array_keys;
use function sprintf;
use function implode;
use function get_class;
use function is_array;
use function count;
use function array_merge;
use function array_map;
use function is_string;
use function dirname;
use function file_exists;
use function mkdir;
use function pathinfo;
use function file_put_contents;
use function json_encode;

use const PATHINFO_EXTENSION;
use const JSON_PRETTY_PRINT;

/**
 * Base interface that all chains should implement.
 */
abstract class Chain
{
    protected ?BaseMemory $memory = null;
    protected BaseCallbackManager $callbackManager;

    protected bool $verbose = false;

    /**
     * @param ?BaseMemory          $memory
     * @param ?BaseCallbackManager $callbackManager
     * @param ?bool                $verbose
     */
    public function __construct(
        ?BaseMemory $memory = null,
        ?BaseCallbackManager $callbackManager = null,
        ?bool $verbose = null
    ) {
        $this->memory = $memory;
        $this->callbackManager = $callbackManager ?? new CallbackManager(
            [new StdOutCallbackHandler() ]
        );

        $this->verbose = $verbose ?? $this->verbose;
    }

    public function getChainType(): string
    {
        throw new NotImplemented('Saving not supported for this chain type.');
    }

    public function __get(string $name)
    {
        return $this->$name;
    }

    /**
     * Input keys this chain expects.
     *
     * @return array
     */
    abstract public function inputKeys(): array;

    /**
     * Output keys this chain expects.
     *
     * @return array
     */
    abstract public function outputKeys(): array;

    /**
     * Check that all inputs are present.
     *
     * @param array $inputs
     *
     * @return void
     */
    protected function validateInputs(array $inputs): void
    {
        $missingKeys = array_diff($this->inputKeys(), array_keys($inputs));
        if (!empty($missingKeys)) {
            throw new InvalidArgumentException(sprintf(
                'Missing some input keys: %s',
                implode(', ', $missingKeys)
            ));
        }
    }

    /**
     * Check that all outputs are present.
     * @param array $outputs
     *
     * @return void
     */
    protected function validateOutputs(array $outputs): void
    {
        if (array_diff($this->outputKeys(), array_keys($outputs))) {
            throw new InvalidArgumentException(sprintf(
                'Did not get output keys that were expected. Got: %s. Expected: %s.',
                implode(', ', array_keys($outputs)),
                implode(', ', $this->outputKeys())
            ));
        }
    }

    /**
     * Run the logic of this chain and return the output.
     *
     * @param array $inputs
     *
     * @return array
     */
    abstract protected function call(array $inputs): array;

    /**
     * Run the logic of this chain and add to output if desired.
     *
     * @param mixed $inputs             Dictionary of inputs, or single input if chain expects only one param.
     * @param bool  $returnOnlyOutputs  boolean for whether to return only outputs in the response.
     *                                  If True, only new keys generated by this chain will be returned.
     *                                  If False, both input keys and new keys generated by this chain will be returned.
     *                                  Defaults to False.
     *
     * @return array
     * @throws Exception
     */
    public function __invoke($inputs, bool $returnOnlyOutputs = false): array
    {
        $inputs = $this->prepInputs($inputs);
        $this->callbackManager->onChainStart(
            ['name' => get_class($this)],
            $inputs,
            ['verbose' => $this->verbose]
        );

        try {
            $outputs = $this->call($inputs);
        } catch (Exception $e) {
            $this->callbackManager->onChainError($e, ['verbose' => $this->verbose]);
            throw $e;
        }

        $this->callbackManager->onChainEnd($outputs, ['verbose' => $this->verbose]);
        return $this->prepOutputs($inputs, $outputs, $returnOnlyOutputs);
    }


    /**
     * Validate and prep inputs.
     *
     * @param mixed $inputs
     *
     * @return array
     */
    protected function prepInputs(mixed $inputs): array
    {
        if (!is_array($inputs)) {
            $inputKeys = $this->inputKeys();

            if ($this->memory !== null) {
                $inputKeys = array_diff(
                    $inputKeys,
                    $this->memory->getMemoryVariables()
                );
            }

            if (count($inputKeys) !== 1) {
                throw new InvalidArgumentException(
                    sprintf(
                        'A single string input was passed in, but this chain expects '
                        . 'multiple inputs (%s). When a chain expects '
                        . 'multiple inputs, please call it by passing in an array, '
                        . "eg `chain(['foo' => 1, 'bar' => 2])`",
                        implode(', ', $inputKeys)
                    )
                );
            }

            $inputs = [$inputKeys[0] => $inputs];
        }

        if ($this->memory !== null) {
            $externalContext = $this->memory->loadMemoryVariables($inputs);
            $inputs = array_merge($inputs, $externalContext);
        }

        $this->validateInputs($inputs);
        return $inputs;
    }

    /**
     * Validate and prep outputs.
     *
     * @param array $inputs
     * @param array $outputs
     * @param bool  $returnOnlyOutputs
     *
     * @return array
     */
    protected function prepOutputs(array $inputs, array $outputs, bool $returnOnlyOutputs = false): array
    {
        $this->validateOutputs($outputs);

        if ($this->memory !== null) {
            $this->memory->saveContext($inputs, $outputs);
        }

        if ($returnOnlyOutputs) {
            return $outputs;
        } else {
            return array_merge($inputs, $outputs);
        }
    }

    /**
     * Call the chain on all inputs in the list.
     *
     * @param array $inputList
     *
     * @return array
     */
    public function apply(array $inputList): array
    {
        return array_map([$this, '__invoke'], $inputList);
    }

    /**
     * Run the chain as text in, text out or multiple variables, text out.
     *
     * @param array|string $text
     * @param array ...$args
     *
     * @return mixed
     * @throws LogicException
     */
    public function run(array|string $text, ...$args)
    {
        if (count($this->outputKeys()) !== 1) {
            throw new LogicException('run not supported when there is not exactly one output key.');
        }

        if ($args !== []) {
            throw new InvalidArgumentException('run supports only one positional argument.');
        }

        if (is_string($text)) {
            $inputs = [$this->inputKeys()[0] => $text];
        } elseif (is_array($text)) {
            $inputs = $text;
        } else {
            throw new InvalidArgumentException('Invalid input format.');
        }

        $outputs = $this($inputs);

        return $outputs[$this->outputKeys()[0]];
    }

    /**
     * Convert the chain to an array.
     * @return array
     */
    abstract protected function toArray(): array;

    /**
     * Save the chain.
     *
     * @param string $filePath Path to file to save the chain to.
     *
     * @return void
     */
    public function save(string $filePath): void
    {
        $chainDict = $this->toArray();
        $directoryPath = dirname($filePath);

        if (!file_exists($directoryPath)) {
            mkdir($directoryPath, 0777, true);
        }

        if (pathinfo($filePath, PATHINFO_EXTENSION) === 'json') {
            file_put_contents($filePath, json_encode($chainDict, JSON_PRETTY_PRINT));
        } elseif (pathinfo($filePath, PATHINFO_EXTENSION) === 'yaml') {
            $yaml = Yaml::dump($chainDict, 4);
            file_put_contents($filePath, $yaml);
        } else {
            throw new InvalidArgumentException('File format not supported.');
        }
    }
}