<?php

if ( ! defined( 'ABSPATH' ) ) {
    exit;
}

class Niwa_AI_Agent_Memory_Index {
    private const INDEX_DIR = 'agent-memory';
    private const INDEX_FILE = 'index.json';
    private const SESSION_EXPORT_DIR = 'sessions';
    private const CHUNK_LINE_WINDOW = 12;
    private const CHUNK_LINE_OVERLAP = 3;
    private const MAX_SEARCH_RESULTS = 12;
    private const VECTOR_WEIGHT = 0.7;
    private const TEXT_WEIGHT = 0.3;
    private const TEMPORAL_HALF_LIFE_DAYS = 30;

    private $embedding_client;
    private $workspace = null;

    public function __construct( $embedding_client = null ) {
        $this->embedding_client = $embedding_client;
    }

    public function set_workspace( $workspace ) {
        $this->workspace = $workspace;
    }

    public function search( $query, $documents, $options = array() ) {
        $query     = trim( (string) $query );
        $documents = is_array( $documents ) ? array_values( $documents ) : array();
        $limit     = isset( $options['limit'] ) ? max( 1, min( self::MAX_SEARCH_RESULTS, intval( $options['limit'] ) ) ) : 6;
        $session_id = isset( $options['session_id'] ) ? trim( (string) $options['session_id'] ) : '';
        if ( '' === $query || empty( $documents ) ) {
            return array();
        }

        $index = $this->sync_documents(
            $documents,
            array(
                'embed_missing'   => ! empty( $options['embed_missing'] ),
                'persist_index'   => ! empty( $options['persist_index'] ),
                'write_artifacts' => ! empty( $options['write_artifacts'] ),
            )
        );
        $chunks = isset( $index['chunks'] ) && is_array( $index['chunks'] ) ? $index['chunks'] : array();
        if ( empty( $chunks ) ) {
            return array();
        }

        $terms      = $this->tokenize_query( $query );
        $query_vec  = $this->embedding_client && method_exists( $this->embedding_client, 'embed_text' )
            ? (array) $this->embedding_client->embed_text( $query )
            : array();
        $candidates = array();

        foreach ( $chunks as $chunk ) {
            if ( ! is_array( $chunk ) || empty( $chunk['text'] ) ) {
                continue;
            }

            $text_score = $this->score_text_chunk( $query, $terms, $chunk );
            $vector_score = $this->score_vector_chunk( $query_vec, $chunk );
            if ( $text_score <= 0 && $vector_score <= 0 ) {
                continue;
            }

            $hybrid = ( self::TEXT_WEIGHT * $text_score ) + ( self::VECTOR_WEIGHT * $vector_score );
            $hybrid *= $this->resolve_source_weight( $chunk, $session_id );
            $hybrid *= $this->resolve_temporal_decay( isset( $chunk['timestamp'] ) ? intval( $chunk['timestamp'] ) : 0 );

            if ( $hybrid < 0.2 ) {
                continue;
            }

            $path       = isset( $chunk['path'] ) ? (string) $chunk['path'] : '';
            $start_line = max( 1, intval( isset( $chunk['start_line'] ) ? $chunk['start_line'] : 1 ) );
            $end_line   = max( $start_line, intval( isset( $chunk['end_line'] ) ? $chunk['end_line'] : $start_line ) );

            $candidates[] = array(
                'entry_id'    => isset( $chunk['id'] ) ? (string) $chunk['id'] : md5( $path . ':' . $start_line . ':' . $end_line ),
                'path'        => $path,
                'start_line'  => $start_line,
                'end_line'    => $end_line,
                'snippet'     => $this->clip_text( (string) $chunk['text'], 900 ),
                'score'       => round( $hybrid, 4 ),
                'text_score'  => round( $text_score, 4 ),
                'vector_score' => round( $vector_score, 4 ),
                'source_type' => isset( $chunk['source_type'] ) ? (string) $chunk['source_type'] : '',
                'session_id'  => isset( $chunk['session_id'] ) ? (string) $chunk['session_id'] : '',
                'timestamp'   => isset( $chunk['timestamp'] ) ? intval( $chunk['timestamp'] ) : 0,
            );
        }

        if ( empty( $candidates ) ) {
            return array();
        }

        usort(
            $candidates,
            function ( $left, $right ) {
                return floatval( $right['score'] ) <=> floatval( $left['score'] );
            }
        );

        return array_slice( $candidates, 0, $limit );
    }

