/**
 * RagAgentService — web-retrieval AI agent built on Claude tool use.
 *
 * Two jobs:
 *
 * Job 1 — Form Retrieval
 *   Find and download current official Thai government PDF forms from government
 *   portals. First run bootstraps the stock library. Can re-run to refresh stale PDFs.
 *   Tools: fetch_page, search_links, download_pdf, store_form
 *
 * Job 2 — Law Retrieval
 *   Given a matter context (practice area, keywords, matter IDs), browse Thai
 *   government and official legal sources for relevant statutes, circulars,
 *   regulations, and policy updates. Download and index them in the law corpus.
 *   Called at matter open, when practice area changes, or on demand.
 *   Tools: fetch_page, search_links, download_pdf, store_law
 */

import { Injectable, Logger } from '@nestjs/common'
import { InjectModel } from '@nestjs/mongoose'
import { Model } from 'mongoose'
import { ConfigService } from '@nestjs/config'
import Anthropic from '@anthropic-ai/sdk'
import type { Tool, MessageParam, ToolUseBlock, ToolResultBlockParam } from '@anthropic-ai/sdk/resources/messages'
import { FormTemplate, FormTemplateDocument, LawCorpusEntry, LawCorpusEntryDocument } from '../common/schemas'
import { FormCalibrationAgent } from '../form-templates/form-calibration.agent'
import { STOCK_FORMS } from '../form-templates/stock-library'
import * as fs from 'fs'
import * as path from 'path'

// ─── Shared run state ─────────────────────────────────────────────────────────

export type RunType = 'form-retrieval' | 'law-retrieval'

// ─── Form retrieval ──────────────────────────────────────────────────────────

export interface FormRetrievalResult {
  stockId: string
  name: string
  status: 'found' | 'not_found' | 'error'
  url?: string
  sizeBytes?: number
  message: string
}

export interface RetrievalRun {
  runId: string
  runType: RunType
  startedAt: string
  finishedAt?: string
  status: 'running' | 'done' | 'error'
  results: FormRetrievalResult[]
  log: string[]
}

// ─── Law retrieval ────────────────────────────────────────────────────────────

export interface LawRetrievalContext {
  practiceArea: string              // e.g. 'immigration' | 'corporate' | 'property'
  keywords?: string[]               // additional search terms
  matterId?: string                 // matter that triggered this retrieval
  jurisdiction?: string             // defaults to 'th'
}

export interface LawRetrievalResult {
  title: string
  lawType: string
  sourceUrl: string
  storageKey?: string
  status: 'stored' | 'linked' | 'error'
  message: string
}

export interface LawRetrievalRun {
  runId: string
  runType: 'law-retrieval'
  context: LawRetrievalContext
  startedAt: string
  finishedAt?: string
  status: 'running' | 'done' | 'error'
  results: LawRetrievalResult[]
  log: string[]
}

// ─── Helpers ──────────────────────────────────────────────────────────────────

const HEADERS = {
  'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36',
  'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
  'Accept-Language': 'th,en;q=0.9',
}

function extractLinks(html: string, baseUrl: string): { text: string; href: string }[] {
  const base = new URL(baseUrl)
  const links: { text: string; href: string }[] = []
  const re = /<a[^>]+href=["']([^"']+)["'][^>]*>([\s\S]*?)<\/a>/gi
  let m: RegExpExecArray | null
  while ((m = re.exec(html)) !== null) {
    const rawHref = m[1].trim()
    const text = m[2].replace(/<[^>]+>/g, '').replace(/\s+/g, ' ').trim()
    if (!rawHref || rawHref.startsWith('#') || rawHref.startsWith('javascript')) continue
    try {
      const href = new URL(rawHref, base).toString()
      links.push({ text, href })
    } catch { /* skip malformed */ }
  }
  return links
}

function stripHtml(html: string): string {
  return html
    .replace(/<style[\s\S]*?<\/style>/gi, '')
    .replace(/<script[\s\S]*?<\/script>/gi, '')
    .replace(/<[^>]+>/g, ' ')
    .replace(/&nbsp;/g, ' ')
    .replace(/&amp;/g, '&')
    .replace(/&lt;/g, '<')
    .replace(/&gt;/g, '>')
    .replace(/\s{3,}/g, '\n\n')
    .trim()
    .slice(0, 6000)
}

