import { NextRequest } from 'next/server'
import { ChatOpenAI } from "@langchain/openai"
import { ChatPromptTemplate } from "@langchain/core/prompts"
import { StructuredOutputParser } from "langchain/output_parsers"
import { RunnableSequence, RunnableLambda } from "@langchain/core/runnables"
import { z } from "zod"
import { AIMessageChunk, BaseMessage, MessageContent } from "@langchain/core/messages"
import db from "@/lib/db"
import { tavily } from "@tavily/core"

// Add this type definition near the top of the file
type ProgressCallback = (step: string, message: string, data?: any) => void;

type GenerateInput = {
  bookTitle: string;
  author: string;
  useTavily: boolean;
  onProgress: ProgressCallback;
};

const model = new ChatOpenAI({ 
  openAIApiKey: process.env.OPENAI_API_KEY,
  temperature: 0.7,
  modelName: "gpt-4o",
})

const listChaptersTemplate = `
You are an expert book analyzer. Your task is to list the chapters of the book "{bookTitle}" by {author}. 
If you don't know the exact chapters, provide a reasonable estimate based on the book's topic and structure.
Return the list as an array of chapter titles.

Book: {bookTitle}
Author: {author}

Chapters:
`

const summarizeChapterTemplate = `
Provide a detailed summary of the following chapter from the book "{bookTitle}" by {author}:

Chapter: {chapterTitle}

Your summary should:
1. Explain the main concepts introduced in this chapter
2. Discuss key examples or studies mentioned
3. Highlight how this chapter relates to the book's overall theme
4. Include any important quotes or memorable ideas

Aim for a comprehensive summary of about 200-250 words.

Summary:
`

const combineChaptersTemplate = `
You are an expert book summarizer for Blinkist. Your task is to create a comprehensive summary of the book "{bookTitle}" by {author} based on the detailed chapter summaries provided.

Chapter Summaries:
{chapterSummaries}

Additional Context:
{additionalContext}

Please structure your final summary as follows:

1. Brief Overview (3-4 sentences about the book's main theme and structure)
2. Key Ideas (5-7 main points from the book, derived from the chapter summaries, with brief explanations)
3. Notable Concepts and Studies (2-3 significant concepts or studies discussed in the book)
4. Actionable Takeaways (3-4 practical lessons readers can apply, with brief explanations)
5. Critical Analysis (1-2 sentences on the book's strengths or potential criticisms)
6. Conclusion (2-3 sentences wrapping up the summary and the book's significance)

Your summary should be informative, engaging, and similar in style to Blinkist summaries. Aim for a length of about 600-800 words.

Important: Do not summarize the chapter summaries again. They will be included in full separately.

{format_instructions}
`

const listChaptersPrompt = ChatPromptTemplate.fromTemplate(listChaptersTemplate)
const summarizeChapterPrompt = ChatPromptTemplate.fromTemplate(summarizeChapterTemplate)

const finalSummaryParser = StructuredOutputParser.fromZodSchema(
  z.object({
    briefOverview: z.string().describe("3-4 sentences about the book's main theme and structure"),
    keyIdeas: z.array(z.object({
      idea: z.string(),
      explanation: z.string()
    })).describe("5-7 main points from the book with brief explanations"),
    notableConceptsAndStudies: z.array(z.string()).describe("2-3 significant concepts or studies discussed in the book"),
    actionableTakeaways: z.array(z.object({
      takeaway: z.string(),
      explanation: z.string()
    })).describe("3-4 practical lessons readers can apply with brief explanations"),
    criticalAnalysis: z.string().describe("1-2 sentences on the book's strengths or potential criticisms"),
    conclusion: z.string().describe("2-3 sentences wrapping up the summary and the book's significance"),
  })
);

const combineChaptersPrompt = ChatPromptTemplate.fromTemplate(combineChaptersTemplate)
  .pipe(model)
  .pipe(finalSummaryParser);

const listChaptersChain = listChaptersPrompt.pipe(model)
const summarizeChapterChain = summarizeChapterPrompt.pipe(model)

type ChaptersResult = AIMessageChunk | BaseMessage[] | string | { content: MessageContent };

// Add Tavily client initialization
const tvly = tavily({ apiKey: process.env.TAVILY_API_KEY });