    public function read_entry( $rel_path, $from = 1, $lines = 20 ) {
        $rel_path = ltrim( str_replace( '\\', '/', trim( (string) $rel_path ) ), '/' );
        $from     = max( 1, intval( $from ) );
        $lines    = max( 1, min( 120, intval( $lines ) ) );
        if ( '' === $rel_path ) {
            return array(
                'path'  => '',
                'text'  => '',
                'error' => 'Agent memory path je prazan.',
            );
        }

        $abs_path = $this->resolve_rel_path( $rel_path );
        if ( '' === $abs_path || ! file_exists( $abs_path ) || ! is_readable( $abs_path ) ) {
            return array(
                'path'  => $rel_path,
                'text'  => '',
                'error' => 'Agent memory entry nije pronadjen.',
            );
        }

        $read = Niwa_AI_Text_Normalizer::normalize_file_contents( $abs_path );
        if ( empty( $read['success'] ) ) {
            return array(
                'path'  => $rel_path,
                'text'  => '',
                'error' => 'Agent memory entry nije citljiv.',
            );
        }

        $raw_lines = Niwa_AI_Text_Normalizer::split_lines( isset( $read['content'] ) ? $read['content'] : '' );
        $slice = array_slice( $raw_lines, $from - 1, $lines );
        return array(
            'path'       => $rel_path,
            'from'       => $from,
            'lines'      => count( $slice ),
            'start_line' => $from,
            'end_line'   => $from + max( 0, count( $slice ) - 1 ),
            'text'       => implode( "\n", $slice ),
        );
    }

    public function build_session_export_document( $session_id, $manifest = array(), $records = array() ) {
        $session_id = trim( (string) $session_id );
        if ( '' === $session_id ) {
            return array();
        }

        $manifest = is_array( $manifest ) ? $manifest : array();
        $records  = is_array( $records ) ? $records : array();

        $lines = array(
            '# Session Memory Export',
            'session_id: ' . $session_id,
            '',
        );

        $worklog = isset( $manifest['worklog_summary'] ) ? trim( (string) $manifest['worklog_summary'] ) : '';
        if ( '' !== $worklog ) {
            $lines[] = '## Worklog';
            foreach ( Niwa_AI_Text_Normalizer::split_lines( $worklog, true ) as $line ) {
                $line = trim( (string) $line );
                if ( '' !== $line ) {
                    $lines[] = $line;
                }
            }
            $lines[] = '';
        }

        $sticky_notes = isset( $manifest['sticky_notes'] ) && is_array( $manifest['sticky_notes'] ) ? array_slice( $manifest['sticky_notes'], -8 ) : array();
        if ( ! empty( $sticky_notes ) ) {
            $lines[] = '## Sticky Notes';
            foreach ( $sticky_notes as $note ) {
                if ( ! is_array( $note ) || empty( $note['text'] ) ) {
                    continue;
                }
                $lines[] = '- ' . Niwa_AI_Text_Normalizer::normalize_inline( $note['text'] );
            }
            $lines[] = '';
        }

        $lines[] = '## Transcript';
        foreach ( $records as $record ) {
            if ( ! is_array( $record ) ) {
                continue;
            }

            $type = isset( $record['type'] ) ? (string) $record['type'] : '';
            if ( 'turn_start' === $type && ! empty( $record['prompt'] ) ) {
                $lines[] = 'User: ' . Niwa_AI_Text_Normalizer::normalize_inline( $record['prompt'] );
            } elseif ( 'turn_end' === $type && ! empty( $record['reply'] ) ) {
                $lines[] = 'Assistant: ' . Niwa_AI_Text_Normalizer::normalize_inline( $record['reply'] );
            } elseif ( 'sticky_note' === $type && ! empty( $record['text'] ) ) {
                $lines[] = 'Note: ' . Niwa_AI_Text_Normalizer::normalize_inline( $record['text'] );
            } elseif ( 'subagent_result' === $type && ! empty( $record['subagent_record']['summary'] ) ) {
                $lines[] = 'Subagent: ' . Niwa_AI_Text_Normalizer::normalize_inline( $record['subagent_record']['summary'] );
            }
        }

        $content = Niwa_AI_Text_Normalizer::normalize(
            implode( "\n", array_values( array_filter( $lines, 'is_string' ) ) ) . "\n",
            array(
                'trim'              => false,
                'preserve_newlines' => true,
            )
        );

        return array(
            'path'       => $this->get_session_export_rel_path( $session_id ),
            'abs_path'   => $this->get_session_export_abs_path( $session_id ),
            'content'    => $content,
            'source_type' => 'session_export',
            'session_id' => $session_id,
            'timestamp'  => isset( $manifest['updated_at'] ) ? intval( $manifest['updated_at'] ) : time(),
        );
    }

