import { useMemo } from "react";
import { CitesTradeRecord } from "@/lib/api";

// Helper types for visualization data
type YearlyData = {
  year: number;
  count: number;
};

type CountryData = {
  country: string;
  count: number;
};

type TermData = {
  term: string;
  count: number;
};

type PurposeData = {
  purpose: string;
  count: number;
  description: string;
};

type SourceData = {
  source: string;
  count: number;
  description: string;
};

type TermYearData = {
  year: number;
  [term: string]: number;
};

type VisualizationData = {
  recordsByYear: YearlyData[];
  topImporters: CountryData[];
  topExporters: CountryData[];
  termsTraded: TermData[];
  tradePurposes: PurposeData[];
  tradeSources: SourceData[];
  termQuantitiesByYear: TermYearData[];
  topTerms: string[];
} | null;

// Purpose code descriptions
const PURPOSE_DESCRIPTIONS: Record<string, string> = {
  'T': 'Commercial',
  'Z': 'Zoo',
  'G': 'Botanical garden',
  'Q': 'Circus or travelling exhibition',
  'S': 'Scientific',
  'H': 'Hunting trophy',
  'P': 'Personal',
  'M': 'Medical',
  'E': 'Educational',
  'N': 'Reintroduction',
  'B': 'Breeding',
  'L': 'Law enforcement'
};

// Source code descriptions
const SOURCE_DESCRIPTIONS: Record<string, string> = {
  'W': 'Wild',
  'R': 'Ranched',
  'D': 'Captive-bred (App. I)',
  'A': 'Artificially propagated (plants)',
  'C': 'Captive-bred (App. II/III)',
  'F': 'Born in captivity (F1)',
  'U': 'Unknown',
  'I': 'Confiscated/seized',
  'O': 'Pre-Convention'
};

export function useTradeDataVisualization(
  tradeRecords: CitesTradeRecord[] | undefined,
  filteredRecords: CitesTradeRecord[] | undefined = undefined
) {
  // Use filtered records for visualization if provided, otherwise use all records
  const recordsToVisualize = filteredRecords || tradeRecords;
  
  // Compute visualization data
  const visualizationData = useMemo((): VisualizationData => {
    if (!recordsToVisualize || recordsToVisualize.length === 0) return null;

    // Records by year
    const yearCounts: Record<number, number> = {};
    recordsToVisualize.forEach(record => {
      yearCounts[record.year] = (yearCounts[record.year] || 0) + 1;
    });
    const recordsByYear: YearlyData[] = Object.entries(yearCounts)
      .map(([year, count]) => ({ year: parseInt(year), count }))
      .sort((a, b) => a.year - b.year);

    // Top importers
    const importerCounts: Record<string, number> = {};
    recordsToVisualize.forEach(record => {
      if (record.importer) {
        importerCounts[record.importer] = (importerCounts[record.importer] || 0) + 1;
      }
    });
    const topImporters: CountryData[] = Object.entries(importerCounts)
      .map(([country, count]) => ({ country, count }))
      .sort((a, b) => b.count - a.count)
      .slice(0, 10);

    // Top exporters
    const exporterCounts: Record<string, number> = {};
    recordsToVisualize.forEach(record => {
      if (record.exporter) {
        exporterCounts[record.exporter] = (exporterCounts[record.exporter] || 0) + 1;
      }
    });
    const topExporters: CountryData[] = Object.entries(exporterCounts)
      .map(([country, count]) => ({ country, count }))
      .sort((a, b) => b.count - a.count)
      .slice(0, 10);

    // Terms traded
    const termCounts: Record<string, number> = {};
    recordsToVisualize.forEach(record => {
      termCounts[record.term] = (termCounts[record.term] || 0) + 1;
    });
    const termsTraded: TermData[] = Object.entries(termCounts)
      .map(([term, count]) => ({ term, count }))
      .sort((a, b) => b.count - a.count);

    // Trade purposes
    const purposeCounts: Record<string, number> = {};
    recordsToVisualize.forEach(record => {
      if (record.purpose) {
        purposeCounts[record.purpose] = (purposeCounts[record.purpose] || 0) + 1;
      }
    });
    const tradePurposes: PurposeData[] = Object.entries(purposeCounts)
      .map(([purpose, count]) => ({ 
        purpose, 
        count,
        description: PURPOSE_DESCRIPTIONS[purpose] || 'Unknown'
      }))
      .sort((a, b) => b.count - a.count);

    // Trade sources
    const sourceCounts: Record<string, number> = {};
    recordsToVisualize.forEach(record => {
      if (record.source) {
        sourceCounts[record.source] = (sourceCounts[record.source] || 0) + 1;
      }
    });
    const tradeSources: SourceData[] = Object.entries(sourceCounts)
      .map(([source, count]) => ({ 
        source, 
        count,
        description: SOURCE_DESCRIPTIONS[source] || 'Unknown'
      }))
      .sort((a, b) => b.count - a.count);

    // Top terms by year
    const topTerms = termsTraded.slice(0, 5).map(t => t.term);
    const termsByYear: Record<number, Record<string, number>> = {};
    
    recordsToVisualize.forEach(record => {
      if (!termsByYear[record.year]) {
        termsByYear[record.year] = {};
      }
      if (topTerms.includes(record.term)) {
        termsByYear[record.year][record.term] = (termsByYear[record.year][record.term] || 0) + (record.quantity || 1);
      }
    });
    
    const termQuantitiesByYear: TermYearData[] = Object.entries(termsByYear)
      .map(([year, terms]) => {
        const yearData: TermYearData = { year: parseInt(year) };
        topTerms.forEach(term => {
          yearData[term] = terms[term] || 0;
        });
        return yearData;
      })
      .sort((a, b) => a.year - b.year);

    return {
      recordsByYear,
      topImporters,
      topExporters,
      termsTraded,
      tradePurposes,
      tradeSources,
      termQuantitiesByYear,
      topTerms
    };
  }, [recordsToVisualize]);

  // Get unique years and terms (from all records, not just filtered ones)
  const years = useMemo(() => {
    if (!tradeRecords) return [];
    const uniqueYears = [...new Set(tradeRecords.map((r: CitesTradeRecord) => r.year))]
      .filter(year => year !== null && year !== undefined)
      .sort((a, b) => Number(a) - Number(b)); // Sort ascending for year ranges
    return uniqueYears;
  }, [tradeRecords]);
  
  const terms = useMemo(() => {
    if (!tradeRecords) return [];
    const uniqueTerms = [...new Set(tradeRecords.map((r: CitesTradeRecord) => r.term))]
      .filter(term => term !== null && term !== undefined && term !== '')
      .sort();
    return uniqueTerms;
  }, [tradeRecords]);
  
  const appendices = useMemo(() => {
    if (!tradeRecords) return [];
    const uniqueAppendices = [...new Set(tradeRecords.map((r: CitesTradeRecord) => r.appendix))]
      .filter(appendix => appendix !== null && appendix !== undefined && appendix !== '')
      .sort();
    return uniqueAppendices;
  }, [tradeRecords]);

  return {
    visualizationData,
    years,
    terms,
    appendices,
    PURPOSE_DESCRIPTIONS,
    SOURCE_DESCRIPTIONS
  };
}