function isPdf(buffer: Buffer): boolean {
  return buffer.slice(0, 4).toString('ascii') === '%PDF'
}

// ─── Service ──────────────────────────────────────────────────────────────────

@Injectable()
export class RagAgentService {
  private readonly logger = new Logger(RagAgentService.name)
  private readonly anthropic: Anthropic
  private readonly MODEL = 'claude-sonnet-4-6'
  private currentRun: RetrievalRun | null = null
  private currentLawRun: LawRetrievalRun | null = null

  // Page cache: url → { links, text } — so agent can call search_links without re-fetching
  private pageCache = new Map<string, { links: { text: string; href: string }[]; text: string }>()

  constructor(
    @InjectModel(FormTemplate.name) private templateModel: Model<FormTemplateDocument>,
    @InjectModel(LawCorpusEntry.name) private lawModel: Model<LawCorpusEntryDocument>,
    private readonly config: ConfigService,
    private readonly calibrationAgent: FormCalibrationAgent,
  ) {
    this.anthropic = new Anthropic({ apiKey: config.get<string>('ANTHROPIC_API_KEY') })
  }

  getCurrentRun(): RetrievalRun | null { return this.currentRun }
  getCurrentLawRun(): LawRetrievalRun | null { return this.currentLawRun }

  // ── Tool definitions ───────────────────────────────────────────────────────

  private get FORM_TOOLS(): Tool[] {
    return [
      ...this.BASE_TOOLS,
      {
        name: 'store_form',
        description: 'Save the last downloaded PDF to disk and associate it with a stock form template in the database. Call this immediately after a successful download_pdf.',
        input_schema: {
          type: 'object' as const,
          properties: {
            stock_id:   { type: 'string', description: 'The stockId of the form template (e.g. "th-immigration-tm7")' },
            source_url: { type: 'string', description: 'The URL the PDF was downloaded from (for record-keeping)' },
            notes:      { type: 'string', description: 'Optional notes about this form version or source' },
          },
          required: ['stock_id', 'source_url'],
        },
      },
    ]
  }

  private get LAW_TOOLS(): Tool[] {
    return [
      ...this.BASE_TOOLS,
      {
        name: 'store_law',
        description: 'Save a retrieved law, circular, regulation, or policy update to the law corpus database. Call after download_pdf (if the document is a PDF) or after fetch_page (if the content is HTML/text only).',
        input_schema: {
          type: 'object' as const,
          properties: {
            title: { type: 'string', description: 'Official title of the law or document' },
            title_local: { type: 'string', description: 'Thai language title if the primary title is English' },
            law_type: {
              type: 'string',
              enum: ['statute', 'circular', 'regulation', 'amendment', 'policy_update', 'gazette', 'other'],
              description: 'Type of legal document',
            },
            source_url: { type: 'string', description: 'URL where this document was found' },
            issued_by: { type: 'string', description: 'Issuing authority (e.g. "Immigration Bureau", "BOI", "Revenue Department")' },
            reference_code: { type: 'string', description: 'Official reference number, section, or order number' },
            effective_date: { type: 'string', description: 'Effective date in YYYY-MM-DD format if known' },
            language: { type: 'string', enum: ['th', 'en', 'bilingual'], description: 'Primary language of the document' },
            summary: { type: 'string', description: 'Brief summary of what the law/circular says (1-3 sentences)' },
            tags: { type: 'array', items: { type: 'string' }, description: 'Relevant keywords for this document' },
          },
          required: ['title', 'law_type', 'source_url'],
        },
      },
    ]
  }

  private get BASE_TOOLS(): Tool[] {
    return [
      {
        name: 'fetch_page',
        description: 'Fetch a webpage and return its visible text content and all hyperlinks. Use this to navigate government websites and find form download pages.',
        input_schema: {
          type: 'object' as const,
          properties: {
            url: { type: 'string', description: 'Full URL of the page to fetch' },
          },
          required: ['url'],
        },
      },
      {
        name: 'search_links',
        description: 'Filter the links from the most recently fetched page by a keyword. Returns matching links with their text and URLs. Use after fetch_page to find specific download links.',
        input_schema: {
          type: 'object' as const,
          properties: {
            page_url: { type: 'string', description: 'URL of the page whose links to search (must have been fetched already)' },
            keyword:  { type: 'string', description: 'Keyword to filter links by (case-insensitive match on link text or URL)' },
          },
          required: ['page_url', 'keyword'],
        },
      },
      {
        name: 'download_pdf',
        description: 'Download a file from a URL and verify it is a valid PDF. Returns size in bytes if successful. The file is held in memory until store_form or store_law is called.',
        input_schema: {
          type: 'object' as const,
          properties: {
            url: { type: 'string', description: 'Direct URL to the PDF file' },
          },
          required: ['url'],
        },
      },
    ]
  }

