<?php
namespace boru\boruai\Vector;

use boru\boruai\BoruAI;
use boru\boruai\Models\Embedding;
use boru\output\Output;
use boru\qdrant\Models\Collection;
use boru\qdrant\Models\Point;
use boru\query\Entity;
use boru\query\models\EntityDefinition;
use boru\query\Query;
use Exception;

class VectorCollection {
    
    
    /** @var Collection */
    private $collection;

    private $recordClassName;

    private $hasInit = false;

    public function __construct($initOptions=[]) {
        $this->fromOptions($initOptions);
        if($this->collection) {
            $this->collection->save();
        }
    }

    /**
     * @param Collection|array|null $collection
     * @param array|null $vector
     * @param array|null $options
     * @return Collection
     */
    public function collection($collection=null,$vector=null,$options=[]) {
        if ($collection !== null) {
            if($collection instanceof Collection) {
                $this->collection = $collection;
            } else {
                $this->collection = new Collection($collection,$vector,$options);
            }
            $this->collection = $collection;
            return $this;
        }
        return $this->collection;
    }
    public function name() {
        if($this->collection) {
            return $this->collection->name();
        }
        return null;
    }
    public function recordClassName($recordClassOrObject=null) {
        if($recordClassOrObject instanceof CollectionRecord) {
            $this->recordClassName = get_class($recordClassOrObject);
        } elseif(is_string($recordClassOrObject)) {
            $this->recordClassName = $recordClassOrObject;
        }
        if($this->recordClassName === null) {
            $this->recordClassName = CollectionRecord::class;
        }
        return $this->recordClassName;
    }

    /**
     * Initializes the collection if it has not been initialized yet. If $force is true, it will reinitialize the collection even if it has already been initialized.
     * @param bool $force 
     * @return void 
     * @throws Exception 
     */
    public function initCollection($force=false) {
        if(!$this->hasInit || $force) {
            if($this->collection) {
                $this->collection->save();
                $this->hasInit = true;
                return;
            }
            throw new \Exception("Collection not set");
        }
    }

    /**
     * Create's the collection if it doesn't exist, returns the collection object.
     * @return Collection
     */
    public function createCollection() {
        if($this->collection) {
            $this->collection->save();
            return $this->collection;
        }
        return null;
    }

    public function deleteCollection() {
        if($this->collection) {
            return $this->collection->delete(null,true);
        }
        return false;
    }

    /**
     * Upserts a collection record into the vector collection.
     * @param CollectionRecord $collectionRecord
     * @return bool
     * @throws \Exception
     */
    public function upsert($collectionRecord,$batchSize=100) {
        $this->initCollection();
        $point = $this->extractPoint($collectionRecord);
        if($point === false) {
            throw new \Exception("Invalid point type");
        }
        return $this->collection->addPoint($point);
    }

    /**
     * Upserts a batch of collection records into the vector collection.
     * @param CollectionRecord[] $collectionRecords
     * @param int $batchSize
     * @return bool
     * @throws \Exception
     */
    public function upsertBatch($collectionRecords,$batchSize=100) {
        $this->initCollection();
        if(!is_array($collectionRecords)) {
            $collectionRecords = [$collectionRecords];
        }
        $points = $this->extractPoints($collectionRecords);
        if($points === false) {
            throw new \Exception("Invalid point type");
        }
        if(count($points) <= $batchSize) {
            return $this->collection->addPoints($points,$batchSize);
        }
        $points = array_chunk($points,$batchSize);
        foreach($points as $chunk) {
            $this->collection->addPoints($chunk,$batchSize);
        }
        return true;
    }

    /**
     * Returns the point object from one CollectionRecord or Point.
     * @param CollectionRecord|Point $record
     * @return Point
     */
    public function extractPoint($record) {
        if($record instanceof CollectionRecord) {
            return $record->point();
        } elseif($record instanceof Point) {
            return $record;
        }
        return false;
    }

    /**
     * Returns the points from the collection records. If the records are not an array, it will return the point directly.
     * @param CollectionRecord|Point|array $records
     * @return Point[]|Point|false
     */
    public function extractPoints($records) {
        if(is_array($records)) {
            $points = [];
            foreach($records as $record) {
                $point = $this->extractPoint($record);
                if($point) {
                    $points[] = $point;
                }
            }
            return $points;
        } elseif($records instanceof CollectionRecord || $records instanceof Point) {
            return $this->extractPoint($records);
        }
        return false;
    }