// Modify the generateChapterSummaries function
const generateChapterSummaries = new RunnableLambda({
  func: async (input: GenerateInput) => {
    let additionalContext = '';
    if (input.useTavily) {
      input.onProgress('tavily_search_start', 'Searching for additional context');
      const searchQuery = `${input.bookTitle} by ${input.author} book summary`;
      const searchResult = await tvly.search(searchQuery, {
        maxResults: 5,
        searchDepth: "advanced",
      });
      
      // Log the entire Tavily search result
      console.log('Tavily search result:', JSON.stringify(searchResult, null, 2));
      
      // Log each result separately for easier reading
      searchResult.results.forEach((result, index) => {
        console.log(`Tavily result ${index + 1}:`, {
          title: result.title,
          url: result.url,
          content: result.content.substring(0, 200) + '...' // Log first 200 characters of content
        });
      });

      additionalContext = searchResult.results.map((r: any) => r.content).join('\\n\\n');
      input.onProgress('tavily_search_complete', 'Additional context retrieved');
    }

    const chaptersResult = await listChaptersChain.invoke({
      bookTitle: input.bookTitle,
      author: input.author,
    }) as ChaptersResult;
    
    console.log('chaptersResult:', JSON.stringify(chaptersResult, null, 2)) // Log the result
    let chaptersContent = '';
    
    if (chaptersResult instanceof AIMessageChunk) {
      chaptersContent = stringifyMessageContent(chaptersResult.content);
    } else if (Array.isArray(chaptersResult)) {
      chaptersContent = chaptersResult.map((chunk: BaseMessage | string) => {
        if (typeof chunk === 'string') return chunk;
        if (chunk instanceof BaseMessage) return stringifyMessageContent(chunk.content);
        return JSON.stringify(chunk);
      }).join('');
    } else if (typeof chaptersResult === 'string') {
      chaptersContent = chaptersResult;
    } else if (chaptersResult && typeof chaptersResult === 'object' && 'content' in chaptersResult) {
      chaptersContent = stringifyMessageContent(chaptersResult.content);
    } else {
      console.error('Unexpected chaptersResult:', chaptersResult) // Log the unexpected result
      throw new Error('Unexpected response format from listChaptersChain');
    }
    
    console.log('chaptersContent:', chaptersContent) // Log the content
    
    // Extract the JSON array from the content
    const jsonMatch = chaptersContent.match(/```json\\n([\\s\\S]*?)\\n```/);
    let chapters: string[] = [];
    if (jsonMatch && jsonMatch[1]) {
      try {
        chapters = JSON.parse(jsonMatch[1]);
      } catch (error) {
        console.error('Error parsing JSON:', error);
        chapters = chaptersContent.split('\\n').map(line => line.trim()).filter(line => line !== '');
      }
    } else {
      chapters = chaptersContent.split('\\n').map(line => line.trim()).filter(line => line !== '');
    }
    
    const chapterSummaries = await Promise.all(
      chapters.map(async (chapterTitle: string) => {
        const summaryResult = await summarizeChapterChain.invoke({
          bookTitle: input.bookTitle,
          author: input.author,
          chapterTitle
        }) as ChaptersResult;
        let summaryContent = '';
        
        if (summaryResult instanceof AIMessageChunk) {
          summaryContent = stringifyMessageContent(summaryResult.content);
        } else if (Array.isArray(summaryResult)) {
          summaryContent = summaryResult.map((chunk: BaseMessage | string) => {
            if (typeof chunk === 'string') return chunk;
            if (chunk instanceof BaseMessage) return stringifyMessageContent(chunk.content);
            return JSON.stringify(chunk);
          }).join('');
        } else if (typeof summaryResult === 'string') {
          summaryContent = summaryResult;
        } else if (summaryResult && typeof summaryResult === 'object' && 'content' in summaryResult) {
          summaryContent = stringifyMessageContent(summaryResult.content);
        } else {
          console.error('Unexpected summaryResult:', summaryResult) // Log the unexpected result
          throw new Error('Unexpected response format from summarizeChapterChain');
        }
        
        return { title: chapterTitle, summary: summaryContent }
      })
    )
    
    return {
      ...input,
      chapters,
      chapterSummaries,
      chapterSummariesText: chapterSummaries.map(ch => `${ch.title}: ${ch.summary}`).join('\\n\\n'),
      additionalContext
    }
  }
})

function stringifyMessageContent(content: MessageContent): string {
  if (typeof content === 'string') {
    return content;
  } else if (Array.isArray(content)) {
    return content.map(item => {
      if (typeof item === 'string') {
        return item;
      } else if (typeof item === 'object' && item !== null) {
        return JSON.stringify(item);
      }
      return '';
    }).join('');
  }
  return '';
}