  // ── Tool implementations ───────────────────────────────────────────────────

  private pendingPdf: { url: string; buffer: Buffer } | null = null

  private async toolFetchPage(url: string): Promise<string> {
    try {
      const res = await fetch(url, { headers: HEADERS, signal: AbortSignal.timeout(15_000) })
      if (!res.ok) return `HTTP ${res.status} from ${url}`
      const ct = res.headers.get('content-type') ?? ''
      if (ct.includes('pdf')) {
        return `This URL returns a PDF directly (content-type: ${ct}). Use download_pdf instead.`
      }
      const html = await res.text()
      const links = extractLinks(html, url)
      const text  = stripHtml(html)
      this.pageCache.set(url, { links, text })
      const linkSummary = links.slice(0, 40)
        .map(l => `  [${l.text.slice(0, 60)}] → ${l.href}`)
        .join('\n')
      return `PAGE TEXT:\n${text}\n\nLINKS (first 40):\n${linkSummary}\n\nTotal links found: ${links.length}`
    } catch (e: any) {
      return `Error fetching ${url}: ${e.message}`
    }
  }

  private toolSearchLinks(pageUrl: string, keyword: string): string {
    const cached = this.pageCache.get(pageUrl)
    if (!cached) return `Page not in cache. Call fetch_page("${pageUrl}") first.`
    const kw = keyword.toLowerCase()
    const matches = cached.links.filter(
      l => l.text.toLowerCase().includes(kw) || l.href.toLowerCase().includes(kw),
    )
    if (!matches.length) return `No links matching "${keyword}" on ${pageUrl}`
    return matches.map(l => `  [${l.text.slice(0, 80)}] → ${l.href}`).join('\n')
  }

  private async toolDownloadPdf(url: string): Promise<string> {
    try {
      const res = await fetch(url, { headers: HEADERS, signal: AbortSignal.timeout(30_000) })
      if (!res.ok) return `HTTP ${res.status} — cannot download from ${url}`
      const ct = res.headers.get('content-type') ?? ''
      const buffer = Buffer.from(await res.arrayBuffer())
      if (!isPdf(buffer)) {
        return `Downloaded ${buffer.length} bytes but it is NOT a valid PDF (content-type: ${ct}). First bytes: ${buffer.slice(0, 20).toString('hex')}`
      }
      this.pendingPdf = { url, buffer }
      return `Downloaded valid PDF: ${buffer.length} bytes from ${url}. Now call store_form() to save it.`
    } catch (e: any) {
      return `Error downloading ${url}: ${e.message}`
    }
  }

  private async toolWebSearch(query: string): Promise<string> {
    const url = `https://lite.duckduckgo.com/lite/?q=${encodeURIComponent(query)}`
    try {
      const res = await fetch(url, {
        headers: { ...HEADERS, 'Accept': 'text/html' },
        signal: AbortSignal.timeout(15_000),
      })
      if (!res.ok) return `Search failed: HTTP ${res.status}`
      const html = await res.text()
      const links = extractLinks(html, url)
      this.pageCache.set(url, { links, text: stripHtml(html) })
      // Filter out DDG navigation links — keep only actual result destinations
      const results = links.filter(l =>
        !l.href.includes('duckduckgo.com') &&
        !l.href.includes('duck.co') &&
        l.text.length > 4,
      ).slice(0, 12)
      if (!results.length) return `No web results found for: ${query}`
      return `Search results for "${query}":\n` +
        results.map(l => `  [${l.text.slice(0, 80)}] → ${l.href}`).join('\n') +
        `\n\n(Use fetch_page on a promising URL, or download_pdf if a URL ends in .pdf)`
    } catch (e: any) {
      return `Search error: ${e.message}`
    }
  }

