/**
 * LLM-based Relationship Classifier for Memory Graph
 *
 * Uses an LLM to intelligently classify how new facts relate to existing facts:
 * - NEW: Completely new, unrelated information
 * - UPDATES: Replaces/contradicts an existing fact (same slot, different value)
 * - EXTENDS: Adds information about an existing entity/relationship
 */

import type { LLMProvider, CompletionRequest, ContentBlock } from '../providers/types.js';

/**
 * A fact to be classified
 */
export interface FactToClassify {
  content: string;
  subject: string;
  category: string;
}

/**
 * An existing fact for comparison
 */
export interface ExistingFact {
  id: string;
  content: string;
  subject: string;
  category: string;
}

/**
 * Classification result
 */
export interface ClassificationResult {
  classification: 'NEW' | 'UPDATES' | 'EXTENDS';
  targetId?: string; // ID of the fact being updated/extended
  confidence: number;
  reason: string;
}

/**
 * Prompt for relationship classification
 */
const CLASSIFICATION_PROMPT = `You are a memory relationship classifier. Given a NEW fact and EXISTING facts, determine how the new fact relates to existing ones.

Classifications:
- NEW: The fact is completely new information, unrelated to existing facts
- UPDATES: The new fact REPLACES or SUPERSEDES an existing fact (same topic/attribute, different value)
- EXTENDS: The new fact adds MORE information about an entity mentioned in existing facts

UPDATES examples (same slot, different value — the old fact becomes outdated):
  - "Lives in Metropolis" UPDATES "Lives in Springfield" (location changed)
  - "Works at Acme Corp" UPDATES "Works at Globex" (employer changed)
  - "Name: Johnny" UPDATES "Name: John" (name correction)
  - "Mood: stressed about deadlines" UPDATES "Mood: happy and relaxed" (mood changed)
  - "Focus: learning Rust" UPDATES "Focus: learning Python" (focus shifted)
  - "Prefers dark mode" UPDATES "Prefers light mode" (preference changed)
  - "Has 2 cats" UPDATES "Has 1 cat" (quantity changed)

EXTENDS examples (adding info about same entity):
  - "Flatmate is Bob" exists → "Bob works at Google" EXTENDS it
  - "Wife is Sarah" exists → "Sarah likes cooking" EXTENDS it

IMPORTANT RULES:
1. Different relationship types are NOT updates: "Flatmate is Bob" and "Wife is Jamie" are BOTH NEW
2. UPDATES requires the SAME attribute/slot with a CHANGED value — look for contradictions
3. When in doubt between NEW and UPDATES, check if both facts can be true simultaneously.
   If they CANNOT both be true → UPDATES. If they CAN → NEW or EXTENDS.

Respond with JSON only:
{
  "classification": "NEW" | "UPDATES" | "EXTENDS",
  "targetId": "id of existing fact if UPDATES or EXTENDS, omit for NEW",
  "confidence": 0.0-1.0,
  "reason": "brief explanation"
}`;

/**
 * Batch classification result for multiple facts
 */
export interface BatchClassificationResult {
  results: ClassificationResult[];
}

/**
 * LLM-based Relationship Classifier
 */
export class RelationshipClassifier {
  private provider: LLMProvider;
  /** Maximum facts to classify in a single batch (to avoid token limits) */
  private maxBatchSize: number;

  constructor(provider: LLMProvider, options?: { maxBatchSize?: number }) {
    this.provider = provider;
    this.maxBatchSize = options?.maxBatchSize ?? 10;
  }

  /**
   * Classify how a new fact relates to existing facts
   */
  async classify(
    newFact: FactToClassify,
    existingFacts: ExistingFact[]
  ): Promise<ClassificationResult> {
    // If no existing facts, it's always NEW
    if (existingFacts.length === 0) {
      return {
        classification: 'NEW',
        confidence: 1.0,
        reason: 'No existing facts to compare against',
      };
    }

    try {
      const prompt = this.buildClassificationPrompt(newFact, existingFacts);

      const request: CompletionRequest = {
        messages: [{ role: 'user', content: prompt }],
        temperature: 0.1, // Low temperature for consistent classification
        maxTokens: 200,
        purpose: 'relation_classify',
      };

      const response = await this.provider.complete(request);
      const content = this.extractTextContent(response.content);
      return this.parseClassificationResponse(content);
    } catch (error) {
      // Default to NEW on error
      return {
        classification: 'NEW',
        confidence: 0.5,
        reason: `Classification failed: ${(error as Error).message}`,
      };
    }
  }

  /**
   * Batch classify multiple facts in a single LLM call
   * This is much more efficient than calling classify() for each fact
   */
  async classifyBatch(
    newFacts: FactToClassify[],
    existingFacts: ExistingFact[]
  ): Promise<ClassificationResult[]> {
    // If no new facts, return empty
    if (newFacts.length === 0) {
      return [];
    }

    // If no existing facts, all are NEW
    if (existingFacts.length === 0) {
      return newFacts.map(() => ({
        classification: 'NEW' as const,
        confidence: 1.0,
        reason: 'No existing facts to compare against',
      }));
    }

    // Split into batches if too many facts
    if (newFacts.length > this.maxBatchSize) {
      const results: ClassificationResult[] = [];
      for (let i = 0; i < newFacts.length; i += this.maxBatchSize) {
        const batch = newFacts.slice(i, i + this.maxBatchSize);
        const batchResults = await this.classifyBatch(batch, existingFacts);
        results.push(...batchResults);
      }
      return results;
    }

    try {
      const prompt = this.buildBatchClassificationPrompt(newFacts, existingFacts);

      const request: CompletionRequest = {
        messages: [{ role: 'user', content: prompt }],
        temperature: 0.1,
        maxTokens: 100 + newFacts.length * 100, // Scale tokens with batch size
        purpose: 'relation_classify',
      };

      const response = await this.provider.complete(request);
      const content = this.extractTextContent(response.content);
      return this.parseBatchClassificationResponse(content, newFacts.length);
    } catch (error) {
      // Default all to NEW on error
      return newFacts.map(() => ({
        classification: 'NEW' as const,
        confidence: 0.5,
        reason: `Batch classification failed: ${(error as Error).message}`,
      }));
    }
  }