type GenerateChapterSummariesResult = {
  bookTitle: string;
  author: string;
  useTavily: boolean;
  onProgress: ProgressCallback;
  chapters: string[];
  chapterSummaries: { title: string; summary: string; }[];
  chapterSummariesText: string;
  additionalContext: string;
};

const fullChain = RunnableSequence.from([
  async (input: GenerateInput) => {
    input.onProgress('chapters_start', 'Generating chapter list');
    const result = await generateChapterSummaries.invoke(input);
    input.onProgress('chapters_complete', 'Chapter list generated', { chapters: result.chapters });
    return result;
  },
  async (input: GenerateChapterSummariesResult) => {
    input.onProgress('summary_start', 'Generating final summary');
    const finalSummaryResult = await combineChaptersPrompt.invoke({
      bookTitle: input.bookTitle,
      author: input.author,
      chapterSummaries: input.chapterSummariesText,
      additionalContext: input.additionalContext,
      format_instructions: finalSummaryParser.getFormatInstructions(),
    })
    input.onProgress('summary_complete', 'Final summary generated');

    return {
      chapters: input.chapters,
      chapterSummaries: input.chapterSummaries,
      finalSummary: finalSummaryResult
    }
  }
])

// Modify the sendSSEUpdate function
function sendSSEUpdate(writer: WritableStreamDefaultWriter<Uint8Array>, data: any) {
  const encoder = new TextEncoder();
  const message = `data: ${JSON.stringify(data)}\\n\\n`;
  writer.write(encoder.encode(message));
}

export async function POST(req: NextRequest) {
  const responseStream = new TransformStream();
  const writer = responseStream.writable.getWriter();

  const encoder = new TextEncoder();

  try {
    const { bookTitle, author, useTavily } = await req.json()

    if (!bookTitle || !author) {
      sendSSEUpdate(writer, { step: 'error', message: 'Book title and author are required' });
      writer.close();
      return new Response(responseStream.readable, {
        headers: {
          'Content-Type': 'text/event-stream',
          'Cache-Control': 'no-cache',
          'Connection': 'keep-alive',
        },
      });
    }

    sendSSEUpdate(writer, { step: 'start', message: 'Starting summary generation' });

    try {
      const result = await fullChain.invoke({ 
        bookTitle, 
        author, 
        useTavily,
        onProgress: (step: string, message: string) => sendSSEUpdate(writer, { step, message })
      });

      sendSSEUpdate(writer, { step: 'summary_complete', message: 'Summary generation complete' });

      // Save the summary to the database
      const book = await db.book.create({
        data: {
          title: bookTitle,
          author: author,
          summary: result.finalSummary.briefOverview,
          detailedSummary: {
            create: {
              briefOverview: result.finalSummary.briefOverview,
              chapterSummaries: result.chapterSummaries,
              keyIdeas: result.finalSummary.keyIdeas,
              notableConceptsStudies: result.finalSummary.notableConceptsAndStudies,
              actionableTakeaways: result.finalSummary.actionableTakeaways,
              criticalAnalysis: result.finalSummary.criticalAnalysis,
              conclusion: result.finalSummary.conclusion,
            },
          },
        },
        include: {
          detailedSummary: true,
        },
      })

      sendSSEUpdate(writer, { step: 'database_save', message: 'Summary saved to database' });
      sendSSEUpdate(writer, { step: 'complete', message: 'Process complete', data: { ...result, savedBook: book } });
    } catch (innerError: unknown) {
      console.error('Error in summary generation or database save:', innerError);
      let errorMessage = 'An unknown error occurred';
      if (innerError instanceof Error) {
        errorMessage = innerError.message;
      } else if (typeof innerError === 'string') {
        errorMessage = innerError;
      }
      sendSSEUpdate(writer, { step: 'error', message: `Error: ${errorMessage}` });
    }

    writer.close();
  } catch (error: unknown) {
    console.error('Error in POST handler:', error)
    let errorMessage = 'An unknown error occurred';
    if (error instanceof Error) {
      errorMessage = error.message;
    } else if (typeof error === 'string') {
      errorMessage = error;
    }
    sendSSEUpdate(writer, { step: 'error', message: errorMessage });
    writer.close();
  }

  return new Response(responseStream.readable, {
    headers: {
      'Content-Type': 'text/event-stream',
      'Cache-Control': 'no-cache',
      'Connection': 'keep-alive',
    },
  });
}

// Remove or comment out the GET function if it exists
// export async function GET(req: NextRequest) { ... }