import { Injectable, Logger, OnModuleInit, Optional } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { QdrantClient } from '@qdrant/js-client-rest';
import { AiGatewayService } from '../ai/gateway/ai-gateway.service';
import type { CitationContext } from './integration-types';

export interface VectorPoint {
  id: string;
  embedding: number[];
  payload: {
    userId: string;
    knowledgeBaseId: string;
    sourceId: string;
    sourceType: string;
    chunkIndex: number;
    text: string;
    [key: string]: any;
  };
}

export interface SearchFilters {
  userId?: string;
  knowledgeBaseId?: string;
  sourceType?: string;
  sourceId?: string;
  mustNotDeleted?: boolean;
}

export interface SearchResult {
  id: string;
  score: number;
  payload: Record<string, any>;
}

export interface RerankedResult {
  id: string;
  score: number;
  relevanceScore: number;
  payload: Record<string, any>;
}

const COLLECTION_NAME = 'zhixi_chunks';
const VECTOR_SIZE = 1024;

@Injectable()
export class VectorService implements OnModuleInit {
  private readonly logger = new Logger(VectorService.name);
  private client: QdrantClient;
  private initialized = false;

  constructor(
    private readonly config: ConfigService,
    @Optional() private readonly aiGateway?: AiGatewayService,
  ) {
    const url = config.get<string>('qdrant.url', 'http://127.0.0.1:6333');
    this.client = new QdrantClient({ url });
  }

  async onModuleInit() {
    // Don't block startup — init lazily on first request
  }

  private async ensureInit() {
    if (this.initialized) return;
    try {
      await this.client.getCollection(COLLECTION_NAME);
      this.logger.log(`Connected to Qdrant collection: ${COLLECTION_NAME}`);
      this.initialized = true;
    } catch {
      this.logger.warn(`Qdrant collection ${COLLECTION_NAME} not available, will retry later`);
    }
  }

  private async ensureCollection() {
    try {
      await this.client.createCollection(COLLECTION_NAME, {
        vectors: { size: VECTOR_SIZE, distance: 'Cosine' },
        hnsw_config: { m: 16, ef_construct: 100 },
      });
      await this.client.createPayloadIndex(COLLECTION_NAME, { field_name: 'userId', field_schema: 'keyword' });
      await this.client.createPayloadIndex(COLLECTION_NAME, { field_name: 'knowledgeBaseId', field_schema: 'keyword' });
      await this.client.createPayloadIndex(COLLECTION_NAME, { field_name: 'deleted', field_schema: 'bool' });
      this.logger.log(`Created Qdrant collection: ${COLLECTION_NAME}`);
    } catch (err: any) {
      this.logger.error(`Failed to create Qdrant collection: ${err.message}`);
    }
  }

  async upsert(points: VectorPoint[]): Promise<void> {
    await this.ensureInit();
    if (points.length === 0) return;
    await this.client.upsert(COLLECTION_NAME, {
      wait: true,
      points: points.map(p => ({
        id: p.id,
        vector: p.embedding,
        payload: { ...p.payload, deleted: false },
      })),
    });
    this.logger.log(`Upserted ${points.length} vectors`);
  }

  async deleteBySource(sourceId: string): Promise<void> {
    await this.ensureInit();
    await this.client.delete(COLLECTION_NAME, {
      filter: { must: [{ key: 'sourceId', match: { value: sourceId } }] },
      wait: true,
    });
  }

  async deleteBySourceIds(sourceIds: string[]): Promise<void> {
    if (sourceIds.length === 0) return;
    await this.client.delete(COLLECTION_NAME, {
      filter: {
        must: [{ key: 'sourceId', match: { any: sourceIds } }],
      },
      wait: true,
    });
  }

  async search(embedding: number[], filters: SearchFilters = {}, topK = 10): Promise<SearchResult[]> {
    await this.ensureInit();
    const must: any[] = [];
    if (filters.userId) must.push({ key: 'userId', match: { value: filters.userId } });
    if (filters.knowledgeBaseId) must.push({ key: 'knowledgeBaseId', match: { value: filters.knowledgeBaseId } });
    if (filters.sourceType) must.push({ key: 'sourceType', match: { value: filters.sourceType } });
    if (filters.sourceId) must.push({ key: 'sourceId', match: { value: filters.sourceId } });
    if (filters.mustNotDeleted !== false) must.push({ key: 'deleted', match: { value: false } });

    const results = await this.client.search(COLLECTION_NAME, {
      vector: embedding,
      limit: topK,
      filter: must.length > 0 ? { must } : undefined,
      with_payload: true,
    });

    return results.map(r => ({
      id: String(r.id),
      score: r.score,
      payload: r.payload as Record<string, any>,
    }));
  }

  async rerank(query: string, results: SearchResult[]): Promise<RerankedResult[]> {
    const reranker = this.config.get<string>('ai.rerank.model', 'BAAI/bge-reranker-v2-m3');
    try {
      const response = await this.aiGateway?.generate({
        feature: 'rerank',
        userId: 'system',
        tier: 'cheap',
        promptKey: 'rerank',
        promptVersion: 'v1',
        messages: [
          { role: 'system', content: 'Rerank search results' },
          { role: 'user', content: JSON.stringify({ query, documents: results.map(r => r.payload.text) }) },
        ],
      });

      // Fallback: use vector scores if AI rerank unavailable
      if (!response?.parsed) {
        return results.map((r, i) => ({ ...r, relevanceScore: r.score }));
      }

      return results.map((r, i) => ({
        ...r,
        relevanceScore: response.parsed.scores?.[i] ?? r.score,
      }));
    } catch {
      return results.map(r => ({ ...r, relevanceScore: r.score }));
    }
  }

  // ── Admin helpers ──

  async getCollectionInfo() {
    try {
      const info = await this.client.getCollection(COLLECTION_NAME);
      const count = await this.client.count(COLLECTION_NAME);
      return {
        name: COLLECTION_NAME,
        vectorSize: info.config.params.vectors?.size || VECTOR_SIZE,
        distance: info.config.params.vectors?.distance || 'Cosine',
        pointsCount: count.count,
        status: info.status,
      };
    } catch (err: any) {
      return { name: COLLECTION_NAME, error: err.message };
    }
  }

  async count(): Promise<number> {
    const result = await this.client.count(COLLECTION_NAME);
    return result.count;
  }

  /** Build citation context from search results */
  async buildCitationContexts(results: SearchResult[]): Promise<CitationContext[]> {
    return results.map(r => ({
      chunkId: r.id,
      sourceId: r.payload?.sourceId || 'unknown',
      sourceTitle: r.payload?.sourceTitle || 'Untitled',
      content: r.payload?.text || '',
      pageNumber: r.payload?.pageNumber,
      sectionTitle: r.payload?.sectionTitle,
      score: r.score,
    }));
  }
}