  private async toolStoreForm(stockId: string, sourceUrl: string, notes?: string): Promise<string> {
    if (!this.pendingPdf) return 'No PDF in memory. Call download_pdf() first.'

    const template = await this.templateModel.findOne({ stockId }).lean()
    if (!template) return `No FormTemplate found with stockId "${stockId}". Available stockIds: ${STOCK_FORMS.map(f => f.stockId).join(', ')}`

    const stockDir = path.join(process.cwd(), 'uploads', 'form-templates', 'stock')
    if (!fs.existsSync(stockDir)) fs.mkdirSync(stockDir, { recursive: true })

    const safeId = stockId.replace(/[^a-zA-Z0-9_-]/g, '_')
    const filePath = path.join(stockDir, `${safeId}.pdf`)
    fs.writeFileSync(filePath, this.pendingPdf.buffer)

    const storageKey = `stock/${safeId}.pdf`
    const updateNotes = notes
      ? `${(template as any).notes ?? ''}\nSource: ${sourceUrl}. ${notes}`.trim()
      : `${(template as any).notes ?? ''}\nSource: ${sourceUrl}`.trim()

    await this.templateModel.findOneAndUpdate(
      { stockId },
      { storageKey, sourceUrl, notes: updateNotes },
    )

    const size = this.pendingPdf.buffer.length
    this.pendingPdf = null
    return `Stored ${size} bytes as "${storageKey}" and updated FormTemplate for ${stockId}.`
  }

  private async toolStoreLaw(
    input: {
      title: string
      title_local?: string
      law_type: string
      source_url: string
      issued_by?: string
      reference_code?: string
      effective_date?: string
      language?: string
      summary?: string
      tags?: string[]
    },
    context: LawRetrievalContext,
  ): Promise<string> {
    let storageKey: string | undefined

    // If there's a pending PDF, save it
    if (this.pendingPdf) {
      const lawDir = path.join(process.cwd(), 'uploads', 'law-corpus')
      if (!fs.existsSync(lawDir)) fs.mkdirSync(lawDir, { recursive: true })

      const safeTitle = input.title
        .replace(/[^a-zA-Z0-9ก-๙\s_-]/g, '')
        .replace(/\s+/g, '_')
        .slice(0, 60)
      const filename = `${Date.now()}_${safeTitle}.pdf`
      fs.writeFileSync(path.join(lawDir, filename), this.pendingPdf.buffer)
      storageKey = `law-corpus/${filename}`
      this.pendingPdf = null
    }

    // Upsert by source_url to avoid duplicates
    const existing = await this.lawModel.findOne({ sourceUrl: input.source_url })
    const relatedMatterIds = context.matterId
      ? [...new Set([...(existing?.relatedMatterIds ?? []), context.matterId])]
      : (existing?.relatedMatterIds ?? [])

    if (existing) {
      await this.lawModel.findByIdAndUpdate(existing._id, {
        title: input.title,
        titleLocal: input.title_local,
        lawType: input.law_type,
        practiceArea: context.practiceArea,
        jurisdiction: context.jurisdiction ?? 'th',
        issuedBy: input.issued_by,
        referenceCode: input.reference_code,
        effectiveDate: input.effective_date ? new Date(input.effective_date) : undefined,
        language: input.language ?? 'th',
        summary: input.summary,
        tags: input.tags ?? [],
        relatedMatterIds,
        retrievedBy: 'rag-agent',
        ...(storageKey ? { storageKey } : {}),
      })
      return `Updated existing law corpus entry for "${input.title}" (sourceUrl already known).`
    }

    await this.lawModel.create({
      title: input.title,
      titleLocal: input.title_local,
      lawType: input.law_type,
      practiceArea: context.practiceArea,
      jurisdiction: context.jurisdiction ?? 'th',
      sourceUrl: input.source_url,
      storageKey,
      issuedBy: input.issued_by,
      referenceCode: input.reference_code,
      effectiveDate: input.effective_date ? new Date(input.effective_date) : undefined,
      language: input.language ?? 'th',
      summary: input.summary,
      tags: input.tags ?? [],
      relatedMatterIds,
      retrievedBy: 'rag-agent',
    })

    return `Stored law corpus entry: "${input.title}" [${input.law_type}]${storageKey ? ` — PDF saved as ${storageKey}` : ' — linked by URL only'}.`
  }