    private function sync_documents( $documents, $options = array() ) {
        $documents = is_array( $documents ) ? array_values( $documents ) : array();
        $options   = is_array( $options ) ? $options : array();
        $embed_missing   = ! empty( $options['embed_missing'] );
        $persist_index   = ! empty( $options['persist_index'] );
        $write_artifacts = ! empty( $options['write_artifacts'] );
        $this->ensure_index_dirs();

        $existing   = $this->load_index();
        $old_chunks = array();
        foreach ( isset( $existing['chunks'] ) && is_array( $existing['chunks'] ) ? $existing['chunks'] : array() as $chunk ) {
            if ( ! is_array( $chunk ) || empty( $chunk['id'] ) ) {
                continue;
            }
            $old_chunks[ (string) $chunk['id'] ] = $chunk;
        }

        $docs   = array();
        $chunks = array();
        foreach ( $documents as $document ) {
            $normalized = $this->normalize_document( $document, $write_artifacts );
            if ( empty( $normalized['path'] ) || empty( $normalized['content'] ) ) {
                continue;
            }

            $doc_chunks = $this->chunk_document( $normalized );
            foreach ( $doc_chunks as &$chunk ) {
                $chunk_id = isset( $chunk['id'] ) ? (string) $chunk['id'] : '';
                if ( '' !== $chunk_id && ! empty( $old_chunks[ $chunk_id ]['vector'] ) && is_array( $old_chunks[ $chunk_id ]['vector'] ) ) {
                    $chunk['vector'] = $old_chunks[ $chunk_id ]['vector'];
                }
            }
            unset( $chunk );

            $docs[] = array(
                'path'        => (string) $normalized['path'],
                'abs_path'    => (string) $normalized['abs_path'],
                'hash'        => (string) $normalized['hash'],
                'timestamp'   => intval( isset( $normalized['timestamp'] ) ? $normalized['timestamp'] : 0 ),
                'source_type' => isset( $normalized['source_type'] ) ? (string) $normalized['source_type'] : '',
                'session_id'  => isset( $normalized['session_id'] ) ? (string) $normalized['session_id'] : '',
            );
            $chunks = array_merge( $chunks, $doc_chunks );
        }

        if ( $embed_missing ) {
            $this->embed_missing_chunks( $chunks );
        }

        $index = array(
            'version'    => 1,
            'updated_at' => time(),
            'documents'  => $docs,
            'chunks'     => $chunks,
        );
        if ( $persist_index ) {
            $this->write_index( $index );
        }

        return $index;
    }

