import { anthropic, MODELS } from '../lib/anthropic.js';
import { DraftKBSchema, type Chunk, type DraftKB, type DraftTheme, type DraftTopic, type DocumentFormat } from '../types.js';

const BATCH_SIZE = 40;
const BATCH_OVERLAP = 5;
const LARGE_DOC_THRESHOLD = 60;

const SYSTEM_PROMPT = `You are a knowledge architect. Your task is to analyse a set of text chunks from a source document and extract a structured knowledge base.

Output ONLY valid JSON matching the schema provided. No preamble, no explanation, no markdown fences.

Rules:
- Group related content into Themes. A Theme is a broad subject area.
- Under each Theme, identify discrete Topics. A Topic covers one specific concept.
- Identify relationships between Topics: related, prerequisite, or contrast.
  - related: Topics that complement each other
  - prerequisite: Topic A must be understood before Topic B
  - contrast: Topics that represent opposing approaches or concepts
- For each Topic, extract key terms suitable for a glossary.
- Assign a complexity weight (1–5) to each Topic.
  1 = introductory, 5 = advanced
- Draft a body for each Topic (2–4 paragraphs) based on the source chunks.
- Draft a description for each Theme (1–2 sentences).
- Every Topic must reference the chunk IDs that contributed to it.

Output schema:
{
  "themes": [
    {
      "title": "string",
      "description": "string",
      "topics": [
        {
          "title": "string",
          "body": "string",
          "difficulty": "introductory" | "intermediate" | "advanced",
          "complexityWeight": 1-5,
          "keyTerms": ["string"],
          "sourceChunkIds": ["chunk-id"],
          "relationships": {
            "related": ["topic title"],
            "prerequisites": ["topic title"],
            "contrasts": ["topic title"]
          }
        }
      ]
    }
  ]
}`;

const STRICT_SUFFIX = '\n\nCRITICAL: Your entire response must be valid JSON only. No text before or after.';

function buildUserPrompt(chunks: Chunk[], filename: string, format: DocumentFormat, strict: boolean): string {
  const chunkText = chunks
    .map(c => `[CHUNK-${c.id}]\n${c.text}`)
    .join('\n\n');

  return `Source document: ${filename}\nFormat: ${format}\n\nChunks:\n${chunkText}\n\nExtract the knowledge base structure from these chunks.${strict ? STRICT_SUFFIX : ''}`;
}

async function callClaude(chunks: Chunk[], filename: string, format: DocumentFormat, strict: boolean): Promise<DraftKB> {
  const response = await anthropic.messages.create({
    model: MODELS.SONNET,
    max_tokens: 8000,
    temperature: 0,
    system: SYSTEM_PROMPT,
    messages: [{ role: 'user', content: buildUserPrompt(chunks, filename, format, strict) }],
  });

  const textBlock = response.content.find(b => b.type === 'text');
  if (!textBlock || textBlock.type !== 'text') {
    throw new Error('structure_extraction_failed: no text block in response');
  }

  let parsed: unknown;
  try {
    parsed = JSON.parse(textBlock.text);
  } catch {
    if (strict) throw new Error('structure_extraction_failed');
    return callClaude(chunks, filename, format, true);
  }

  const result = DraftKBSchema.safeParse(parsed);
  if (!result.success) {
    if (strict) throw new Error('structure_extraction_failed');
    return callClaude(chunks, filename, format, true);
  }

  if (result.data.themes.length === 0) {
    throw new Error('no_structure_found');
  }

  return result.data;
}

// ---------------------------------------------------------------------------
// DraftKB merge (for large documents processed in batches)
// ---------------------------------------------------------------------------

interface MergedTheme {
  title: string;
  description: string;
  topicMap: Map<string, DraftTopic>;
}

function mergeDraftKBs(batches: DraftKB[]): DraftKB {
  const themeMap = new Map<string, MergedTheme>();

  for (const kb of batches) {
    for (const theme of kb.themes) {
      const key = theme.title.toLowerCase().trim();
      const existing = themeMap.get(key);

      if (!existing) {
        themeMap.set(key, {
          title: theme.title,
          description: theme.description,
          topicMap: new Map(theme.topics.map(t => [t.title.toLowerCase().trim(), { ...t }])),
        });
      } else {
        if (theme.description.length > existing.description.length) {
          existing.description = theme.description;
        }
        for (const topic of theme.topics) {
          const tKey = topic.title.toLowerCase().trim();
          const existingTopic = existing.topicMap.get(tKey);
          if (!existingTopic) {
            existing.topicMap.set(tKey, { ...topic });
          } else {
            existingTopic.body =
              existingTopic.body.length >= topic.body.length ? existingTopic.body : topic.body;
            existingTopic.sourceChunkIds = [
              ...new Set([...existingTopic.sourceChunkIds, ...topic.sourceChunkIds]),
            ];
            existingTopic.relationships = {
              related: [...new Set([...existingTopic.relationships.related, ...topic.relationships.related])],
              prerequisites: [...new Set([...existingTopic.relationships.prerequisites, ...topic.relationships.prerequisites])],
              contrasts: [...new Set([...existingTopic.relationships.contrasts, ...topic.relationships.contrasts])],
            };
          }
        }
      }
    }
  }

  const themes: DraftTheme[] = [...themeMap.values()].map(t => ({
    title: t.title,
    description: t.description,
    topics: [...t.topicMap.values()],
  }));

  return { themes };
}

// ---------------------------------------------------------------------------
// Public entry
// ---------------------------------------------------------------------------

export async function extractStructure(
  chunks: Chunk[],
  filename: string,
  format: DocumentFormat,
): Promise<DraftKB> {
  if (chunks.length <= LARGE_DOC_THRESHOLD) {
    return callClaude(chunks, filename, format, false);
  }

  const batches: DraftKB[] = [];
  let start = 0;

  while (start < chunks.length) {
    const end = Math.min(start + BATCH_SIZE, chunks.length);
    const batchChunks = chunks.slice(start, end);
    const batchKB = await callClaude(batchChunks, filename, format, false);
    batches.push(batchKB);
    if (end >= chunks.length) break;
    start = end - BATCH_OVERLAP;
  }

  return mergeDraftKBs(batches);
}