  // ── Tool dispatcher ────────────────────────────────────────────────────────

  private async dispatchFormTool(name: string, input: any): Promise<string> {
    switch (name) {
      case 'web_search':    return this.toolWebSearch(input.query)
      case 'fetch_page':    return this.toolFetchPage(input.url)
      case 'search_links':  return this.toolSearchLinks(input.page_url, input.keyword)
      case 'download_pdf':  return this.toolDownloadPdf(input.url)
      case 'store_form':    return this.toolStoreForm(input.stock_id, input.source_url, input.notes)
      default:              return `Unknown tool: ${name}`
    }
  }

  private async dispatchLawTool(name: string, input: any, context: LawRetrievalContext): Promise<string> {
    switch (name) {
      case 'web_search':    return this.toolWebSearch(input.query)
      case 'fetch_page':    return this.toolFetchPage(input.url)
      case 'search_links':  return this.toolSearchLinks(input.page_url, input.keyword)
      case 'download_pdf':  return this.toolDownloadPdf(input.url)
      case 'store_law':     return this.toolStoreLaw(input, context)
      default:              return `Unknown tool: ${name}`
    }
  }

  // ── Main agent loop ────────────────────────────────────────────────────────

  async runFormRetrieval(): Promise<RetrievalRun> {
    const runId = Date.now().toString(36)
    const run: RetrievalRun = {
      runId, runType: 'form-retrieval', startedAt: new Date().toISOString(), status: 'running', results: [], log: [],
    }
    this.currentRun = run
    this.pageCache.clear()
    this.pendingPdf = null

    const log = (msg: string) => {
      this.logger.log(msg)
      run.log.push(`[${new Date().toTimeString().slice(0, 8)}] ${msg}`)
    }

    log('RAG Form Retrieval Agent starting')

    // Build the task list for the agent
    const formList = STOCK_FORMS.map((f, i) =>
      `${i + 1}. stockId="${f.stockId}" | "${f.name}" (${f.nameLocal})` + (f.sourceUrl ? `\n   Authority: ${f.sourceUrl.split('/')[2]} | Hint URL: ${f.sourceUrl}` : '')
    ).join('\n\n')

    const systemPrompt = `You are a legal document retrieval agent for a Thai law firm's practice management system.

Your task: find and download the current official PDF versions of Thai government legal forms from their official government websites.

You have four tools:
- fetch_page: browse any URL and see its text + links
- search_links: filter links from a fetched page by keyword
- download_pdf: download a file and verify it is a valid PDF
- store_form: save the downloaded PDF and record it in the database

STRATEGY:
- Start by fetching the government portal's forms/downloads page (not just the homepage)
- Thai government sites commonly have a "เอกสาร" (documents) or "ดาวน์โหลด" (download) section
- Look for links containing the form code (e.g. "TM7", "tm7", "ตม.7") or "pdf"
- If a URL redirects to HTML, try the page's /download/ or /forms/ subdirectory
- Immigration forms: try https://www.immigration.go.th and navigate to their forms section
- DOE work permit forms: try https://www.doe.go.th
- DBD forms: try https://www.dbd.go.th
- Always verify with download_pdf before calling store_form
- If you can't find the exact form, note why and move on — do not loop forever
- Process all ${STOCK_FORMS.length} forms

After retrieving all forms, summarise what you found and what you could not find.`

    const userMessage = `Find and download these ${STOCK_FORMS.length} Thai government forms. For each one, navigate to the official government website, find the current PDF download link, download it, and store it using store_form with the exact stockId shown.

FORMS TO RETRIEVE:
${formList}

Start with the immigration forms since there are the most. Good luck.`

    const messages: MessageParam[] = [{ role: 'user', content: userMessage }]

    // Agentic loop — max 80 turns to avoid runaway
    let turns = 0
    const MAX_TURNS = 80

    while (turns < MAX_TURNS) {
      turns++
      log(`Turn ${turns}`)

      const response = await this.anthropic.messages.create({
        model: this.MODEL,
        max_tokens: 4096,
        system: systemPrompt,
        tools: this.FORM_TOOLS,
        messages,
      })

      // Append assistant message
      messages.push({ role: 'assistant', content: response.content })

      if (response.stop_reason === 'end_turn') {
        // Extract final summary text
        const summary = response.content
          .filter((b: any) => b.type === 'text')
          .map((b: any) => b.text)
          .join('\n')
        log(`Agent finished:\n${summary}`)
        break
      }

      if (response.stop_reason !== 'tool_use') break

      // Execute tool calls
      const toolResults: ToolResultBlockParam[] = []
      for (const block of response.content) {
        if (block.type !== 'tool_use') continue
        const tb = block as ToolUseBlock
        log(`Tool: ${tb.name}(${JSON.stringify(tb.input).slice(0, 120)})`)
        const result = await this.dispatchFormTool(tb.name, tb.input)
        log(`  → ${result.slice(0, 200)}`)
        toolResults.push({ type: 'tool_result', tool_use_id: tb.id, content: result })

        // If a form was just stored, record the result
        if (tb.name === 'store_form' && result.startsWith('Stored')) {
          const stockId = (tb.input as any).stock_id
          const form = STOCK_FORMS.find(f => f.stockId === stockId)
          run.results.push({
            stockId,
            name: form?.name ?? stockId,
            status: 'found',
            url: (tb.input as any).source_url,
            message: result,
          })
        }
      }

      messages.push({ role: 'user', content: toolResults })
    }

    if (turns >= MAX_TURNS) log('Max turns reached — stopping agent')

    // Record not-found forms
    const foundIds = new Set(run.results.map(r => r.stockId))
    for (const f of STOCK_FORMS) {
      if (!foundIds.has(f.stockId)) {
        run.results.push({ stockId: f.stockId, name: f.name, status: 'not_found', message: 'Not retrieved in this run' })
      }
    }

    run.finishedAt = new Date().toISOString()
    run.status = 'done'

    // Kick off calibration for every form that was just retrieved
    const foundTemplates = await this.templateModel
      .find({ stockId: { $in: [...foundIds] } })
      .select('_id')
      .lean()

    if (foundTemplates.length > 0) {
      log(`Triggering calibration agent for ${foundTemplates.length} retrieved form(s)`)
      setTimeout(async () => {
        for (const tpl of foundTemplates) {
          await this.calibrationAgent.calibrateOne(String((tpl as any)._id))
        }
      }, 1000)
    }

    log(`Run complete. Found: ${foundIds.size}/${STOCK_FORMS.length}`)
    return run
  }