    private function normalize_document( $document, $write_artifacts = false ) {
        $document = is_array( $document ) ? $document : array();
        $path     = isset( $document['path'] ) ? ltrim( str_replace( '\\', '/', trim( (string) $document['path'] ) ), '/' ) : '';
        $abs_path = isset( $document['abs_path'] ) ? (string) $document['abs_path'] : '';
        $content  = Niwa_AI_Text_Normalizer::normalize(
            isset( $document['content'] ) ? (string) $document['content'] : '',
            array(
                'trim'              => false,
                'preserve_newlines' => true,
            )
        );
        if ( '' === $path || '' === $abs_path || '' === trim( $content ) ) {
            return array();
        }

        $source_type  = isset( $document['source_type'] ) ? (string) $document['source_type'] : '';
        $should_write = 'session_export' === $source_type;
        if ( $write_artifacts && file_exists( $abs_path ) && is_writable( $abs_path ) ) {
            $current = file_get_contents( $abs_path );
            if ( false !== $current && (string) $current !== $content ) {
                @file_put_contents( $abs_path, $content, LOCK_EX );
            }
        } elseif ( $write_artifacts && $should_write && ! file_exists( $abs_path ) ) {
            @file_put_contents( $abs_path, $content );
        }

        return array(
            'path'        => $path,
            'abs_path'    => $abs_path,
            'content'     => $content,
            'hash'        => md5( $content ),
            'timestamp'   => isset( $document['timestamp'] ) ? intval( $document['timestamp'] ) : ( file_exists( $abs_path ) ? intval( @filemtime( $abs_path ) ) : time() ),
            'source_type' => $source_type,
            'session_id'  => isset( $document['session_id'] ) ? (string) $document['session_id'] : '',
        );
    }

    private function chunk_document( $document ) {
        $lines         = Niwa_AI_Text_Normalizer::split_lines( isset( $document['content'] ) ? $document['content'] : '' );
        $window        = self::CHUNK_LINE_WINDOW;
        $overlap       = self::CHUNK_LINE_OVERLAP;
        $chunks        = array();
        $line_count    = count( $lines );
        $start_index   = 0;

        while ( $start_index < $line_count ) {
            $slice = array_slice( $lines, $start_index, $window );
            $text  = trim( implode( "\n", $slice ) );
            if ( '' !== $text ) {
                $start_line = $start_index + 1;
                $end_line   = $start_line + count( $slice ) - 1;
                $chunks[]   = array(
                    'id'         => md5( $document['path'] . ':' . $document['hash'] . ':' . $start_line . ':' . $end_line ),
                    'path'       => (string) $document['path'],
                    'start_line' => $start_line,
                    'end_line'   => $end_line,
                    'text'       => $text,
                    'vector'     => array(),
                    'source_type' => isset( $document['source_type'] ) ? (string) $document['source_type'] : '',
                    'session_id' => isset( $document['session_id'] ) ? (string) $document['session_id'] : '',
                    'timestamp'  => intval( isset( $document['timestamp'] ) ? $document['timestamp'] : 0 ),
                );
            }

            if ( $start_index + $window >= $line_count ) {
                break;
            }

            $start_index += max( 1, $window - $overlap );
        }

        return $chunks;
    }

    private function embed_missing_chunks( &$chunks ) {
        if ( ! $this->embedding_client || ! method_exists( $this->embedding_client, 'embed_texts' ) ) {
            return;
        }

        $pending_indexes = array();
        $pending_texts   = array();
        foreach ( $chunks as $index => $chunk ) {
            if ( ! is_array( $chunk ) || ! empty( $chunk['vector'] ) || empty( $chunk['text'] ) ) {
                continue;
            }
            $pending_indexes[] = $index;
            $pending_texts[]   = (string) $chunk['text'];
        }

        if ( empty( $pending_texts ) ) {
            return;
        }

        $vectors = (array) $this->embedding_client->embed_texts( $pending_texts );
        foreach ( $pending_indexes as $position => $chunk_index ) {
            if ( ! isset( $vectors[ $position ] ) || ! is_array( $vectors[ $position ] ) ) {
                continue;
            }
            $chunks[ $chunk_index ]['vector'] = array_values( $vectors[ $position ] );
        }
    }