  /**
   * Build the classification prompt
   */
  private buildClassificationPrompt(
    newFact: FactToClassify,
    existingFacts: ExistingFact[]
  ): string {
    const existingFactsStr = existingFacts
      .map((f) => `- [${f.id}] (${f.subject}, ${f.category}): "${f.content}"`)
      .join('\n');

    return `${CLASSIFICATION_PROMPT}

EXISTING FACTS:
${existingFactsStr}

NEW FACT:
Subject: ${newFact.subject}
Category: ${newFact.category}
Content: "${newFact.content}"

Classify this new fact:`;
  }

  /**
   * Build batch classification prompt for multiple facts
   */
  private buildBatchClassificationPrompt(
    newFacts: FactToClassify[],
    existingFacts: ExistingFact[]
  ): string {
    const existingFactsStr = existingFacts
      .map((f) => `- [${f.id}] (${f.subject}, ${f.category}): "${f.content}"`)
      .join('\n');

    const newFactsStr = newFacts
      .map((f, i) => `${i + 1}. Subject: ${f.subject}, Category: ${f.category}, Content: "${f.content}"`)
      .join('\n');

    return `${CLASSIFICATION_PROMPT}

EXISTING FACTS:
${existingFactsStr}

NEW FACTS TO CLASSIFY:
${newFactsStr}

Classify ALL new facts. Respond with a JSON array:
{
  "classifications": [
    {"index": 1, "classification": "NEW|UPDATES|EXTENDS", "targetId": "id or null", "confidence": 0.0-1.0, "reason": "brief explanation"},
    ...
  ]
}`;
  }

  /**
   * Extract text content from LLM response (handles ContentBlock[])
   */
  private extractTextContent(content: ContentBlock[] | string): string {
    if (typeof content === 'string') {
      return content;
    }
    return content
      .map((block) => ('text' in block ? block.text : ''))
      .join('');
  }

  /**
   * Parse the classification response
   */
  private parseClassificationResponse(content: string): ClassificationResult {
    try {
      // Extract JSON from response
      const jsonMatch = content.match(/\{[\s\S]*\}/);
      if (!jsonMatch) {
        throw new Error('No JSON found in response');
      }

      const parsed = JSON.parse(jsonMatch[0]);

      // Validate classification
      const validClassifications = ['NEW', 'UPDATES', 'EXTENDS'];
      if (!validClassifications.includes(parsed.classification)) {
        throw new Error(`Invalid classification: ${parsed.classification}`);
      }

      return {
        classification: parsed.classification,
        targetId: parsed.targetId,
        confidence: parsed.confidence ?? 0.8,
        reason: parsed.reason ?? 'No reason provided',
      };
    } catch {
      // Default to NEW on parse failure
      return {
        classification: 'NEW',
        confidence: 0.5,
        reason: 'Failed to parse classification response',
      };
    }
  }

  /**
   * Parse batch classification response
   */
  private parseBatchClassificationResponse(content: string, expectedCount: number): ClassificationResult[] {
    const defaultResult = (): ClassificationResult => ({
      classification: 'NEW',
      confidence: 0.5,
      reason: 'Failed to parse batch response',
    });

    try {
      const jsonMatch = content.match(/\{[\s\S]*\}/);
      if (!jsonMatch) {
        return Array(expectedCount).fill(null).map(defaultResult);
      }

      const parsed = JSON.parse(jsonMatch[0]);
      const classifications = parsed.classifications;

      if (!Array.isArray(classifications)) {
        return Array(expectedCount).fill(null).map(defaultResult);
      }

      const validClassifications = ['NEW', 'UPDATES', 'EXTENDS'];
      const results: ClassificationResult[] = [];

      // Build results array in order
      for (let i = 0; i < expectedCount; i++) {
        const item = classifications.find((c: { index: number }) => c.index === i + 1);
        if (item && validClassifications.includes(item.classification)) {
          results.push({
            classification: item.classification,
            targetId: item.targetId || undefined,
            confidence: item.confidence ?? 0.8,
            reason: item.reason ?? 'No reason provided',
          });
        } else {
          results.push(defaultResult());
        }
      }

      return results;
    } catch {
      return Array(expectedCount).fill(null).map(defaultResult);
    }
  }

}

/**
 * Options for creating a RelationshipClassifier
 */
export interface RelationshipClassifierOptions {
  /** Maximum facts to classify in a single batch (default: 10) */
  maxBatchSize?: number;
}

/**
 * Create a RelationshipClassifier instance
 */
export function createRelationshipClassifier(
  provider: LLMProvider,
  options?: RelationshipClassifierOptions
): RelationshipClassifier {
  return new RelationshipClassifier(provider, options);
}