  // ── Law Retrieval ──────────────────────────────────────────────────────────

  /**
   * Retrieve laws, circulars, regulations, and policy updates relevant to a
   * given practice area and matter context. Browses official Thai government
   * and regulatory websites. Stores results in the law_corpus collection.
   *
   * Called at matter open, when a new practice area is added to an existing
   * matter, or on-demand when the attorney wants a current law sweep.
   */
  async runLawRetrieval(context: LawRetrievalContext): Promise<LawRetrievalRun> {
    const runId = Date.now().toString(36)
    const run: LawRetrievalRun = {
      runId,
      runType: 'law-retrieval',
      context,
      startedAt: new Date().toISOString(),
      status: 'running',
      results: [],
      log: [],
    }
    this.currentLawRun = run
    this.pageCache.clear()
    this.pendingPdf = null

    const log = (msg: string) => {
      this.logger.log(msg)
      run.log.push(`[${new Date().toTimeString().slice(0, 8)}] ${msg}`)
    }

    log(`Law Retrieval Agent starting — practice area: ${context.practiceArea}`)

    // Source map per practice area — key official portals to start from
    const sourcesMap: Record<string, string[]> = {
      immigration: [
        'https://www.immigration.go.th',
        'https://www.boi.go.th',
      ],
      corporate: [
        'https://www.dbd.go.th',
        'https://www.boi.go.th',
        'https://www.sec.or.th',
      ],
      property: [
        'https://www.dol.go.th',
        'https://www.reic.or.th',
      ],
      employment: [
        'https://www.doe.go.th',
        'https://www.mol.go.th',
      ],
      tax: [
        'https://www.rd.go.th',
      ],
      civil_litigation: [
        'https://www.coj.go.th',
        'https://krisdika.go.th',
      ],
      family: [
        'https://krisdika.go.th',
        'https://www.coj.go.th',
      ],
    }

    const startingSources = sourcesMap[context.practiceArea] ?? ['https://krisdika.go.th']
    const keywordsStr = context.keywords?.length
      ? `Additional keywords to focus on: ${context.keywords.join(', ')}`
      : ''

    const systemPrompt = `You are a legal research agent for a Thai law firm specialising in ${context.practiceArea}.

Your task: browse official Thai government and regulatory websites to find current laws, circulars, regulations, policy updates, and official guidance relevant to ${context.practiceArea} practice in Thailand.

You have four tools:
- fetch_page: browse any URL and see its text + links
- search_links: filter links from a fetched page by keyword
- download_pdf: download a file and verify it is a valid PDF (optional — use when document is a PDF you want to store)
- store_law: record a law or document in the firm's law corpus database

WHAT TO LOOK FOR:
- Primary legislation (acts, codes) governing ${context.practiceArea}
- Recent amendments or revisions (last 2 years particularly important)
- Circulars and policy updates from relevant authorities
- Official guidance notes and interpretations
- Any notifications or orders affecting current practice

STRATEGY:
- Start with the provided official sources
- Navigate to their laws/regulations/downloads sections
- Look for Thai keywords: กฎหมาย (law), ระเบียบ (regulation), ประกาศ (announcement/circular), คำสั่ง (order), แก้ไข (amendment)
- Use store_law for each relevant document found — you do NOT need to download every PDF, a URL link is sufficient
- If a document is clearly a PDF, download_pdf first then store_law
- Aim for 5-15 relevant entries per run — quality over quantity
- Do not store duplicates (same source_url)
${keywordsStr}

After completing retrieval, summarise what was found.`

    const userMessage = `Find and index current laws, regulations, and official guidance relevant to ${context.practiceArea} practice in Thailand.
${context.matterId ? `This retrieval was triggered by matter ID: ${context.matterId}` : ''}

Start with these official sources:
${startingSources.map((s, i) => `${i + 1}. ${s}`).join('\n')}

Browse each source, find relevant legal documents, and store them with store_law. Focus on what is current and directly applicable.`

    const messages: MessageParam[] = [{ role: 'user', content: userMessage }]

    let turns = 0
    const MAX_TURNS = 60

    while (turns < MAX_TURNS) {
      turns++
      log(`Turn ${turns}`)

      const response = await this.anthropic.messages.create({
        model: this.MODEL,
        max_tokens: 4096,
        system: systemPrompt,
        tools: this.LAW_TOOLS,
        messages,
      })

      messages.push({ role: 'assistant', content: response.content })

      if (response.stop_reason === 'end_turn') {
        const summary = response.content
          .filter((b: any) => b.type === 'text')
          .map((b: any) => b.text)
          .join('\n')
        log(`Agent finished:\n${summary}`)
        break
      }

      if (response.stop_reason !== 'tool_use') break

      const toolResults: ToolResultBlockParam[] = []
      for (const block of response.content) {
        if (block.type !== 'tool_use') continue
        const tb = block as ToolUseBlock
        log(`Tool: ${tb.name}(${JSON.stringify(tb.input).slice(0, 120)})`)
        const result = await this.dispatchLawTool(tb.name, tb.input, context)
        log(`  → ${result.slice(0, 200)}`)
        toolResults.push({ type: 'tool_result', tool_use_id: tb.id, content: result })

        // Record stored law in run results
        if (tb.name === 'store_law' && result.startsWith('Stored')) {
          run.results.push({
            title: (tb.input as any).title,
            lawType: (tb.input as any).law_type,
            sourceUrl: (tb.input as any).source_url,
            status: 'stored',
            message: result,
          })
        } else if (tb.name === 'store_law' && result.startsWith('Updated')) {
          run.results.push({
            title: (tb.input as any).title,
            lawType: (tb.input as any).law_type,
            sourceUrl: (tb.input as any).source_url,
            status: 'linked',
            message: result,
          })
        }
      }

      messages.push({ role: 'user', content: toolResults })
    }

    if (turns >= MAX_TURNS) log('Max turns reached — stopping agent')

    run.finishedAt = new Date().toISOString()
    run.status = 'done'
    log(`Law retrieval complete. Entries stored/linked: ${run.results.length}`)
    return run
  }
}
