package net.micode.notes.ui;

import android.content.Context;
import android.media.AudioFormat;
import android.media.AudioRecord;
import android.media.MediaRecorder;
import android.os.Handler;
import android.os.Looper;
import android.util.Log;

import org.json.JSONObject;
import net.micode.notes.vosk.Model;
import net.micode.notes.vosk.Recognizer;
import net.micode.notes.vosk.LibVosk;
import net.micode.notes.vosk.LogLevel;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class VoskRecognitionService {
    private static final String TAG = "VoskRecognitionService";
    private static final String MODEL_DIR = "vosk-model-small-cn-0.22";
    private static final int SAMPLE_RATE = 16000;
    private static final int CHANNEL_IN = AudioFormat.CHANNEL_IN_MONO;
    private static final int ENCODING = AudioFormat.ENCODING_PCM_16BIT;
    private static final int BUFFER_SIZE_FACTOR = 2;

    private Context mContext;
    private Model mModel;
    private Recognizer mRecognizer;
    private AudioRecord mAudioRecord;
    private ExecutorService mExecutor;
    private Handler mMainHandler;

    private boolean mIsListening = false;
    private boolean mIsInitialized = false;

    private List<RecognitionListener> mListeners;

    public interface RecognitionListener {
        void onReadyForSpeech();
        void onBeginningOfSpeech();
        void onEndOfSpeech();
        void onResult(String text);
        void onPartialResult(String text);
        void onError(String error);
    }

    public VoskRecognitionService(Context context) {
        mContext = context.getApplicationContext();
        mExecutor = Executors.newSingleThreadExecutor();
        mMainHandler = new Handler(Looper.getMainLooper());
        mListeners = new ArrayList<>();
    }

    public void addListener(RecognitionListener listener) {
        if (listener != null && !mListeners.contains(listener)) {
            mListeners.add(listener);
        }
    }

    public void removeListener(RecognitionListener listener) {
        mListeners.remove(listener);
    }

    public void initialize() {
        mExecutor.execute(() -> {
            try {
                LibVosk.setLogLevel(LogLevel.INFO);
                
                String modelPath = getModelPath();
                if (modelPath == null) {
                    copyModelFromAssets();
                    modelPath = getModelPath();
                }

                if (modelPath != null && new File(modelPath).exists()) {
                    mModel = new Model(modelPath);
                    mRecognizer = new Recognizer(mModel, SAMPLE_RATE);
                    mRecognizer.setMaxAlternatives(0);
                    mRecognizer.setWords(true);
                    mIsInitialized = true;
                    Log.i(TAG, "Vosk recognition initialized successfully");
                } else {
                    Log.e(TAG, "Model path not found: " + modelPath);
                }
            } catch (Exception e) {
                Log.e(TAG, "Failed to initialize Vosk: " + e.getMessage());
            }
        });
    }

    private String getModelPath() {
        File modelDir = new File(mContext.getFilesDir(), MODEL_DIR);
        if (modelDir.exists()) {
            return modelDir.getAbsolutePath();
        }
        return null;
    }

    private void copyModelFromAssets() {
        try {
            String[] files = mContext.getAssets().list("");
            boolean modelExists = false;
            
            if (files != null) {
                for (String file : files) {
                    if (file.equals(MODEL_DIR)) {
                        modelExists = true;
                        break;
                    }
                }
            }

            if (!modelExists) {
                Log.w(TAG, "Model not found in assets. Please download the model manually.");
                Log.w(TAG, "Download from: https://alphacephei.com/vosk/models/vosk-model-small-cn-0.22.zip");
                return;
            }

            File modelDir = new File(mContext.getFilesDir(), MODEL_DIR);
            if (!modelDir.exists()) {
                modelDir.mkdirs();
            }

            copyAssets(MODEL_DIR, modelDir.getAbsolutePath());
            Log.i(TAG, "Model copied from assets to: " + modelDir.getAbsolutePath());
        } catch (IOException e) {
            Log.e(TAG, "Failed to copy model from assets: " + e.getMessage());
        }
    }

    private void copyAssets(String assetDir, String outDir) throws IOException {
        String[] files = mContext.getAssets().list(assetDir);
        if (files != null) {
            for (String file : files) {
                String fullAssetPath = assetDir + "/" + file;
                String fullOutPath = outDir + "/" + file;
                
                File outFile = new File(fullOutPath);
                
                if (mContext.getAssets().list(fullAssetPath).length > 0) {
                    outFile.mkdirs();
                    copyAssets(fullAssetPath, fullOutPath);
                } else {
                    InputStream in = mContext.getAssets().open(fullAssetPath);
                    OutputStream out = new FileOutputStream(outFile);
                    byte[] buffer = new byte[1024];
                    int read;
                    while ((read = in.read(buffer)) != -1) {
                        out.write(buffer, 0, read);
                    }
                    in.close();
                    out.flush();
                    out.close();
                }
            }
        }
    }

    public void startListening() {
        if (!mIsInitialized) {
            notifyError("语音识别未初始化，请确保模型文件已正确放置");
            return;
        }

        if (mIsListening) {
            stopListening();
        }

        mExecutor.execute(() -> {
            try {
                int minBufferSize = AudioRecord.getMinBufferSize(SAMPLE_RATE, CHANNEL_IN, ENCODING);
                int bufferSize = minBufferSize * BUFFER_SIZE_FACTOR;

                mAudioRecord = new AudioRecord(
                        MediaRecorder.AudioSource.MIC,
                        SAMPLE_RATE,
                        CHANNEL_IN,
                        ENCODING,
                        bufferSize
                );

                if (mAudioRecord.getState() != AudioRecord.STATE_INITIALIZED) {
                    notifyError("无法初始化录音设备");
                    return;
                }

                mIsListening = true;
                mAudioRecord.startRecording();

                mMainHandler.post(() -> notifyReadyForSpeech());

                byte[] buffer = new byte[bufferSize];

                while (mIsListening && mAudioRecord != null) {
                    int readSize = mAudioRecord.read(buffer, 0, buffer.length);
                    
                    if (readSize > 0) {
                        if (mRecognizer.acceptWaveForm(buffer, readSize)) {
                            String result = mRecognizer.getResult();
                            String text = parseResultText(result);
                            if (!text.isEmpty()) {
                                mMainHandler.post(() -> notifyResult(text));
                            }
                        } else if (mRecognizer.getPartialResult() != null) {
                            String partialResult = mRecognizer.getPartialResult();
                            String text = parsePartialResultText(partialResult);
                            if (!text.isEmpty()) {
                                mMainHandler.post(() -> notifyPartialResult(text));
                            }
                        }
                    }
                }
            } catch (Exception e) {
                Log.e(TAG, "Error during recognition: " + e.getMessage());
                notifyError("语音识别错误: " + e.getMessage());
            }
        });
    }

    public void stopListening() {
        mIsListening = false;
        
        if (mAudioRecord != null) {
            try {
                mAudioRecord.stop();
                mAudioRecord.release();
            } catch (Exception e) {
                Log.e(TAG, "Error stopping audio record: " + e.getMessage());
            }
            mAudioRecord = null;
        }

        if (mRecognizer != null) {
            String finalResult = mRecognizer.getFinalResult();
            if (finalResult != null && !finalResult.isEmpty()) {
                String text = parseResultText(finalResult);
                if (!text.isEmpty()) {
                    notifyResult(text);
                }
            }
        }

        mMainHandler.post(() -> notifyEndOfSpeech());
    }

    public boolean isListening() {
        return mIsListening;
    }

    public boolean isInitialized() {
        return mIsInitialized;
    }

    private String parseResultText(String json) {
        try {
            JSONObject jsonObject = new JSONObject(json);
            if (jsonObject.has("text") && !jsonObject.isNull("text")) {
                return jsonObject.getString("text").trim();
            }
        } catch (Exception e) {
            Log.e(TAG, "Error parsing result: " + e.getMessage());
        }
        return "";
    }

    private String parsePartialResultText(String json) {
        try {
            JSONObject jsonObject = new JSONObject(json);
            if (jsonObject.has("partial") && !jsonObject.isNull("partial")) {
                return jsonObject.getString("partial").trim();
            }
        } catch (Exception e) {
            Log.e(TAG, "Error parsing partial result: " + e.getMessage());
        }
        return "";
    }

    private void notifyReadyForSpeech() {
        for (RecognitionListener listener : mListeners) {
            listener.onReadyForSpeech();
        }
    }

    private void notifyBeginningOfSpeech() {
        for (RecognitionListener listener : mListeners) {
            listener.onBeginningOfSpeech();
        }
    }

    private void notifyEndOfSpeech() {
        for (RecognitionListener listener : mListeners) {
            listener.onEndOfSpeech();
        }
    }

    private void notifyResult(String text) {
        for (RecognitionListener listener : mListeners) {
            listener.onResult(text);
        }
    }

    private void notifyPartialResult(String text) {
        for (RecognitionListener listener : mListeners) {
            listener.onPartialResult(text);
        }
    }

    private void notifyError(String error) {
        for (RecognitionListener listener : mListeners) {
            listener.onError(error);
        }
    }

    public void release() {
        stopListening();
        
        if (mRecognizer != null) {
            mRecognizer.close();
            mRecognizer = null;
        }
        
        if (mModel != null) {
            mModel.close();
            mModel = null;
        }
        
        if (mExecutor != null) {
            mExecutor.shutdown();
            mExecutor = null;
        }
        
        mIsInitialized = false;
        Log.i(TAG, "Vosk recognition released");
    }
}