    private function fromOptions($options=[]) {
        $vectorParams = $collectionName = null;
        if(!is_array($options) && is_string($options)) {
            $check = json_decode($options,true);
            if($check) {
                $options = $check;
            }
        }
        if(is_array($options)) {
            if(isset($options["vectorParams"])) {
                $vectorParams = $options["vectorParams"];
                unset($options["vectorParams"]);
            } elseif(isset($options["vector"])) {
                $vectorParams = $options["vector"];
                unset($options["vector"]);
            }
            if(isset($options["collectionName"])) {
                $collectionName = $options["collectionName"];
                unset($options["collectionName"]);
            } elseif(isset($options["collection"])) {
                $collectionName = $options["collection"];
                unset($options["collection"]);
            }
        } elseif($options instanceof Collection) {
            $collectionName = $options;
        } elseif(is_string($options)) {
            $collectionName = $options;
        }
        if($collectionName && $collectionName instanceof Collection) {
            $this->collection = $collectionName;
        } elseif($collectionName) {
            $this->collection = new Collection($collectionName,$vectorParams,$options);
        }
    }

    /**
     * @param string $id
     * @return CollectionRecord|null
     */
    public function retrieve($id) {
        if($this->collection) {
            try {
                $point = $this->collection->getPoint($id);
                if($point) {
                    return CollectionRecord::fromPoint($point);
                }
            } catch (\Exception $e) {
                Output::outLine("Error retrieving point: " . $e->getMessage());
            }
        }
        return null;
    }

    /**
     * 
     * @param string|CollectionRecord|Point $idOrRecord 
     * @return boru\qdrant\ApiBaseResponse|false 
     */
    public function delete($idOrRecord) {
        if($this->collection) {
            if($idOrRecord instanceof CollectionRecord) {
                $id = $idOrRecord->uid();
            } elseif($idOrRecord instanceof Point) {
                $id = $idOrRecord->id();
            } elseif(is_string($idOrRecord)) {
                $id = $idOrRecord;
            } else {
                return false;
            }
            try {
                return $this->collection->deletePoint($id);
            } catch (\Exception $e) {
                Output::outLine("Error deleting point: " . $e->getMessage());
            }
        }
        return false;
    }

    /**
     * Returns the class name to utilize for the record. If $this->recordClassName is set, it will be used. If the recordClassOrObject is a CollectionRecord, its class name will be used. If the recordClassOrObject is a string, it will be used as the class name. Otherwise, CollectionRecord::class will be returned.
     * @param mixed|null $recordClassOrObject
     * @return string|null
     */
    protected function getRecordClassName($recordClassOrObject=null) {
        if($this->recordClassName) {
            return $this->recordClassName;
        }
        if($recordClassOrObject instanceof CollectionRecord) {
            return get_class($recordClassOrObject);
        } elseif(is_string($recordClassOrObject) && class_exists($recordClassOrObject)) {
            return $recordClassOrObject;
        }
        return CollectionRecord::class;
    }

    public function newRecord($recordClassOrObject=null) {
        $className = $this->getRecordClassName($recordClassOrObject);
        if(class_exists($className)) {
            $object = new $className();
            $object->collection($this);
            return $object;
        }
        return null;
    }


    public function search($pointRecordOrString,$filter=[],$limit=10,$withPayload=true,$withVector=false) {
        if($this->collection()) {
            $vector = null;
            $className = $this->getRecordClassName($pointRecordOrString);
            if($pointRecordOrString instanceof CollectionRecord) {
                $point = $pointRecordOrString->point();
                $vector = $pointRecordOrString->vector();
            } elseif($pointRecordOrString instanceof Point) {
                $point = $pointRecordOrString;
                $vector = $point->vector();
            } elseif(is_array($pointRecordOrString)) {
                $vector = $pointRecordOrString;
            } elseif(is_string($pointRecordOrString)) {
                $vector = Embedding::embed($pointRecordOrString);
            }
            if($vector) {
                $points = $this->collection->search($vector,$filter,$limit,$withPayload,$withVector);
                if($points) {
                    $records = [];
                    foreach($points as $point) {
                        $record = $className::fromPoint($point);
                        if($record) {
                            $records[] = $record;
                        }
                    }
                    return $records;
                }
            }
        }
        return [];
    }
}