    private function score_text_chunk( $query, $terms, $chunk ) {
        $text = mb_strtolower( (string) ( isset( $chunk['text'] ) ? $chunk['text'] : '' ) );
        if ( '' === $text ) {
            return 0.0;
        }

        $hits = 0;
        foreach ( (array) $terms as $term ) {
            if ( '' !== $term && false !== mb_strpos( $text, $term ) ) {
                $hits++;
            }
        }
        if ( $hits <= 0 ) {
            return 0.0;
        }

        $coverage = $hits / max( 1, count( (array) $terms ) );
        $phrase   = false !== mb_strpos( $text, mb_strtolower( (string) $query ) ) ? 0.35 : 0.0;
        return min( 1.0, $coverage + $phrase );
    }

    private function score_vector_chunk( $query_vec, $chunk ) {
        if ( empty( $query_vec ) || empty( $chunk['vector'] ) || ! is_array( $chunk['vector'] ) ) {
            return 0.0;
        }

        $sim = $this->cosine_similarity( $query_vec, $chunk['vector'] );
        return max( 0.0, min( 1.0, ( $sim + 1 ) / 2 ) );
    }

    private function resolve_source_weight( $chunk, $current_session_id ) {
        $source_type = isset( $chunk['source_type'] ) ? (string) $chunk['source_type'] : '';
        $path        = strtoupper( isset( $chunk['path'] ) ? (string) $chunk['path'] : '' );
        $session_id  = isset( $chunk['session_id'] ) ? (string) $chunk['session_id'] : '';
        $weight      = 1.0;

        if ( 'session_export' === $source_type ) {
            $weight = 1.14;
        } elseif ( 'companion_diary' === $source_type ) {
            $weight = 1.18;
        } elseif ( 'companion_decision' === $source_type ) {
            $weight = 1.16;
        } elseif ( 'companion_session' === $source_type ) {
            $weight = 1.13;
        } elseif ( 'companion_profile' === $source_type ) {
            $weight = 1.11;
        } elseif ( 'daily_memory' === $source_type || 0 === strpos( $path, 'MEMORY/' ) ) {
            $weight = 1.12;
        } elseif ( false !== strpos( $path, 'SOUL.MD' ) || false !== strpos( $path, 'IDENTITY.MD' ) || false !== strpos( $path, 'USER.MD' ) ) {
            $weight = 1.1;
        } elseif ( false !== strpos( $path, 'MEMORY.MD' ) ) {
            $weight = 1.08;
        }

        if ( '' !== $current_session_id && $session_id === $current_session_id ) {
            $weight += 0.08;
        }

        return $weight;
    }

    private function resolve_temporal_decay( $timestamp ) {
        $timestamp = intval( $timestamp );
        if ( $timestamp <= 0 ) {
            return 1.0;
        }

        $age_days = max( 0, ( time() - $timestamp ) / DAY_IN_SECONDS );
        if ( $age_days <= 0 ) {
            return 1.1;
        }

        $decay = pow( 0.5, $age_days / self::TEMPORAL_HALF_LIFE_DAYS );
        return max( 0.75, min( 1.1, 0.75 + ( 0.35 * $decay ) ) );
    }

    private function cosine_similarity( $left, $right ) {
        $left  = is_array( $left ) ? array_values( $left ) : array();
        $right = is_array( $right ) ? array_values( $right ) : array();
        if ( empty( $left ) || empty( $right ) ) {
            return 0.0;
        }

        $length = min( count( $left ), count( $right ) );
        if ( $length <= 0 ) {
            return 0.0;
        }

        $dot = 0.0;
        $norm_left = 0.0;
        $norm_right = 0.0;
        for ( $i = 0; $i < $length; $i++ ) {
            $lv = floatval( $left[ $i ] );
            $rv = floatval( $right[ $i ] );
            $dot += $lv * $rv;
            $norm_left += $lv * $lv;
            $norm_right += $rv * $rv;
        }

        if ( $norm_left <= 0 || $norm_right <= 0 ) {
            return 0.0;
        }

        return $dot / ( sqrt( $norm_left ) * sqrt( $norm_right ) );
    }

    private function tokenize_query( $query ) {
        $query = Niwa_AI_Text_Normalizer::normalize_inline( $query );
        $terms = preg_split( '/[^\p{L}\p{N}_-]+/u', mb_strtolower( $query ) );
        return array_values(
            array_filter(
                array_unique( (array) $terms ),
                function ( $term ) {
                    return mb_strlen( (string) $term ) >= 3;
                }
            )
        );
    }

    private function load_index() {
        $path = $this->get_index_path();
        if ( '' === $path || ! file_exists( $path ) ) {
            return array();
        }

        $raw = file_get_contents( $path );
        if ( false === $raw ) {
            return array();
        }

        $decoded = json_decode( (string) $raw, true );
        return is_array( $decoded ) ? $decoded : array();
    }

    private function write_index( $index ) {
        $path = $this->get_index_path();
        if ( '' === $path ) {
            return false;
        }

        $dir = dirname( $path );
        if ( ! file_exists( $dir ) ) {
            wp_mkdir_p( $dir );
        }

        $json = wp_json_encode( $index, JSON_PRETTY_PRINT | JSON_UNESCAPED_SLASHES );
        if ( ! is_string( $json ) || '' === $json ) {
            return false;
        }

        return false !== file_put_contents( $path, $json, LOCK_EX );
    }

    private function ensure_index_dirs() {
        $root = $this->get_index_root();
        if ( '' === $root ) {
            return;
        }

        wp_mkdir_p( $root );
        wp_mkdir_p( trailingslashit( $root ) . self::SESSION_EXPORT_DIR );
    }

    private function get_index_root() {
        if ( ! $this->workspace || ! method_exists( $this->workspace, 'get_paths' ) ) {
            return '';
        }

        $paths = (array) $this->workspace->get_paths();
        if ( empty( $paths['root'] ) ) {
            return '';
        }

        return trailingslashit( (string) $paths['root'] ) . self::INDEX_DIR;
    }

    private function get_index_path() {
        $root = $this->get_index_root();
        if ( '' === $root ) {
            return '';
        }

        return trailingslashit( $root ) . self::INDEX_FILE;
    }

    private function get_session_export_rel_path( $session_id ) {
        $filename = basename( $this->get_session_export_abs_path( $session_id ) );
        return self::INDEX_DIR . '/' . self::SESSION_EXPORT_DIR . '/' . $filename;
    }

    private function get_session_export_abs_path( $session_id ) {
        $root = $this->get_index_root();
        if ( '' === $root ) {
            return '';
        }

        $stem = 'session-' . substr( md5( (string) $session_id ), 0, 12 ) . '.md';
        if ( $this->workspace && method_exists( $this->workspace, 'get_session_transcript_path' ) ) {
            $transcript_path = (string) $this->workspace->get_session_transcript_path( $session_id );
            if ( '' !== $transcript_path ) {
                $stem = preg_replace( '/\.jsonl$/i', '.md', basename( $transcript_path ) );
            }
        }

        return trailingslashit( $root ) . self::SESSION_EXPORT_DIR . '/' . $stem;
    }

    private function resolve_rel_path( $rel_path ) {
        if ( ! $this->workspace || ! method_exists( $this->workspace, 'get_paths' ) ) {
            return '';
        }

        $paths = (array) $this->workspace->get_paths();
        $root  = isset( $paths['root'] ) ? realpath( (string) $paths['root'] ) : false;
        if ( false === $root || '' === $root ) {
            return '';
        }

        $candidate = $root . DIRECTORY_SEPARATOR . str_replace( '/', DIRECTORY_SEPARATOR, ltrim( (string) $rel_path, '/' ) );
        $real      = realpath( $candidate );
        if ( false === $real ) {
            return $candidate;
        }

        if ( 0 !== strpos( $real, $root ) ) {
            return '';
        }

        return $real;
    }

    private function clip_text( $text, $limit = 220 ) {
        $text = Niwa_AI_Text_Normalizer::normalize_inline( $text );
        if ( mb_strlen( $text ) <= $limit ) {
            return $text;
        }

        return rtrim( mb_substr( $text, 0, $limit ) ) . '...';
    }
}
